Few-Step Generation via Data-Space Iteration

1 INTRODUCTION

• Flow matching learns a continuous probability flow between noise and data, but high-quality sampling often requires many network evaluations, motivating few-step distillation.Distillation aims to reduce inference to one or a few evaluations.Source: §1 · p. 1
• Few-step generation is a practical balance because one-step distillation often degrades quality, while a few extra evaluations can substantially improve sample quality.Four-step sampling is commonly used as a latency–fidelity balance in deployed high-resolution systems.Source: §1 · p. 2
• Existing few-step generators organize refinement around manually chosen flow-time discretizations, and self-rollout methods must fix the schedule during training, making schedule search costly.Different samples or spatial regions may favor different refinement, but a global schedule advances them together.Source: §1 · p. 2
• Data-space iteration refines predictions directly from an initial noise sample, training every step to produce the best sample its capacity permits without a flow-time discretization hyperparameter.The generator uses the current prediction and original noise to refine the sample.Source: §1 · p. 2
• The method retains DMD’s real and fake flow models, objective, and self-rollout structure, changing only the generator’s rollout from probability-flow to data-space iteration.This holds the training recipe fixed while changing the iterative sampling space.Source: §1 · p. 2
• On class-conditional ImageNet 256 × 256, data-space iteration outperforms standard discretization baselines and matches or improves upon schedule-searched variants under matched DMD settings.The comparison uses a pretrained SiT-XL/2 flow-matching teacher and does not require schedule-specific training.Source: §1 · p. 2

2.1 FLOW MATCHING BACKGROUND

• Flow matching learns a velocity field between data and noise, and generation solves its ODE backward; accurate numerical solutions often require tens or hundreds of evaluations.Distillation replaces that costly process with a generator requiring only a few evaluations.Source: §2.1 · pp. 2–3

2.2 EXISTING FEW-STEP GENERATION PARADIGMS

• Existing few-step methods iterate in probability-flow space using either flow-renoise transitions or direct flow-map transitions, both organized by a chosen schedule.For self-rollout distillation, the schedule fixes training states before training, and testing another generally requires training another model.Source: §2.2 · p. 3

2.3 FEW-STEP GENERATION VIA DATA-SPACE ITERATION

• The shared generator receives the same source noise at every refinement step, while the refinement index lets it use a different update at each stage.The index r denotes the refinement step, and R is the maximum number of network evaluations.Source: §2.3 · p. 3
• Each data-space step outputs a valid sample prediction and can use the previous prediction to improve it, without adding noise or constructing intermediate flow-time states.Sampling therefore needs no flow-time schedule.Source: §2.3 · p. 3

2.4 TRAINING VIA DISTRIBUTION MATCHING DISTILLATION

• DMD training retains the pretrained real flow model and auxiliary fake flow model, and uses the difference between their predictions as a fixed generator update direction toward the data distribution.The fake model tracks the generator distribution at each refinement step; its training alternates with generator updates on noisy generated samples.Source: §2.4 · p. 4

2.5 GENERATOR IMPLEMENTATION

• The generator is built by minimally modifying a pretrained velocity network, adding a projection for the current estimate and refinement-index conditioning while fixing the original flow-time input at t = 1.The construction preserves the pretrained parameterization; added parameters are often negligible, and the model size can decrease with a constant timestep embedding.Source: §2.5 · p. 4

2.6 CONNECTION TO FIXED-POINT MODELS

• Unlike a fixed-point model, the generator is refinement-index conditioned and distribution matching is applied at every step, so every rollout prefix up to R is encouraged to produce a valid sample without requiring convergence.An optional optimal-transport regularizer intended to encourage successive predictions to stabilize did not improve results.Source: §2.6 · p. 5

3.1 EXPERIMENTAL SETUP

• On ImageNet 256 × 256, evaluation uses latent-space FID over 50k class-balanced samples and measures sampling cost by function evaluations.The experiments use DiT-XL/2 and a pretrained SiT-XL/2 flow-matching backbone; DMD generators support up to four steps and are trained for 40 epochs.Source: §3.1 · p. 5

3.2 ABLATION STUDY ON DATA-SPACE ITERATION GENERATION

• Greater emphasis on high-noise timesteps improves DMD weighting, importance sampling outperforms loss weighting, and conditioning on refinement step r consistently improves performance.The ablation also reports that β1 = 0.0 consistently improves performance, while the optimal-transport regularizer slightly degrades FID and is omitted from the final method.Source: §3.2 · pp. 5–6, Table 1 · p. 6

3.3 ABLATION STUDY ACROSS DIFFERENT SAMPLING METHODS

• Data-space iteration beats flow-renoise under every compared DMD weighting and beats flow-map under the compared weighting settings; both data-space and flow-renoise improve with greater high-noise emphasis.Across compared CFG configurations, data-space iteration outperforms both alternatives; full-timestep CFG worsens performance, while scale 1.1 over τ ∈ [0.0, 0.6] yields the best overall performance across methods.Source: §3.3 · p. 6, Table 2 · p. 7
• Changing the flow-time discretization substantially affects flow-renoise intermediate-step FID and flow-map CFG FID, but the shift parameter does not predictably track final performance.The default flow-space baseline is the uniform trailing discretization [1.0, 0.75, 0.5, 0.25].Source: §3.3 · pp. 6–7

3.4 MAIN COMPARISONS

• Across reported NFE settings, data-space iteration beats flow-renoise and flow-map with uniform discretization, and remains ahead of schedule-searched flow-map while flow-renoise lags.The schedule-search comparison requires training a separate generator for each candidate schedule; all unguided DMD-distilled models also substantially outperform the flow-matching teacher.Source: §3.4 · pp. 7–8, Table 3 · p. 8

3.5 VISUALIZATION

• Data-space iteration refines structures and details from its previous prediction, whereas flow-renoise can discard satisfactory regions and regenerate different content.The data-space generator was trained without an optimal-transport regularizer or another explicit loss constraint on its refinement trajectory.Source: §3.5 · p. 8

Few-step distillation

• DMD is a widely adopted distribution-level distillation approach, while improvements to its scaling, parameterization, and score estimation are complementary to this paper’s focus on iteration space.Pointwise regression objectives often produce blurry samples.Source: §4 · p. 9

Optimizing timestep discretization

• Prior schedule optimization either shares one schedule across samples or predicts sample-specific schedules, but the latter is not directly compatible with self-rollout distillation.In self-rollout distillation, the schedule must be specified during training.Source: §4 · p. 9

Self-conditioning

• Unlike self-conditioning, which refines the velocity field while generation still integrates an ODE, this method iterates directly on data predictions.The related methods are cited as Chen et al. (2023), Strudel et al. (2022), and Dieleman et al. (2022).Source: §4 · p. 9

Deep equilibrium models

• A related equilibrium approach solves a fixed point in hidden representation and requires an equilibrium-specific architecture and solver, whereas this method requires minimal architectural changes.The equilibrium approach is cited as Geng et al. (2023).Source: §4 · p. 9

Looped architectures

• Looped-transformer distillation backpropagates through the full unrolled loop, with computational and memory costs that can be a practical hurdle; data-space iteration is a lightweight drop-in modification.The looped-transformer work is cited as Lin et al. (2026a).Source: §4 · p. 9

Conclusion

• On class-conditional ImageNet 256 × 256, data-space iteration outperforms uniform-discretization baselines and matches or improves upon schedule-searched variants without schedule-specific training.The method directly refines data predictions and integrates into DMD with minimal changes.Source: §5 · p. 9

Limitations

• Data-space iteration requires self-rollout training, which adds a disadvantage when self-rollout is not otherwise required.When ODE-based distillation objectives are preferred, paired data and noise samples can be used, but more efficient training approaches remain to be explored.Source: §5 · p. 9

Additional contents and future work

• The main paper focuses on ImageNet 256 × 256, provides CIFAR-10 results with a UNet in the appendix, and leaves scaling to text-to-image and other modalities for future work.Hyperparameters, architectures, and code are provided in the appendix.Source: §5 · p. 9

Disclosure

• The authors used generative AI to assist with drafting, table and figure organization, and code implementation and checking, but not to run training jobs or curate citations.They state that they reviewed AI-assisted work and take responsibility for the final content.Source: AI USE STATEMENT · p. 10

A IMPLEMENTATION DETAILS

• Table 4 summarizes the final architecture and training settings for the data-space iteration generator.The generator and fake flow model use the XL/2 configuration of the pretrained SiT backbone.Source: Table 4 · p. 15, A IMPLEMENTATION DETAILS · p. 15
• The generator and fake model use the same constant learning rate, with four fake-model updates per generator update; training uses TF32 and generator EMA decay 0.9999 for evaluation.The DMD teacher timestep is sampled from Beta(2.5, 1) with constant weighting w(τ) = 1; no weight decay or gradient clipping is used.Source: A IMPLEMENTATION DETAILS · p. 15
• During training, the refinement index is sampled as r = i mod R, and the fake velocity network receives r and is identical across iteration-method comparisons.The added r embedding is initialized from N(0, 0.02) for both the generator and fake velocity network.Source: A IMPLEMENTATION DETAILS · p. 15
• The final model has 673,527,200 parameters, fewer than the original SiT-XL/2 model’s 675,129,632 parameters.The added input projection and refinement embedding add 24,192 parameters, while replacing the timestep projection MLP with a constant t = 1 embedding decreases the parameter count by 1,626,624.Source: A IMPLEMENTATION DETAILS · p. 15
• For classifier-free-guided results, guidance scale is 1.1 over teacher timesteps [0.0, 0.6], and the fake model is warmed up for 100 iterations.The warmup is stated to prevent training from diverging when CFG is enabled.Source: A IMPLEMENTATION DETAILS · p. 15

B REPRODUCIBILITY

• Results are reproducible within one day on eight 80 GB HBM3 GPUs with a high-bandwidth interconnect; close final FIDs use three training runs and separate evaluation seeds.The reported standard deviation is below 0.05 for both data-space iteration and flow-map.Source: B REPRODUCIBILITY · p. 16

• FID is the primary evaluation metric, while sFID, IS, precision, and recall are included for reference.Source: C FULL METRICS · p. 16
• FID and the additional metrics are evaluated on 50,000 class-balanced generated samples.Source: C FULL METRICS · p. 16
• Table 5 presents the full evaluation metrics.Source: Table 5 · p. 16, C FULL METRICS · p. 16

D MORE VISUALIZATIONS

• Figure 3 places generated images in the first row and signed differences between consecutive steps beneath the later image, with zero shown as mid-gray.The same fixed normalization is used for all images.Source: Figure 3 · p. 17, D MORE VISUALIZATIONS · p. 17
• Table 6 reports average L2 differences between rollout steps over 50,000 samples, measured after VAE decoding in RGB pixel space with values normalized to [0, 1].Source: Table 6 · p. 17, D MORE VISUALIZATIONS · p. 17

E CODE

• Algorithm 1 uses data-space and flow-renoise rollout functions with a shared DMD objective; the sampled refinement index selects the gradient-tracked step.Earlier rollout steps are evaluated without gradient tracking.Source: E CODE · p. 18
• Phased DMD rolls out to the phase endpoint, samples a DMD noising time from that phase's valid interval, and re-noises the intermediate flow state before applying the objective.Source: E CODE · p. 19

• An unweighted score-based DMD update through xτ corresponds to w(τ)=(1−τ)^2/τ on bx0, while applying the score difference directly to bx0 corresponds to w(τ)=(1−τ)/τ.Source: F DMD WEIGHTING FUNCTIONS AND THEIR… · p. 20
• Both score-derived weightings diverge as τ approaches zero; restricting samples to τ≥0.02 did not prevent the corresponding ablation runs from diverging.Source: F DMD WEIGHTING FUNCTIONS AND THEIR… · p. 20
• For the linear path, w(τ)=1 applies the velocity difference directly to bx0, while w(τ)=1−τ applies it directly to bxτ.Source: F DMD WEIGHTING FUNCTIONS AND THEIR… · p. 20
• Increasing emphasis on high-noise timesteps empirically improves FID, although the w(τ)=τ and w(τ)=τ^2 weightings have no direct theoretical connection.Source: F DMD WEIGHTING FUNCTIONS AND THEIR… · p. 20
• The original DMD weight retains adaptive gradient normalization (C channels and S spatial locations); omitting it yields effective velocity weighting (1−τ)τ, while Beta(2.00, 1) and Beta(3.00, 1) correspond up to scaling to τ and τ².Beta(2.50, 1) interpolates between those distributions.Source: F DMD WEIGHTING FUNCTIONS AND THEIR… · p. 21

• The unconditional CIFAR-10 experiment uses 32×32 images and the 55M SongUNet architecture, without additional hyperparameter ablations or auxiliary techniques such as data augmentation.Source: G CIFAR-10 EXPERIMENT · p. 21
• The flow-matching teacher achieves FID-50k of 3.54 at 250 NFE, and Table 8 reports FIDs for DMD, Phased DMD, and data-space DMD.Source: G CIFAR-10 EXPERIMENT · p. 21, Table 8 · p. 21
• At 1 NFE, data-space iteration outperforms flow-renoise on ImageNet but performs slightly worse on CIFAR-10; the paper attributes the discrepancy to optimization dynamics.The analysis focuses primarily on NFE greater than 1, where observed differences are attributed to the distinct iteration strategies.Source: G CIFAR-10 EXPERIMENT · p. 21