An End-to-End Latent-Rollout Approach for Pushing Few-Step ImageNet-$256$ Generation to FID $1.11$ without Fréchet Losses

Xiaoran Xu; Yujing Wang

cs.CVSep 26, 2026

An End-to-End Latent-Rollout Approach for Pushing Few-Step ImageNet- $256$ Generation to FID $1.11$ without Fréchet Losses

Authors: Xiaoran Xu, Yujing Wang

Organizations: ByteDance

Abstract

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\times256$ , 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\times$
Augmented	2	677	239.8	119.9	$1.01\times$
Augmented	3	677	359.7	119.9	$1.01\times$

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.

Many-step generation without Fréchet losses
Method	Space	NFE / Stg.	FID $\downarrow$	IS $\uparrow$	sFID $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$
DiT-XL/2 ( Peebles and Xie, 2023 )	SD	$250\times 2$	2.27	278.2	4.60	0.83	0.57
SiT-XL/2 ( Ma et al., 2024 )	SD	$250\times 2$	2.06	277.5	4.49	0.83	0.59
SiT-XL/2 + REPA ( Yu et al., 2025 )	SD	$250\times 2$	1.42	305.7	4.70	0.80	0.65
SiT-XL/1 + REPA-E ( Leng et al., 2025 )	E2E	$250\times 2$	1.12	302.9	4.09	0.79	0.66
RAE + DiT DH -XL ( Zheng et al., 2025 )	RAE	$50\times 2$	1.13	262.6	–	0.78	0.67

Table 2: System-level comparison on class-conditional ImageNet

256\times 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\times 256$ . Samples are generated from randomly drawn Gaussian noise and class labels.

Steps	FID $\downarrow$	IS $\uparrow$
4K	$60.58\rightarrow 1.69$	$35.7\rightarrow 239.7$
10K	$8.76\rightarrow 1.19$	$172.1\rightarrow 271.7$
40K	$5.23\rightarrow 1.11$	$205.7\rightarrow 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

\rightarrow

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 $\downarrow$	IS $\uparrow$
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.

\dagger

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\times 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.

Few-Step Diffusion Sampling

Few-step diffusion sampling develops ways to generate data with far fewer denoising evaluations while retaining quality and diversity.

See all 175 papers

AdvFidelity: Unifying Preference Alignment and Few-Step Image Generation within Reinforcement Learning
Oct 21, 2026Sicong Kang, Fangao Zeng, Yuhan Li +6RL for Image Generation Diffusion-Based Image Generation
Taobao & Tmall group of Alibaba · Shanghai Jiao Tong University
Reinforcing Few-step Generators via Reward-Tilted Distribution Matching
May 25, 2026Yushi Huang, Xiangxin Zhou, Ruoyu Wang +3Diffusion Model Alignment Diffusion Model Distillation
Tencent Hunyuan · Hong Kong University of Science and Technology · Westlake University
Few-Step Generation via Data-Space Iteration
Oct 8, 2026Shanchuan Lin, Yansong Peng, Fu-Yun Wang +1Iterative Refinement Flow-Based Generative Modeling
ByteDance Seed

Oct 21, 2026J/K move · Enter open · S save

AdvFidelity: Unifying Preference Alignment and Few-Step Image Generation within Reinforcement Learning

Sicong Kang, Fangao Zeng, Yuhan Li, Hao Zhou+5

Taobao & Tmall group of Alibaba · Shanghai Jiao Tong University

RL for Image Generation Diffusion-Based Image Generation Diffusion Model Alignment Few-Step Diffusion Sampling Reward Modeling

Open paper AskPDF arXiv

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.

Abstract

Few-Step Diffusion Sampling

Explore similar work

AdvFidelity: Unifying Preference Alignment and Few-Step Image Generation within Reinforcement Learning

Reinforcing Few-step Generators via Reward-Tilted Distribution Matching

Few-Step Generation via Data-Space Iteration