PaperlayerView exact revision on arXiv

Source-linked AI summary

Transition Matching Distillation for Fast Video Generation

Weili Nie, Julius Berner, Nanye Ma, Chao Liu, Saining Xie, Arash Vahdat

arXiv:2601.09881v2cs.CVcs.AIcs.LG

TL;DR

Large video diffusion models are difficult to use interactively because their multi-step sampling is slow and costly. TMD replaces long denoising trajectories with few-step probability transitions using a decoupled backbone and conditional flow head, and reports better fidelity and prompt adherence than existing distilled methods at comparable inference costs.

  • Problem

    Large video diffusion models require inefficient multi-step sampling, limiting their use in real-time interactive applications.

  • Method

    TMD distills diffusion models through few-step probability transitions, separating semantic backbone features from a lightweight flow head trained with transition matching and distribution matching.

  • Results

    TMD outperforms existing distilled methods in visual fidelity and prompt adherence at comparable inference costs, with Wan2.1 14B reaching 84.24 on VBench at NFE=1.38.

  • Takeaways & Limitations

    TMD provides a flexible trade-off between generation speed and video quality for distilled Wan2.1 1.3B and 14B text-to-video models.

  • Takeaways & Limitations

    The method uses conditional velocity approximations for inner flow training and leaves teacher-velocity representations for specific targets to future work.

Abstract

from arXiv · show

Large video diffusion and flow models have achieved remarkable success in high-quality video generation, but their use in real-time interactive applications remains limited due to their inefficient multi-step sampling process. In this work, we present Transition Matching Distillation (TMD), a novel framework for distilling video diffusion models into efficient few-step generators. The central idea of TMD is to match the multi-step denoising trajectory of a diffusion model with a few-step probability transition process, where each transition is modeled as a lightweight conditional flow. To enable efficient distillation, we decompose the original diffusion backbone into two components: (1) a main backbone, comprising the majority of early layers, that extracts semantic representations at each outer transition step; and (2) a flow head, consisting of the last few layers, that leverages these representations to perform multiple inner flow updates. Given a pretrained video flow model, we first introduce a flow head to the model, and adapt it into a conditional flow map. We then apply distribution matching distillation to the student model with flow head rollout in each transition step. Extensive experiments on distilling Wan2.1 1.3B and 14B text-to-video models demonstrate that TMD provides a flexible and strong trade-off between generation speed and visual quality. In particular, TMD outperforms existing distilled models under comparable inference costs in terms of visual fidelity and prompt adherence. Project page: https://research.nvidia.com/labs/genair/tmd

1. Introduction

TMD distills large video diffusion models into few-step generators by matching broad denoising transitions while retaining semantic and fine-grained refinement. Experiments on Wan2.1 models show strong speed–quality trade-offs and improved fidelity and prompt adherence.

  • Motivation and contribution: TMD approximates many-step denoising with a compact few-step probability transition process across widely separated noise levels.This enables large transition steps that match the teacher model’s distribution.
  • Architecture: The student separates a main backbone for semantic representations from a lightweight flow head for fine-grained refinement.The flow head performs multiple inner updates within each outer transition step.
  • Training: TMD uses trajectory-based flow-head pretraining followed by distribution matching distillation with flow-head rollout in each transition step.The two stages align student probability transitions with the teacher’s multi-step diffusion distribution.
  • Results: 84.24 overall VBench score is achieved by the distilled Wan2.1 14B model at near-one-step generation with NFE=1.38.The evaluation covers Wan2.1 1.3B and 14B text-to-video models on VBench and a user preference study.
  • Results: TMD achieves state-of-the-art trade-offs between speed and quality in few-step video generation.The paper reports comprehensive validation on Wan2.1 1.3B and 14B text-to-video models.

2. Background

The background contrasts continuous flow matching, transition matching, and distribution distillation as ways to shorten diffusion trajectories. Transition matching models direct probabilistic transitions, while MeanFlow provides few-step flow maps and DMD-style methods align distributions.

  • Flow and transition matching: Flow matching learns instantaneous velocities along trajectories but generally requires many small time-steps to generate data from noise.Its computational burden motivates direct transition modeling.
  • Transition matching: Transition matching directly models probabilistic transitions between discretized noisy states and can make the preceding state easy to sample from an auxiliary latent variable.Difference Transition Matching uses a deterministic update example based on the predicted latent variable.
  • Transition matching: In transition matching, the main backbone predicts features at each state while a lightweight head approximates the inner velocity.The inner rectified-flow schedule is used to predict the auxiliary variable.
  • MeanFlow: MeanFlow learns a flow map from a point at time s to any preceding point r, using average velocity over the trajectory segment.This formulation targets few-step mapping rather than many small integration steps.
  • Distribution distillation: DMD-style distribution distillation aligns student and teacher distributions through score or velocity differences, using a learned fake score when the student score is unavailable.The student output and noisy interpolates are passed to the relevant score functions.

3. Method

TMD combines a decoupled backbone–flow-head architecture with two-stage training: transition-matching pretraining followed by distribution distillation with rolled-out inner refinement. Video-specific DMD2-v changes and efficient rollout support few-step generation.

  • Training strategy: TMD first pretrains a flow head for iterative feature refinement, then applies DMD2-v distribution distillation with flow-head rollout at every transition step.The second stage aligns the student and teacher distributions during generated-sample training.
  • Decoupled architecture: The decoupled architecture uses a pretrained main backbone as a feature extractor and a lightweight flow head that iteratively predicts the transition target.The design approximates many teacher denoising steps with larger student transitions.
  • Transition matching pretraining: The flow head uses MeanFlow to become a few-step conditional flow map rather than requiring many inner updates.Its map is parameterized by the current state, two times, and main-backbone features.
  • Transition matching pretraining: The method avoids detaching main-backbone features during pretraining and targets the conditional expectation E[x1 − x | x] to remain aligned with teacher velocity predictions.These choices are presented as necessary for flexibility and teacher-aligned output.
  • Video distribution distillation: DMD2-v uses Conv3D discriminators, restricts KD warm-up to one-step distillation, and applies timestep shifting to improve video distillation.The ablations report that KD warm-up can introduce coarse artifacts in multi-step generation, while timestep shifting improves performance and prevents mode collapse.
  • Flow-head rollout: Rolling out two inner flow steps with five final DiT blocks from a 30-block DiT adds less than 17% computation when updating student parameters.Gradients are backpropagated through all inner flow steps while retaining an efficient lightweight head.

4. Experiments

Experiments evaluate TMD on Wan2.1 1.3B and 14B using VBench, visual comparisons, user preferences, and ablations. TMD improves quality and prompt adherence at comparable costs while offering a controllable speed–quality trade-off.

  • 4.1. Experiment setup: Experiments distill Wan2.1 1.3B and 14B text-to-video models and evaluate VBench scores, visual quality, prompt adherence, and user preferences.The setup uses 500k text–video pairs and effective NFE to compare inference costs.
  • 4.2. Comparison with existing methods: TMD improves visual quality and prompt adherence over DMD2-v at comparable cost, and users consistently prefer it in both one-step and two-step settings.The preference advantage is reported as especially significant for prompt alignment.
  • 4.2. Comparison with existing methods: 84.68 overall VBench score makes two-step TMD-N2H5 with NFE = 2.33 outperform all other distilled 1.3B models, including rCM at NFE = 4 with 84.43.TMD-N2H5 uses two denoising steps and five DiT blocks in its flow head.
  • 4.2. Comparison with existing methods: 84.24 overall VBench score makes one-step TMD-N4H5 with NFE = 1.38 outperform other one-step 14B distillation methods and improve over one-step rCM by +1.22.The 14B two-step TMD configuration does not outperform two-step baselines, although it outperforms four-step DMD-v.
  • 4.3. Ablation studies: Increasing effective NFE generally improves overall VBench scores, giving TMD fine-grained control over the generation speed–quality trade-off.The number of inner flow steps N and flow-head layers H control inner-flow computational cost.
  • 4.3. Ablation studies: MeanFlow pretraining and flow-head rollout improve distillation, with rollout producing faster convergence and better performance.The rollout closes the training–inference gap by backpropagating through the unrolled inner-flow trajectory.

5. Related works

Prior work accelerates video diffusion through trajectory or distribution matching, while other approaches reduce cost through autoregressive designs, system optimizations, or backbone decoupling. TMD advances this landscape with a decoupled architecture for reusing semantic representations.

  • Diffusion acceleration: Diffusion acceleration methods primarily follow trajectory matching or distribution matching to reduce the number of generation steps.Trajectory methods learn mappings along the ODE trajectory, whereas distribution methods match student and teacher distributions with fewer steps.
  • Video generation: Video-domain extensions apply these acceleration techniques to the higher-cost setting of video generation.
  • TMD positioning: TMD extends prior acceleration work by introducing a decoupled architecture that more effectively leverages teacher-learned semantic representations.
  • Efficient video generation: Autoregressive–diffusion hybrids improve efficiency by generating frames sequentially, reducing attention complexity through shorter temporal contexts.
  • Decoupled backbone: Backbone decoupling separates feature extraction from decoding and can reuse extracted features across denoising steps to accelerate sampling.

6. Conclusions

TMD addresses the inference latency of large-scale video diffusion models with a decoupled student and two-stage training strategy. Experiments on Wan2.1 show a flexible speed–quality trade-off and stronger fidelity and prompt adherence at comparable or lower budgets.

  • Conclusion: TMD addresses large-scale video diffusion latency with a main backbone for semantic features and a lightweight recurrent flow head for iterative refinement.
  • Conclusion: Its two-stage training combines transition matching pretraining with distribution-based distillation.
  • Conclusion: Experiments on Wan2.1 show that TMD flexibly balances generation speed and video quality while outperforming existing techniques at comparable or reduced computational budgets.
  • Conclusion: Future work targets a unified single-stage pipeline and integration with efficient attention or feature caching.

A. Implementation details

The implementation uses a decoupled main backbone and flow head with time-conditioned inner refinement, specialized sampling, and practical approximations for training. Additional design choices address initialization, timestep shifting, deterministic inference, and discriminator stability.

  • Flow-head conditioning: The flow head conditions on the original time embedding and an additional embedding for the time difference s − r.The summed embeddings condition adaptive normalization layers in the flow-head DiT blocks.
  • Flow-head fusion: Flow-head inputs are fused with main-backbone features through patch embedding, an AdaLN-style block, and gated interpolation.
  • Inference: TMD inference applies the main backbone at each transition and updates the sample using INNERFLOW before returning generated data.
  • Training choices: Training shifts sampled time grids and uses classifier-free guidance, condition dropout, adaptive loss normalization, and a 75% probability of setting r = s.
  • JVP approximation: A central-difference approximation replaces forward-mode JVP computation because the latter is incompatible with flash attention and FSDP.The approximation uses δ = 0.005 throughout the experiments.
  • Initialization: The method initializes the student and teacher with stop-gradient parameterization, while the baseline initially approximates E[x|x_t].
  • Sampling: The implementation uses deterministic outer transitions with independent Gaussian noise for the inner flow during multi-step inference.
  • Discriminator: The discriminator uses unpatchified teacher features at multiple layers and separate heads, which the authors found more stable than fake-score features.

A.4. VBench evaluation

VBench evaluation follows the official protocol with standardized, augmented prompts and aggregates 16 text-to-video dimensions into quality, semantic, and overall scores. Additional visual results compare early and later DMD2 training with and without KD warm-up.

  • VBench protocol: TMD is evaluated with VBench prompt lists and prompt augmentation to preserve fair comparisons across methods.The evaluation summarizes 16 T2V dimensions into quality, semantic, and overall scores.
  • KD warm-up: Figure 10 compares two-step DMD2 distillation with and without KD warm-up at training iterations 0 and 1000 using first, middle, and last video frames.
  • Prompt processing: Extended prompts rather than the short prompts shown in the main figures are passed to the models during evaluation.
  • Configuration: The experiments use default hyperparameters unless otherwise specified.

B. Additional experiments

TMD offers a tunable performance–efficiency tradeoff: increasing effective NFE generally improves performance, while Wan2.1 trajectories are hardest to match near the high-noise regime because of large curvature.

  • Larger effective NFE generally produces better performance, giving TMD finer control over the performance–efficiency tradeoff than DMD2-v.The tradeoff varies the outer and inner steps M and N and the flow-head layers H.
  • Wan2.1 sampling trajectories exhibit extremely large curvature near t = 1, corresponding to the high-noise regime.The curvature is evaluated using a discretized version on Wan’s fixed 50-step sampling grid.
  • The high curvature near t = 1 makes it difficult for trajectory-matching methods to learn mappings along the ODE path.This observation motivates larger time-shifting values in TM-MF and DMD2-v.

B.3. Flow head conditioning

The additional experiments examine flow-head conditioning, inner-flow targets, timestep shifts, recurrence, and architecture efficiency across TMD settings. They show that recurrent inner refinement and effective NFE are central to high-quality, efficient generation.

  • Flow head conditioning: The flow head is conditioned on main-backbone features m = m(x, t) through the gating mechanism used in the main experiments.An alternative concatenation-and-projection fusion is also evaluated for Wan2.1 1.3B with M = 2.
  • Additional settings: TMD’s additional experiments cover extended prompts, visual examples, and alternative inner-flow targets while retaining the same Wan-based evaluation setting.The appendix includes short and long prompts, an alternative y := x target, and comparisons of fusion types.
  • Flow head recurrence: Recurrent flow-head refinement substantially improves quality over a single inner step when distilling Wan2.1 1.3B into a two-step generator.The non-recurrent N1H5 variant shows more artifacts and blurriness than the recurrent N4H5 setting.
  • Timestep shifting: Timestep shifting remains important for DMD2-v with at least three steps, although slightly lower tstudent shifts can produce better scores.The supplied experiments compare DMD2-v across different timestep shifts for tstudent.
  • Inference efficiency: The decoupled architecture’s per-step inference-time increase closely matches effective NFE, supporting effective NFE as an efficiency measure.The timing study uses batch size 1 and a single H100 GPU with M = 1.

B.8. Training iterations

TMD uses relatively few training iterations compared with rCM and remains competitive when TM-MF pretraining is substantially reduced, according to the supplied training-efficiency evidence.

  • Training efficiency: 3k TM-MF pretraining iterations plus 1k distribution-matching iterations are used for Wan2.1 14B with M = 2, versus 10k iterations for rCM.The TMD iterations are described as lightweight TM-MF pretraining followed by distribution matching.
  • Training efficiency: For Wan2.1 1.3B, TMD still outperforms rCM and DMD2-v with a significantly reduced number of TM-MF pretraining iterations.Table 13 compares varying TM-MF iteration counts for TMD with two-step baselines.

B.9. More visual comparison results

Additional visual comparisons present TMD outputs against the teacher and DMD2-v for Wan2.1 14B and 1.3B, alongside examples spanning varied prompts and effective NFE settings.

  • Visual comparisons: Figures 14–19 provide further visual comparisons among the 50-step teacher, DMD2-v, and TMD outputs.The comparisons cover both Wan2.1 14B and Wan2.1 1.3B models.
  • Prompt coverage: The visual examples use prompts describing scenes including a child on a swing, a person eating a burger, a giraffe, roller skating, Bigfoot, and a volcano.These prompts span human actions, animals, fantasy subjects, and dramatic natural scenes.
  • Flow head recurrence: Figure 13 compares recurrent N4H5 and non-recurrent N1H5 flow-head inference, with the latter producing more artifacts and blurriness.This comparison concerns Wan2.1 1.3B distilled into a two-step generator.
  • User study: A user preference study interface is provided as an additional evaluation artifact.The interface is described as belonging to the study in Section 4.2.
Loading 2601.09881v2…