Few-Step Generation via Data-Space Iteration

The gist The introduction introduces data-space iteration, a few-step generation framework that removes flow discretization altogether, which provides a simple and effective alternative to discretized flow-space iteration for fast generation

Flow Matching Background

Flow matching has become a scalable pretraining paradigm for modern generative models by learning a continuous probability flow between simple noise and the data distribution A neural network vϕ is trained to predict the velocity of this flow, dxt/dt = vϕ(xt, t) To generate a sample, one draws x1 ∼ pnoise and solves this ordinary differential equation (ODE) backward from t = 1 to t = 0 An accurate numerical solution often requires tens or hundreds of network evaluations Distillation replaces this costly process with a generator that produces a sample in only a few evaluations

Existing Few-Step Generation Paradigms

Existing methods perform iterations in the probability-flow space, which are grouped into two paradigms

** Flow-renoise iteration involves predicting the data endpoint at schedule point tr, and then mixing it with noise at the next schedule point tr+1 **

** Flow-map iteration learns a direct transition between consecutive schedule points, where xbtr+1 = Gθ(xtr, tr, tr+1) Dependence on flow-time discretization is a limitation because existing methods use flow time to define progress during generation This discretization is usually treated as a fixed design choice rather than a learned component of the generator Data-space iteration removes the flow-time schedule by passing a data prediction, rather than a flow state, between steps Every refinement step is trained to produce the best sample it can, given the current estimate and the original noise This formulation removes the timestep-discretization hyperparameter from the generator’s iterative sampling process and lets the model spend all of its computation on directly improving the sample Our method can be trained with almost no modification to standard distribution matching distillation (DMD) framework, retaining the real and fake flow models, DMD objective, and self-rollout training structure The core DMD formulation remains unchanged, so our method adds no additional computation or memory relative to a matched DMD baseline Our generator output is xb(r+1) = Gθ(xb(r), z, r):= z - gθ(xb(r), z, c, i) The added parameters are often negligible, and the timestep projection MLP can be replaced with a constant t = 1 embedding Our formulation differs in two ways from fixed-point models: first, our generator is conditioned on r, so its update can change across refinement steps; second, we apply the distribution-matching objective at every step rather than only after convergence 3 Experiment Data Space Iteration requires only a small change to standard DMD training by retaining the frozen pretrained flow model vreal = vϕ and an auxiliary fake flow model vfake = vψ, which is initialized from the pretrained model and conditioned on the refinement index The generator update involves sampling a refinement index r ∈ 0,..., R-1 and noise z ∼ pnoise, first running r refinement steps without tracking gradients to obtain xb(r)0, followed by one step with gradients to produce xb(r+1)0 The DMD loss is Ldmd(θ) = Ez,r,τ,ε w(τ) sgh[vrealxb(r+1)τ, τ− vfakexb(r+1)τ, τ, ri⊤xb(r+1)0] 4 Training via Distribution Matching Distillation involves an update for the generator where xb(r+1)0 = Gθ(xb(r)0, z, r) The fake-model update is trained on noisy generator samples with Lfake(ψ) = Ez,r,τ,ε vfakexb(r+1)τ, τ, ri⊤xb(r+1)0−ε − xb(r+1)0 The evaluation on class-conditional ImageNet 256 × 256 using a pretrained SiT-XL/2 flow-matching teacher backbone shows that data-space iteration outperforms standard discretization baselines and matches or improves upon variants selected through schedule search without requiring schedule-specific training 3.4 Main Comparisons Under the standard uniform discretization, data-space iteration consistently outperforms both flowrenoise and flowmap across all reported NFE settings Even when allowing an oracle selection of the best of three discretizations, flowmap only narrows the gap and remains slightly behind 3.5 Visualization Figure 3 visualizes the differences between consecutive refinement steps, showing that data-space iteration learns to refine structures and details from its previous prediction, whereas flow-renoise exhibits large content changes across iterations Table 6 quantifies the change between consecutive rollout steps, showing that for each step r, the average L2 difference E∥x(r)0 − x(r+1)0∥2 over 50k class-balanced generated samples is lower for data-space than for flow-renoise The final results show that data-space iteration achieves stronger performance without schedule search or schedule-specific training 5 Conclusion We introduced data-space iteration for few-step generation by directly refining data predictions, which removes the need for a fixed flow-time discretization and integrates into DMD with minimal changes On class-conditional ImageNet 256 × 256, it outperforms standard uniform-discretization baselines and matches or improves upon schedule-searched variants without requiring schedulespecific training Limitations include the requirement for self-rollout training, which incurs no additional cost when self-rollout is already required but is a disadvantage otherwise Additional contents and future work focus on scaling data-space iteration to text-to-image generation and other modalities 7 Table 3 Main results reports the main results, showing that data-space iteration outperforms both flowrenoise and flowmap across all reported NFE settings 8 Table 5 Full metrics reports full evaluation results with and without classifier-free guidance, demonstrating superior performance for data-space iteration The final results demonstrate that data-space iteration achieves stronger performance without schedule search or schedule-specific training 9 Table 8 CIFAR-10 results report the final FIDs for different methods, showing that data-space iteration outperforms flowrenoise at NFE > 1 The paper concludes that data-space refinement is a simple and effective approach to fast generation> The gist. This paper introduces data-space iteration, a few-step generation framework that removes flow discretization altogether, which provides a simple and effective alternative to discretized flow-space iteration The gist. This paper introduces data-space iteration, a few-step generation framework that removes flow discretization altogether, which provides a simple and effective alternative to discretized flow-space iteration 5 Conclusion We introduced data-space iteration for few-step generation by directly refining data predictions, which removes the need for a fixed flow-time discretization and integrates into DMD with minimal changes The gist.

1. Data-space iteration removes the timestep-discretization hyperparameter from the generator’s iterative sampling process, allowing for a simple and effective alternative to discretized flow-space iteration for fast generation. This means AI systems can perform refinement directly in data space, eliminating the need for a fixed, manually chosen timestep discretization which is often expensive to tune.

2. The method enables controlled comparison between different few-step generation methods by allowing a fixed training recipe while changing the sampling space. This facilitates testing of iteration strategies without requiring schedule-specific training, such as when comparing data-space iteration against variants selected through schedule search.

3. The ability to condition the shared generator on the refinement step, as shown in Table 1b (Conditioning the generator on r), allows for a stage-appropriate correction rather than using the same update rule at every step, leading to improved performance across different evaluation settings.

4. The framework supports a flow-map iteration paradigm where maps are learned directly between consecutive schedule points, which can be beneficial when training is constrained, as it avoids the need for a fixed design choice of discretization.

5. The formulation integrates with distribution matching distillation (DMD) with minimal changes, meaning existing high-quality teacher models can be adapted to generate fast samples by simply changing the rollout process from probability-flow iteration to data-space iteration.

Flow matching has emerged as a scalable paradigm for training high-quality generative models, but sampling from the learned probability flow requires many network evaluations. Distillation can reduce this cost to one or a few evaluations; however, one-step generation often sacrifices quality, making few-step generation the practical operating regime. Existing few-step methods perform their iterative computation along the probability flow and therefore require a fixed, manually chosen timestep discretization. This discretization is often chosen heuristically and is expensive to tune; it may also be restrictive when refinement difficulty differs across samples or spatial locations. We introduce data-space iteration, a few-step generation framework that removes flow discretization altogether. Starting from noise, a shared generator directly refines its prediction in data space, with every iteration trained to produce the best sample permitted by its capacity. Our formulation integrates with distribution matching distillation (DMD) with minimal changes, enabling a controlled comparison between iteration methods under matched training settings. On class-conditional ImageNet 256x256, data-space iteration outperforms standard discretization baselines and matches or improves upon variants selected through schedule search, without requiring schedule-specific training. These results show that data-space iteration provides a simple and effective alternative to discretized flow-space iteration for fast generation.