Trajectory Models for Few-Step Diffusion

Existing diffusion models handle sampling by splitting it into many small Gaussian denoising steps, but reducing the generation stage to a few coarse transitions breaks that assumption. Existing few-step approaches have boosted speed through distillation, consistency training, and adversarial objectives, but in the process they had to give up the likelihood framework.

Normalizing Trajectory Models (NTM) turns each reverse step into an expressive conditional normalizing flow and trains it with exact likelihood. The architecture is built as follows:

• Each step contains shallow invertible blocks,
• A deep parallel predictor is attached across the whole trajectory,
• Both training from scratch and initialization from pretrained flow-matching are supported.

The trajectory likelihood obtained this way leads to self-distillation. A lightweight denoiser trained on the model's score produces high-quality samples in 4 steps, and on text-to-image benchmarks it matches or exceeds strong image generation baselines while still preserving the exact likelihood of the generation trajectory.