Few-Step Generation via Data-Space Iteration

TL;DR

Few-step distilled generators (DMD, DMD2, flow-renoise, flow-map variants) iterate along the probability flow, which forces training and inference to commit to a manually chosen timestep discretization that is expensive to tune and may be a poor fit when refinement difficulty varies across samples or spatial locations. This paper replaces the flow-space iteration with data-space iteration: starting from noise, a shared generator directly refines its data-space prediction over N steps, with every iteration trained to produce the best sample permitted by its capacity — no timestep schedule involved. The framework drops into DMD with minimal changes and, on class-conditional ImageNet 256×256, matches or beats discretized flow-space baselines (including schedule-searched variants) without any schedule-specific training.

Key claims

Existing few-step methods — both flow-renoise (DMD / DMD2) and flow-map (consistency, mean-flow) variants — perform iterative computation along the probability flow and therefore require a fixed, manually chosen timestep discretization that is often selected heuristically and expensive to tune [§1, Abstract].
Fixed flow-space discretization is also restrictive when refinement difficulty differs across samples or spatial locations — a single schedule can’t adapt per-input [§1].
Data-space iteration starts from noise and refines the data-space prediction directly across iterations with a shared generator; every iteration is trained to produce the best sample its capacity permits, removing flow discretization altogether [Abstract, §3].
The formulation integrates with Distribution Matching Distillation (DMD) with minimal changes to the DMD training code, which enables a controlled comparison between iteration methods under matched training settings [§3, §4].
On class-conditional ImageNet 256×256, data-space iteration outperforms standard discretization baselines and matches or improves upon variants selected via schedule search, without requiring schedule-specific training [Abstract, §4].

Method

The paper reframes the N-step student’s forward pass: instead of x_t → x_{t'} updates along the probability flow (requiring a chosen {t_0, t_1, …, t_N} schedule), a shared generator G repeatedly refines its data-space prediction x̂_0^{(k)} = G(x̂_0^{(k-1)}, k) starting from noise. Each iteration is trained to produce the best sample the model’s capacity allows at that iteration index — supervision does not depend on a timestep variable because there isn’t one.

Training inherits DMD’s two-score (real-score + fake-score) distribution matching objective. The only change to the DMD training loop is swapping how the student generates its samples at each inner step; the loss, the score networks, and the optimizer schedule are unchanged. This matched setting is what the authors emphasize for the apples-to-apples comparison against flow-space iteration variants.

Experiments target class-conditional ImageNet 256×256.

Results

On class-conditional ImageNet 256×256, data-space iteration outperforms standard timestep-discretization baselines under matched DMD training settings [Abstract, §4].
Data-space iteration matches or improves upon flow-space variants whose timestep schedules were explicitly tuned via schedule search — without doing any schedule search itself [Abstract, §4].
Because the framework drops into DMD with no training-code changes, the authors isolate iteration method as the sole variable; prior work’s gains from schedule tuning do not survive this controlled comparison [§3, §4].

(Specific FID / IS numbers and schedule-search baselines are reported in the paper’s Table / Figure set; the arxiv abstract does not quote scalar values.)

Why it’s interesting

Data-space iteration is a drop-in replacement for the per-step state update inside DMD that eliminates one of the longest-standing pain points of the DMD/DMD2 pipeline — picking a timestep schedule. This is directly relevant to the ongoing few-step distillation thread in Diffusion Distillation: ArcFlow (ArcFlow: Unleashing 2-Step Text-to-Image Generation via High-Precision Non-Linear Flow Distillation) and TMD (Transition Matching Distillation for Fast Video Generation) also attack the discretization axis, but each inherits a flow-space iteration primitive and then tries to patch it (closed-form non-linear velocity for ArcFlow; flow-head rollout inside DMD for TMD). Data-space iteration instead removes flow-space iteration altogether while leaving the DMD outer loop intact — a different, more foundational attack on the same problem.

The author tweet notes that this came out of autoregressive streaming video with self-rollout distillation, where picking a schedule is especially painful. If the recipe transfers, data-space iteration would compose cleanly with the AR-chunk-asymmetric schedules of DiagDistill (Streaming Autoregressive Video Generation via Diagonal Distillation) and the cross-resolution DMD of RMD (Cross-Resolution Distribution Matching for Diffusion Distillation) by removing the inner per-chunk discretization that both still inherit.