Few-Step Generation via Data-Space Iteration

Few-Step Generation via Data-Space Iteration. By Shanchuan Lin and three others, from ByteDance Seed. Published on the eighth of October. No code repository is linked to it.

Flow matching has emerged as a scalable way to train high-quality generative models, but sampling from the learned probability flow requires many network evaluations. Distillation can bring that down to one evaluation or a few. One step often sacrifices quality, though, so few-step generation is the practical operating regime. Existing few-step methods carry out their iterative computation along the probability flow, and therefore need a fixed timestep discretization, chosen by hand.

That discretization is often picked heuristically, and it is expensive to tune. It may also be restrictive. Refinement difficulty can differ from one sample to another, and from one spatial location to another, while a fixed schedule is the same for all of them. A sample that is easy to refine and one that is hard are given the same steps at the same times.

Data-space iteration is a few-step generation framework that removes flow discretization altogether. There is no schedule of timesteps left to choose. Starting from noise, a shared generator directly refines its own prediction in data space, one iteration after another. The steps no longer correspond to points along a flow. Each one simply takes the current prediction and makes it better, using the same network every time.

Every iteration is trained to produce the best sample permitted by the generator's capacity. The formulation integrates with distribution matching distillation, or DMD, with minimal changes. That compatibility is what makes a controlled comparison possible: different iteration methods can be set against each other under matched training settings, so the iteration scheme is the thing that varies.

On class-conditional ImageNet at two hundred and fifty-six by two hundred and fifty-six, data-space iteration outperforms standard discretization baselines. It also matches or improves on variants selected through schedule search, and it does so without requiring schedule-specific training. That second comparison is the telling one. Schedule search is the expensive tuning the method was meant to make unnecessary, and data-space iteration reaches the level of the searched schedules without performing any search of its own.

The authors present data-space iteration as a simple and effective alternative to discretized flow-space iteration for fast generation. What changes is where the steps live. Instead of walking along a flow at chosen times, the generator keeps improving the sample itself, and the question of which times to choose no longer arises.