Phased DMD: Few-step Distribution Matching Distillation via Score Matching within Subintervals

TL;DR

Phased DMD is a few-step distillation framework that splits the SNR range into nested subintervals and trains one expert per subinterval, progressively distilling the teacher toward higher SNR. Each phase optimizes reverse-KL at an intermediate timestep, so the computational graph depth never exceeds a single-step distillation regardless of total NFE — and the resulting student is structurally a MoE generator even when the teacher is dense. The authors derive the correct score-matching objective for the fake-score estimator inside a subinterval (clean samples are inaccessible in non-final phases) and validate on Qwen-Image-20B and Wan2.1-14B / Wan2.2-A14B with a 4-step / 2-phase configuration, reporting better diversity preservation than DMD and DMD-with-SGTS while keeping key teacher behaviors (text rendering, motion dynamics).

Key claims

Directly extending DMD to multi-step distillation increases computational-graph depth and memory linearly in the step count, degrading training stability [§1, §2.2, Fig. 1a].
Stochastic Gradient Truncation (SGTS), the prior fix for multi-step DMD, can terminate after a single step within a training iteration and thereby collapses the few-step student’s distribution toward that of a one-step student — reducing generative diversity [§1, §2.2, Fig. 1b].
Phased DMD partitions the SNR range into nested subintervals and trains one expert per phase, supervising each expert at an intermediate timestep rather than only at the final clean sample [§2.3.1, Eq. 11, Fig. 1c].
The fake-score-estimator’s standard DMD training objective is invalid in non-final phases because no clean data exists at the intermediate timestep; the paper derives an equivalent objective defined purely within the subinterval and adds a clamping function to handle the t→s singularity [§2.3.2, Eq. 12, Eq. 13].
A 1-D toy experiment shows that the subinterval-only objective yields sampling trajectories that overlap with those of a model trained with the standard full-interval objective, while a naively formulated subinterval objective is visibly biased [§2.3.2, Fig. 2].
Phased DMD produces an MoE student even from a dense teacher, and experts can share a backbone with per-phase LoRA weights, bounding the memory overhead [§2.3.1].
Phased DMD composes with SGTS: 4-step inference across 2 phases uses one gradient-recorded sampling step per phase [§3, Fig. 1d].
On Wan2.1-T2V-14B, Wan2.2-T2V-A14B and Qwen-Image-20B at 4 steps, Phased DMD reports lower DINOv3 pairwise cosine similarity and higher LPIPS pairwise distance than DMD-with-SGTS — i.e. better diversity at matched NFE [Table 2].
On Wan2.2 T2V / I2V at 4 steps, Phased DMD restores motion amplitude that DMD-with-SGTS attenuates: optical-flow and Dynamic-Degree metrics move from 65.45% / 7.87 (SGTS) to 74.55% / 9.84 on T2V and from 80.00% to 83.64% on I2V’s optical-flow rate [Table 3].
The diversity improvement on Qwen-Image-20B is marginal; authors attribute this to the base model’s intrinsically limited output diversity rather than a method limitation [§3.1].
Closest prior work TDM uses disjoint timestep intervals and lacks the subinterval-correct fake-flow objective; Phased DMD uses reverse-nested SNR subintervals and derives the correct fake-flow loss [§4].

Method

A teacher diffusion / flow model is distilled into a $k$ -step, $P$ -phase student. The student’s $P$ experts cover overlapping SNR subintervals $[t_{p+1}, 1] \supset [t_{p}, 1] \supset \cdots$ in a reverse-nested arrangement; phase $p$ trains the $p$ -th expert only, mapping samples at timestep $t_p$ to samples at timestep $t_{p+1}$ (higher SNR). Per phase the generator objective is reverse-KL at the intermediate timestep $t_{p+1}$ (Eq. 11) rather than at $t=0$ .

The fake score estimator $f_\phi$ — DMD’s running approximation of the score of the current generated distribution — cannot be trained with the standard objective because no clean samples exist at $t_p$ when $p<P-1$ . The paper derives, from denoising-score-matching equivalence, an equivalent objective that uses only samples diffused from $t_p$ to $t \in [t_p, t_{p+1}]$ (Eq. 12), then adds a clamping function to prevent the $t \to t_p$ singularity (Eq. 13). A 1-D toy experiment validates that this within-subinterval objective produces flow-trained models whose sampling trajectories match those of a model trained with the standard full-interval objective.

The fake-score estimator is re-initialized from the teacher at the start of each phase and trained independently of prior phases. Combined with SGTS the configuration becomes 4 steps / 2 phases for all reported experiments. Experts share a backbone with per-phase LoRA weights.

Results

Validated by distilling Qwen-Image-20B (T2I), Wan2.1-T2V-14B (T2I), and Wan2.2-T2V-A14B (T2I, T2V, I2V), all in 4-step / 2-phase configuration. Vanilla DMD is only feasible at the smallest config (Wan2.1-T2V-14B) due to compute cost [§3, Table 1].

Diversity at matched 4 NFE — lower DINOv3 cosine = more diverse, higher LPIPS = more diverse [Table 2]:

Wan2.1-T2V-14B: DMD 0.825 / 0.522 → SGTS 0.826 / 0.521 → Phased DMD 0.782 / 0.544. Base model 0.708 / 0.607.
Wan2.2-T2V-A14B: SGTS 0.828 / 0.447 → Phased DMD 0.768 / 0.481. Base model 0.732 / 0.531.
Qwen-Image-20B: SGTS 0.941 / 0.309 → Phased DMD 0.958 / 0.322 (DINOv3 slightly worse, LPIPS slightly better — marginal, as noted).

Motion on Wan2.2 [Table 3]: T2V optical-flow rate 65.45% (SGTS) → 74.55% (Phased), Dynamic Degree 7.87 → 9.84; I2V optical-flow 80.00% → 83.64%. Base model gets 79.55% / 9.32 on T2V and 82.27% on I2V — i.e. Phased DMD matches or exceeds the teacher on motion at 4 steps.

The paper additionally argues that because the low-SNR expert in Wan2.2’s teacher is itself trained on a low-SNR subinterval, Phased DMD’s first phase aligns naturally with it and inherits the teacher’s motion dynamics; this is the structural reason Phased DMD avoids SGTS’s motion-attenuation failure mode [§3.2].

Why it’s interesting

This is a structural fix for the diversity-loss complaint that the distillation thread keeps tripping over. The closest sibling on the wiki is Diversity-Preserved Distribution Matching Distillation for Fast Visual Synthesis (DP-DMD), which solves the same DMD reverse-KL mode-seeking problem by giving the first student step a different objective (v-prediction) than the remaining steps. Phased DMD takes the same intuition — “DMD applied uniformly across steps over-collapses modes” — and runs it along a different axis (SNR-subinterval partitioning instead of per-step objective swap), and as a side effect produces an MoE student even from a dense teacher.

It also slots cleanly next to Cross-Resolution Distribution Matching for Diffusion Distillation (RMD) and Streaming Autoregressive Video Generation via Diagonal Distillation (DiagDistill) as a third “cascade something along the trajectory” recipe: RMD cascades resolution along noise level, DiagDistill cascades step count along AR chunks, and Phased DMD cascades expert assignment along SNR. All three buy quality at fixed total NFE; whether they compose remains an open question. The mathematically-grounded subinterval score-matching derivation also addresses one of the wiki’s existing open questions about TMD, which the paper criticizes for an “incorrect fake flow training” objective.