Iterative generation poses a joint optimization problem across steps, as intermediate predictions shape subsequent computations and ultimately determine the final output distribution. Few-step generators distilled from pretrained diffusion and flow-matching models make such optimization computationally practical end to end. We build on this opportunity with a distill-then-refine approach that uses teacher imitation to establish a strong initialization for a few-step rollout in latent space, then shifts to end-to-end refinement of the complete latent rollout against real data. We introduce FiST (Flow-in-Stage Transformer), an architecture that composes learned latent-state transitions in a few stages using a shared Transformer, with optional cross-stage hidden communication. Distillation applies teacher-forced regression to selected states along teacher trajectories; refinement replaces this supervision with adversarial and auxiliary classification objectives on the final latent output. A trainable discriminator module operates on semantically rich features extracted from clean real and generated latents by a frozen SiT backbone pretrained with REPA. All training takes place in latent space, without image decoding. During refinement, FiST consumes its own intermediate predictions, and endpoint gradients pass through every generation stage. For class-conditional generation on ImageNet at 256×256, our approach achieves FID 1.11 (IS 282) with three stages and FID 1.15 (IS 280) with two. These results demonstrate competitive few-step generation through learned distribution-level supervision, without explicit Fréchet-distance minimization. Ablations characterize how distillation, pretrained checkpoint choices, refinement supervision, and cross-stage hidden communication affect generation quality.
Figures & tables
Type
Stages
Parameters (M)
GFLOPs per image
Avg. GFLOPs per stage
SiT-equivalent compute per stage
Basic
3
673
356.1
118.7
1.00×
Augmented
2
677
239.8
119.9
1.01×
Augmented
3
677
359.7
119.9
1.01×
Table 1: Model size and inference cost of FiST. Augmented FiST includes the cross-stage hidden router. Parameter counts and FLOPs are approximate. Inference costs exclude VAE decoding. For reference, SiT-XL/2 has approximately 675M parameters and costs 118.7 GFLOPs per evaluation. SiT-equivalent compute per stage is the average per-stage cost relative to one such evaluation.
Method
Space
NFE / Stg.
FID ↓
IS ↑
sFID ↓
Prec. ↑
Rec. ↑
Many-step generation without Fréchet losses
DiT-XL/2 ( Peebles and Xie, 2023 )
SD
250×2
2.27
278.2
4.60
0.83
0.57
SiT-XL/2 ( Ma et al., 2024 )
SD
250×2
2.06
277.5
4.49
0.83
0.59
SiT-XL/2 + REPA ( Yu et al., 2025 )
SD
250×2
1.42
305.7
4.70
0.80
0.65
SiT-XL/1 + REPA-E ( Leng et al., 2025 )
E2E
250×2
1.12
302.9
4.09
0.79
0.66
RAE + DiT DH -XL ( Zheng et al., 2025 )
RAE
50×2
1.13
262.6
–
0.78
0.67
Table 2: System-level comparison on class-conditional ImageNet 256×256 . Baseline results follow the cited works and official releases. NFE denotes nominal network evaluations; Stg. denotes FiST stages (see Table 1 for costs). SD: SD-VAE; E2E: VAE trained jointly with the generator; Sem: SD-VAE + SemVAE; RAE: representation autoencoder. SIM: SigLIP + Inception + MAE.
Figure 1: Uncurated samples from three-stage augmented FiST on ImageNet 256×256 . Samples are generated from randomly drawn Gaussian noise and class labels.
Steps
FID ↓
IS ↑
4K
60.58→1.69
35.7→239.7
10K
8.76→1.19
172.1→271.7
40K
5.23→1.11
205.7→282.0
Table 3: Effects of distillation duration and teacher guidance. Results use three-stage augmented FiST. Teacher guidance uses REPA’s best-FID and qualitative-sampling scales, 1.8 and 4.0, with guidance intervals [0,0.7] and [0,1] , respectively ( Yu et al., 2025 ) . Arrows indicate before → after refinement; metrics are measured at the lowest-FID checkpoint within 20K refinement iterations.
Figure 2: Refinement dynamics under different distillation and teacher guidance settings. Experimental settings match Table 3 ; horizontal axes show refinement iterations. FID curves begin at 1K to prevent large pre-refinement values from obscuring later differences.
Refinement objectives
FID ↓
IS ↑
Adversarial only
1.22
277.0
Adversarial + classification
1.11
282.0
Table 4: Effects of auxiliary classification and pretrained checkpoints. Results use three-stage augmented FiST. Left: FID and IS at the lowest-FID checkpoint within 20K refinement iterations. Right: FID with different pretrained weights for distillation (teacher and FiST initialization) and the discriminator’s frozen backbone. SiT and REPA denote vanilla and REPA-trained SiT-XL/2 checkpoints, respectively. † Refinement does not improve the distilled generator’s FID.
Figure 3: Auxiliary classification during refinement. Horizontal axes show refinement iterations; iteration 0 denotes the distilled checkpoint before refinement.
Appendix figures & tables21 assets
Supplementary material from the paper’s appendix.
Appendix
Setting
Distillation
Refinement
Training duration
40K steps
20K iterations
Batch size
1024
1024
Learning rate
10−4
10−5 (G), 4×10−5 (D)
Optimizer
AdamW
AdamW
AdamW betas
(0.9,0.99)
(0.9,0.99)
Weight decay
0
0
Appendix
Table 5: Training hyperparameters. Each refinement iteration consists of one discriminator update followed by one generator update. G and D denote the generator and discriminator, respectively.
Figure 4: Additional uncurated samples. Class 246: Great Dane.
Figure 5: Additional uncurated samples. Class 287: lynx.
Figure 6: Additional uncurated samples. Class 353: gazelle.
Figure 7: Additional uncurated samples. Class 388: giant panda.
Figure 8: Additional uncurated samples. Class 393: anemone fish.
Figure 9: Additional uncurated samples. Class 407: ambulance.
Figure 10: Additional uncurated samples. Class 425: barn.
Figure 11: Additional uncurated samples. Class 483: castle.
Figure 12: Additional uncurated samples. Class 497: church.
Figure 13: Additional uncurated samples. Class 586: half track.
Figure 14: Additional uncurated samples. Class 625: lifeboat.
Figure 15: Additional uncurated samples. Class 630: Loafer.
Figure 16: Additional uncurated samples. Class 717: pickup.
Figure 17: Additional uncurated samples. Class 793: shower cap.
Figure 18: Additional uncurated samples. Class 871: trimaran.
Figure 19: Additional uncurated samples. Class 873: triumphal arch.
Figure 20: Additional uncurated samples. Class 888: viaduct.
Figure 21: Additional uncurated samples. Class 903: wig.
Figure 22: Additional uncurated samples. Class 949: strawberry.
Figure 23: Additional uncurated samples. Class 958: hay.
Practical image generation requires both human preference alignment and few-step efficiency, yet combining reinforcement learning (RL) with distillation is often multi-stage and may erode RL gains We observe that few-step RL progressively improves structural coherence, even with specialized rewards, but leaves low-level artifacts that these rewards may fail to penalize. To address this gap, we propose AdvFidelity, which integrates artifact-sensitive fidelity supervision into RL through an online adversarial reward. This enables joint reward alignment and high-fidelity few-step generation within a single RL stage, without distribution matching distillation (DMD) or other explicit distillation objectives. We train a discriminator to rank high-quality references (positives) above content-matched reconstructions (negatives) produced by the current generator. Reconstructing moderately noised references largely preserves style and scene layout while exposing low-level degradations. This encourages the discriminator to assess artifacts rather than content differences. With negatives reconstructed online, alternating updates adapt fidelity supervision to the generator's evolving failure modes. References can come from existing target-domain data or full-step model samples. Built on LeapAlign, AdvFidelity backpropagates gradients from both rewards through the few-step sampling trajectories used at inference. Across two model backbones and three task settings, AdvFidelity matches full-step LeapAlign in alignment, task performance, and overall generation quality and remains highly competitive with other evaluated full-step RL methods, using only 4 or 8 steps without CFG.
Sicong Kang, Fangao Zeng, Yuhan Li +6
Taobao & Tmall group of Alibaba · Shanghai Jiao Tong University
Recent advances in few-step diffusion distillation have enabled efficient image generation, yet aligning these models with human preferences remains challenging. We propose Reward-Tilted Distribution Matching Distillation (RTDMD), a two-stage framework that unifies distribution matching distillation with reward-guided reinforcement learning for few-step flow generators. We show that minimizing the KL divergence to a reward-tilted teacher distribution naturally decomposes into a distribution matching term and a reward maximization term. In the first stage, we introduce Ambient-Consistent Distribution Matching Distillation (AC-DMD), which performs subinterval-wise distribution matching and augments the fake score objective with a consistency regularizer to help the fake score model track the shifting generator distribution under limited updates. In the second stage, we jointly optimize both terms: for the reward maximization term, we derive a hybrid policy gradient that combines a GRPO-style estimator for the stochastic intermediate transitions with direct reward backpropagation through the deterministic final step, and further introduce step-subset GRPO (SubGRPO) to reduce variance. Experiments on SD3, SD3.5, and FLUX.2 demonstrate that RTDMD establishes new state-of-the-art results across preference, aesthetic, and compositional metrics with only 4 inference steps, outperforming previous few-step text-to-image generation methods. Code and models are available at https://github.com/Harahan/RTDMD.
Yushi Huang, Xiangxin Zhou, Ruoyu Wang +3
Tencent Hunyuan · Hong Kong University of Science and Technology · Westlake University
Flow matching has emerged as a scalable paradigm for training high-quality generative models, but sampling from the learned probability flow requires many network evaluations. Distillation can reduce this cost to one or a few evaluations; however, one-step generation often sacrifices quality, making few-step generation the practical operating regime. Existing few-step methods perform their iterative computation along the probability flow and therefore require a fixed, manually chosen timestep discretization. This discretization is often chosen heuristically and is expensive to tune; it may also be restrictive when refinement difficulty differs across samples or spatial locations. We introduce data-space iteration, a few-step generation framework that removes flow discretization altogether. Starting from noise, a shared generator directly refines its prediction in data space, with every iteration trained to produce the best sample permitted by its capacity. Our formulation integrates with distribution matching distillation (DMD) with minimal changes, enabling a controlled comparison between iteration methods under matched training settings. On class-conditional ImageNet 256x256, data-space iteration outperforms standard discretization baselines and matches or improves upon variants selected through schedule search, without requiring schedule-specific training. These results show that data-space iteration provides a simple and effective alternative to discretized flow-space iteration for fast generation.