Refine, Don't Re-roll

Flow-matching image models need tens to hundreds of network evaluations per sample, so we distill them into generators that need only a few, say four. The usual recipe fixes a timestep schedule, for example the noise levels 1.0, 0.75, 0.5 and 0.25, and lets each step move the sample from one level to the next. The schedule is hard to justify, and changing it means training a new generator.

Our paper asks whether it is needed at all. Instead, each step refines the image prediction directly, starting from the same noise, with no schedule.

1. The idea

Flow matching [6] connects data \(x_0\) and noise \(x_1\) by the line \(x_t=(1-t)x_0+t x_1\), and a network predicts its velocity \(v=x_1-x_0\). Any velocity prediction gives a clean-image prediction:

At \(t=1\) the state is the noise \(z\) itself, so \(\hat x_0=z-v\). Data-space iteration keeps only this prediction as the state and refines it with one shared generator:

• The same noise \(z\) enters every step, not a fresh draw.
• The step index \(r\) tells the generator which stage it is at.
• Nothing is re-noised, so no state at an intermediate noise level is ever built.

The network \(g_\theta\) is the pretrained velocity network with a new input projection for \(\hat x_0^{(r)}\) and an embedding of \(r\), and its flow-time input is fixed to \(t=1\) (for compatibility with the pretrained velocity parameterization).

2. What the schedule costs

Existing methods pass a noisy image from step to step, so they need a schedule \(1=t_0>t_1>\dots>t_R=0\) of noise levels, fixed before training. They differ in how a step moves from \(t_r\) to \(t_{r+1}\):

• Flow-renoise is what multi-step DMD does: predict a clean image \(\hat x_0\), then re-noise it, \(x_{t_{r+1}}=(1-t_{r+1})\hat x_0+t_{r+1}\varepsilon\) with fresh noise \(\varepsilon\). Consistency models follow the same pattern.
• Flow-map methods skip the clean prediction and learn the transition \(x_{t_{r+1}}=G_\theta(x_{t_r},t_r,t_{r+1})\) directly. Phased DMD [4] is the DMD-based example; like Phased Consistency Models [10] for consistency distillation, it distills phase by phase. Its training target needs a different derivation, which my notes on what DMD and Phased DMD estimate[5] cover.

No single schedule is best. We reran both baselines with shifted versions of the uniform grid \([1, 0.75, 0.5, 0.25]\) (Table 2c of our paper [1]): for flow-renoise at 4 NFE, a shift of \(\gamma=2\) is best without guidance and \(\gamma=1\) with guidance. Each candidate costs a full training run.

3. Training

We train with distribution matching distillation (DMD) [2][3], replacing only the rollout by (2). DMD matches the distribution of generator outputs to the teacher's using two flow models: a frozen real model (the teacher, SiT-XL/2 [7]) and a fake model trained on the generator's own samples. The generator, a DiT-XL/2 network [8], is pushed towards where the teacher would denoise its output and away from where the fake model would. Multi-step training also needs self-rollout: to train step \(r\), run the first \(r\) steps without gradients, then take step \(r\) with gradients, so the generator trains on its own intermediate states.

def rollout_data_space(gen, z, c, r, track_grad):
    x_hat = torch.zeros_like(z)
    for i in range(r + 1):
        with torch.set_grad_enabled(track_grad and i == r):   # gradients only at step r
            x_hat = z - gen(x_hat, z, c, i)                   # eq. (3)
    return x_hat

def dmd_step(gen, v_real, v_fake, data, r, rollout_fn, gen_step=True):
    x, c = data
    z, eps = torch.randn_like(x), torch.randn_like(x)
    tau = torch.distributions.Beta(2.5, 1.0).sample((x.shape[0],)).to(x.device)
    tau4 = tau[:, None, None, None]
    x_hat = rollout_fn(gen, z, c, r, track_grad=gen_step)
    x_tau = (1 - tau4) * x_hat + tau4 * eps                   # noise the prediction
    if not gen_step:                                          # fake-model update
        return F.mse_loss(v_fake(x_tau, tau, c, r), eps - x_hat)
    with torch.no_grad():                                     # fixed direction, no gradient
        direction = v_real(x_tau, tau, c) - v_fake(x_tau, tau, c, r)
    return (x_hat * direction).mean()                         # eq. (4) with w(tau)=1

4. Results

We compare all samplers on class-conditional ImageNet 256×256 with the same teacher and DMD recipe, so only the iteration differs. FID (Fréchet Inception Distance) is lower for better samples. Baselines use the uniform schedule, and "oracle" is the best of three schedules for flow-map.

• Data-space iteration has the lowest FID at every step count, with and without guidance, and needs no schedule search. Flow-renoise gets worse from 3 to 4 steps; ours keeps improving or stays level.
• At 4 NFE it matches or beats the oracle flow-map: 2.46 against 2.48 without guidance (parity, std < 0.05 over three runs) and 2.29 against 2.37 with guidance.
• One caveat. At 1 NFE data-space iteration and flow-renoise compute the same function, so that gap comes from optimization, which we attribute to training dynamics. Perhaps iterating on the image prediction makes different step counts more compatible with each other when they are trained together, but this is a hypothesis we have not tested.

5. What the steps look like

FID scores only the final image, so Figure 1 at the top shows the output of every step. Data-space iteration keeps the composition from the first step and sharpens details afterwards; flow-renoise changes pose, background and viewpoint between steps. The evidence that structure persists is qualitative. Measured as mean pixel-space L2 change between consecutive steps, data-space iteration changes about 28% less in the first transition but about 20% more in the last, so "changes less everywhere" is not what the numbers say.

6. Limits and related ideas

• The saving is in training, not sampling. All three samplers use up to four evaluations at inference; data-space iteration removes the schedule search. We report no extra computation or memory over a matched DMD baseline and did not measure wall-clock time.
• Self-rollout training is required. That is free if the distillation recipe already uses it, and a disadvantage otherwise.
• The evaluation is narrow. Our main experiments are class-conditional ImageNet 256×256 with FID. On unconditional CIFAR-10 the method beats DMD at 2–4 NFE but is worse at 1 NFE (12.20 against 11.21). Text-to-image and other modalities are future work, as is efficient training for ODE-based distillation objectives with paired data and noise.

If the generator is trained on its own rollouts, nothing requires those rollouts to pass through noisy states. Dropping the schedule removes one hyperparameter and, in our controlled setting, matches or beats the tuned baselines. Text-to-image models and streaming video generation, where self-rollout is already central, remain to be tested.

@article{lin2026fewstep,
  title   = {Few-Step Generation via Data-Space Iteration},
  author  = {Lin, Shanchuan and Peng, Yansong and Wang, Fu-Yun and Fan, Haoqi},
  journal = {arXiv preprint arXiv:2610.12102},
  year    = {2026}
}