Moving Iteration from Inference to Training: A New Approach to One-Step Image Generation
Generating an image with a diffusion model means running the same network dozens to hundreds of times. DiT-XL/2 reaches FID 2.27 on ImageNet 256x256 with 500 forward passes — 250 steps, doubled. It's slow, but the slowness has a reason: the model only learned to "take a small step" during training, and inference chains hundreds of such small steps together. Two existing routes can make generation one-shot. Distillation: first train a multi-step teacher, then compress hundreds of steps into one for a student — but the teacher's quality is the ceiling. Consistency models: directly learn the noise-to-data mapping without a teacher, but training is notoriously hard to tune — with the same 675M parameters, iCT only reached FID 34.24, and iMeanFlow, after two generations of improvements, reached 1.72.
A new paper from Kaiming He and collaborators (MIT; Mingyang Deng, He Li, Tianhong Li, Yilun Du, et al.) takes a third path, published as arXiv 2602.04770, *Generative Modeling via Drifting*. The idea fits in one sentence: iteration doesn't have to happen at inference time — it can happen at training time. The optimizer is already updating parameters step by step, and each update nudges "the image generated from a given noise" slightly — step i generates f_i(ε), step i+1 generates f_i+1(ε), and the difference Δx is the small step by which the sample is pushed by the parameter update. Diffusion models train to take one step but run hundreds at inference; this method runs hundreds of thousands of steps during training (the optimizer's step count is the drift step count) and only one at inference. The total drift is unchanged — it's just paid through a different window.
The remaining question is direction: where should Δx go? The paper defines a drift field, computed from samples of two distributions — samples drawn from the dataset are positive, samples drawn from the generator itself are negative. A kernel function pulls each generated sample toward the data and pushes it away from its peers. The entire design rests on an antisymmetry condition: swap the two distributions and the field flips sign exactly. The payoff is a one-line proof: when the generated distribution coincides with the data distribution, the field equals its own negation, hence is zero everywhere — all samples are at rest, and the model has reached equilibrium. The training objective is then written as a fixed point: pull the network's output toward "itself plus the drift," with gradients not propagated through the drift step. This frozen-target trick isn't new — it's a classic technique from self-supervised learning (the SimSiam family), here transplanted from representation learning into generation.
The equilibrium looks elegant but is actually fragile. An ablation deliberately breaks the antisymmetry: keeping only attraction (toward positive samples) without repelling the generator's own samples crashes FID from 8.46 to 177 — generated samples collapse onto a few points of the dataset. Attraction and repulsion must be strictly balanced, equal and opposite; this is the non-negotiable skeleton of the method.
Results
Numbers are on ImageNet 256x256:
- Latent space, one-step generation: FID 1.54 — beating all prior one-step methods (best was iMeanFlow at 1.72) and most multi-step models. DiT-XL/2 needed 500 forward passes for 2.27; Drifting's L/2 model (460M parameters) surpasses it in a single pass.
- Pixel space: FID 1.61 — matching PixelDiT, which uses 400x the compute, and leaving StyleGAN-XL (2.30) far behind.
- A hidden knob: the drift is computed not in pixel space but in feature space, using an encoder pretrained with their own MAE. Scaling the encoder width from 256 to 640 and pretraining from 192 to 1280 epochs — without touching a line of generator code — drops FID from 8.46 to 3.36. Part of the generation-quality ceiling now lives in the feature encoder. The encoder exists only during training and is discarded at inference, so it adds no deployment cost.