Image generation reaches near perfect quality with fewer steps
An End-to-End Latent-Rollout Approach for Pushing Few-Step ImageNet-$256$ Generation to FID $1.11$ without Fréchet Losses
Computer Vision and Pattern Recognition
Summary
Creating realistic images step by step is hard because early guesses affect the final picture. This paper presents a new way to generate images in just a few steps by first copying a powerful teacher model and then improving the process using feedback from real images. They use a special Transformer model to handle the steps in a compact, hidden space, letting the system learn efficiently without showing images until the end. This method produces very high-quality ImageNet pictures at 256 by 256 pixels, beating previous approaches that used more steps or complicated losses.
What this means in practice
- •For computer vision engineers: Develop fast, high-quality image generators that need fewer steps to produce realistic images, reducing computation in visual AI tasks.
- •For graphics software developers: Integrate efficient latent space refinement techniques into image editing tools to enhance the quality of generated visuals with less processing time.
Authors
Xiaoran Xu, Yujing Wang
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.