Feature-Guided Rectified Flow Model

• The paper introduces a client-specific FG-RF that shifts generative modeling from pixel-level to feature-level, reducing privacy risks and computational overhead.
• FG-RF leverages deterministic ODE-based rectified flow to synthesize compressed features from non-IID, heterogeneous medical data in a one-shot federated framework.
• Experiments on Brain Tumor, Chest X-Ray TB, and HAM10000 datasets show substantial gains in accuracy and efficiency over DDPM and baseline methods.

Feature-Guided Rectified Flow Model (FG-RF) most directly denotes the client-side generative component introduced for one-shot federated learning (OSFL) in medical imaging, where pixel-level generation is replaced by feature-level generation: each client extracts compressed image features from private data, trains a rectified flow model on those features, uploads the trained model, and the server samples synthetic features for downstream aggregation (Ma et al., 25 Jul 2025). The term is not a universal label across rectified-flow research; rather, it names one concrete instantiation within federated classification, while adjacent work uses related mechanisms such as multi-scale feature conditioning, vector-field steering, mutual-information-based selection, and mid-step feature injection without standardizing the same nomenclature (Djouama et al., 21 Apr 2026, Patel et al., 2024, Wang et al., 18 Mar 2025, Kim et al., 2 Jul 2025).

1. Definition and problem setting

FG-RF was introduced to address three constraints in generative-model-based OSFL for medical imaging. First, pixel-level generators are treated as privacy-risky because diffusion models can memorize training data, so directly generating medical images can leak private information. Second, DDPM-based OSFL is slow in both training and sampling. Third, one-shot aggregation under non-IID medical data requires synthetic data that preserve client-specific distributions well enough to support a single global distillation stage (Ma et al., 25 Jul 2025).

In that formulation, FG-RF is a feature-space rectified flow model guided by local feature extractors, not a vanilla pixel-space RF model. Each client trains its own feature extractor and its own rectified flow model, so the learned flow matches that client’s data distribution. Compared with DDPMs, FG-RF uses a deterministic ODE trajectory rather than iterative denoising with many steps. Compared with GAN-, VAE-, and LDM-based OSFL baselines, it is designed specifically for medical feature synthesis rather than natural-image generation or pixel reconstruction (Ma et al., 25 Jul 2025).

A common misunderstanding is to interpret “feature-guided” as an inference-time auxiliary feature loss. In the named FG-RF formulation, the guidance is more basic and more structural: the generator is trained on features extracted from raw medical images, and those features are the objects being synthesized and aggregated. The model therefore relocates generative modeling from pixel space to feature space.

2. Rectified-flow foundations and the FG-RF objective

Rectified flow, in the broader literature, formulates generative modeling as learning an ODE

$dy_t = v_\Theta(y_t,t)\,dt$

with a straight-path interpolation such as

$z_t = (1-t)x_0 + t\epsilon,$

and learns a velocity field by conditional flow matching or an equivalent noise-prediction objective (Esser et al., 2024). In this sense, rectified flow replaces stochastic reverse diffusion with deterministic transport along a learned vector field.

FG-RF instantiates the same principle in client-local feature space. The client feature representation is denoted by $z_k$, and the local rectified flow model is

$\frac{dz_k}{dt}=g_{\phi_k}(z_{k,t},t).$

The interpolation path is defined by

$z_{k,t}=(1-t)z_{k,0}+t z_{k,T}, \qquad t\in[0,1],$

where $z_{k,0}$ is the target data feature and $z_{k,T}\sim\mathcal N(\mathbf 0,\mathbf I)$ is Gaussian noise. The training loss is

$$\mathcal{L}_{\mathrm{RFM}}=\mathbb{E}_{z_{k,0},z_{k,T},t}\left[\left\|g_{\phi_k}(z_{k,t},t)-(z_{k,T}-z_{k,0})\right\|^2\right].$$

This makes the model learn the straight-line target velocity $z_{k,T}-z_{k,0}$ rather than a long denoising chain (Ma et al., 25 Jul 2025).

After training, the server samples synthetic features from each uploaded client model by integrating the learned velocity field backward:

$$\hat z_k = z_{k,T}+\int_T^0 -g_{\phi_k}(z_k,t)\,dt.$$

The resulting synthetic features are assembled into a server-side dataset

$z_t = (1-t)x_0 + t\epsilon,$0

Within the paper’s interpretation, the principal gain is not a novel rectified-flow ODE but the relocation of transport to a privacy-buffered feature domain (Ma et al., 25 Jul 2025).

3. Client-side architecture and one-shot federated workflow

The client-side FG-RF pipeline has two stages: feature extraction and rectified-flow training. Each client first trains a local feature extractor $z_t = (1-t)x_0 + t\epsilon,$1 on private data. The feature extractor consists of a convolutional layer followed by tanh activation, and the classifier used in the experiments is ResNet-18. Optimization uses the cross-entropy loss

$z_t = (1-t)x_0 + t\epsilon,$2

The convolutional layer extracts informative image features, while tanh maps them into $z_t = (1-t)x_0 + t\epsilon,$3, matching diffusion/flow input conventions and compressing information (Ma et al., 25 Jul 2025).

The second stage trains the local rectified flow model $z_t = (1-t)x_0 + t\epsilon,$4 on the extracted features. The paper implements $z_t = (1-t)x_0 + t\epsilon,$5 using DiT, chosen because self-attention captures global structure better than U-Net-style local convolutions. This choice is notable because it aligns the federated medical-imaging setting with the broader movement toward transformer-based rectified flow, including MM-DiT-style multimodal backbones in text-to-image modeling (Ma et al., 25 Jul 2025, Esser et al., 2024).

The full OSFL workflow has three stages. On the client side, each client trains the feature extractor $z_t = (1-t)x_0 + t\epsilon,$6, trains a teacher classifier $z_t = (1-t)x_0 + t\epsilon,$7 on extracted features, trains the local RFM $z_t = (1-t)x_0 + t\epsilon,$8, and uploads $z_t = (1-t)x_0 + t\epsilon,$9 and $z_k$0 to the server. On the server side, synthetic features are sampled from each client RFM, client classifiers are averaged into a teacher ensemble

$z_k$1

and a global student model $z_k$2 is trained on the synthetic feature dataset. After one communication round, the trained student is sent back to clients (Ma et al., 25 Jul 2025).

This design makes FG-RF the data-synthesis engine of the OSFL framework. Its task is not end-user image generation but client-distribution-preserving synthetic feature production for one-round aggregation.

4. Privacy preservation and dual-layer knowledge distillation

FG-RF incorporates several privacy-preserving mechanisms. The most important is feature-level generation rather than pixel-level generation. The paper also emphasizes irreversible feature extraction through the lossy convolutional front-end and additional compression through tanh. The formal privacy argument uses the data processing inequality. Since

$z_k$3

forms a Markov chain,

$z_k$4

and because feature extraction is lossy,

$z_k$5

Therefore,

$z_k$6

This is the paper’s theoretical basis for the claim that feature-level synthesis leaks less than pixel-level synthesis (Ma et al., 25 Jul 2025).

The synthetic features are consumed by Dual-Layer Knowledge Distillation (DLKD). In compact form, the global student objective is

$z_k$7

The expanded objective is

$z_k$8

The KL term distills teacher logits, while the feature-alignment term matches intermediate representations:

$z_k$9

In the experiments, the aligned intermediate representation is the third layer of ResNet-18 (Ma et al., 25 Jul 2025).

The role of DLKD is to close the gap between synthetic features and real client data. This is especially important under non-IID heterogeneity, where output-logit mimicry alone is insufficient and intermediate-layer alignment transfers richer representational structure. In the paper’s design, FG-RF and DLKD are therefore complementary: FG-RF supplies the synthetic feature distribution, and DLKD regularizes how that distribution is translated into a global model.

5. Empirical behavior, efficiency, and ablation evidence

The reported experiments cover three non-IID medical imaging datasets: Brain Tumor, Chest X-Ray TB, and HAM10000. The framework is reported to outperform multi-round federated learning approaches by up to 21.73%, to exceed the baseline FedISCA by an average of 21.75%, and to improve over DDPM by an average of 24.27%. The reported accuracies are 60.15% on Brain Tumor, 86.32% on Chest X-Ray TB, and 70.01% on HAM10000, with the largest margin over FedISCA reaching 48.51% on HAM10000 (Ma et al., 25 Jul 2025).

Dataset	Reported accuracy	DDPM vs FG-RF time
Brain Tumor	60.15%	4.51 h / 22.68 s vs 1.87 h / 0.45 s
Chest X-Ray TB	86.32%	1.36 h / 29.32 s vs 0.58 h / 0.69 s
HAM10000	70.01%	12.92 h / 18.88 s vs 5.41 h / 0.29 s

These timings support the paper’s central efficiency claim: moving from DDPM to rectified flow materially reduces both training and sampling cost in the one-shot federated setting. The argument is consistent with broader rectified-flow evidence that straight-path ODE transport can be integrated with far fewer effective steps than reverse diffusion (Esser et al., 2024).

The ablation studies isolate the contribution of feature extraction, DLKD, and intermediate-layer alignment. Removing the feature extractor and DLKD yields 54.56, 77.09, and 68.97 on Brain Tumor, Chest X-Ray, and HAM10000, respectively. Removing DLKD alone yields 55.08, 83.16, and 64.91. Removing intermediate-layer alignment yields 55.84, 84.74, and 68.99. The full model reaches 60.15, 86.32, and 70.01. The pattern is consistent across datasets: feature extraction matters, DLKD matters, and intermediate-layer alignment adds measurable gains (Ma et al., 25 Jul 2025).

The privacy evaluation uses $\frac{dz_k}{dt}=g_{\phi_k}(z_{k,t},t).$0 distance between generated and original samples with memorization threshold $\frac{dz_k}{dt}=g_{\phi_k}(z_{k,t},t).$1. Feature-level generation yields larger distances than pixel-level generation, including Brain Tumor Client 0 with 0.3793 for pixel-level versus 0.5195 for feature-level, and Chest X-Ray Client 2 with 0.6304 versus 0.7734. Within the paper’s interpretation, this supports lower memorization risk for feature-level synthesis (Ma et al., 25 Jul 2025).

6. Related formulations, neighboring methods, and scope of the term

Only one paper in the cited set explicitly names a Feature-Guided Rectified Flow Model, but several adjacent works instantiate related ideas. RF-HiT, a rectified-flow hierarchical transformer for medical image segmentation, integrates an hourglass transformer backbone with a hierarchical feature encoder for anatomically guided feature conditioning, fuses multi-scale features via learnable interpolation, requires only 10.14 GFLOPs and 13.6M parameters, runs in as few as three steps, and reports 91.27% mean Dice on ACDC and 87.40% on BraTS 2021 (Djouama et al., 21 Apr 2026). In that case, “feature-guided” means multi-resolution conditioning inside the velocity network rather than feature-space federation.

The text-to-image scaling study on rectified flow establishes an important architectural and training backdrop. It formulates RF as an ODE with straight-line interpolation, introduces biased timestep sampling such as logit-normal and CosMap to emphasize harder intermediate timesteps, and proposes MM-DiT with separate weights for image and text streams and bidirectional attention mixing; however, it explicitly is not a separate FG-RF method paper (Esser et al., 2024). Its relevance is enabling rather than terminological.

FlowChef studies controlled generation by steering the rectified-flow vector field directly. It argues that RFMs have smoother vector fields, fewer NFEs, and fewer trajectory crossovers than diffusion models, and uses a deterministic, gradient-skipping update for classifier guidance, linear inverse problems, and image editing without extra training or inversion (Patel et al., 2024). RFMI, by contrast, does not define FG-RF but derives a prompt-image alignment signal from the conditional and unconditional RF velocity fields through a mutual-information estimator, then uses high point-wise MI samples for self-supervised LoRA fine-tuning (Wang et al., 18 Mar 2025). ReFlex performs training-free real-image editing in FLUX by extracting I2T-CA, I2I-SA, and residual features from a mid-step latent and adapting/injecting them during early sampling, again without naming the method FG-RF (Kim et al., 2 Jul 2025). RestoRect treats feature distillation as a latent rectified-flow process in which the student synthesizes teacher-quality features through learnable trajectories and a FLEX alignment loss, providing perhaps the closest restoration-oriented analogue of an FG-RF design (Verma et al., 27 Sep 2025).

A frequent misconception is therefore to treat FG-RF as a settled, canonical rectified-flow subtype. The cited literature supports a narrower statement: the explicit FG-RF name belongs to a client-side feature-space generator for one-shot federated medical classification (Ma et al., 25 Jul 2025), while the broader research area contains several FG-RF-like mechanisms—multi-scale feature injection, feature-space generation, vector-field steering, and internal-feature-guided editing—that are technically adjacent but not terminologically identical. This suggests that “feature-guided rectified flow” is better understood as a recurring design pattern built atop rectified-flow transport rather than as a single universally standardized algorithm.