Normalizing Trajectory Models (NTM)
- Field: Computer Vision
- Authors: Jiatao Gu, Tianrong Chen, Ying Shen
- Published: 2025-05-07
- arXiv: 2505.05129
- Replaces the many-small-Gaussian-steps assumption with per-step conditional normalizing flows supporting exact likelihood training.
- Shallow invertible blocks per step + deep parallel predictor across the trajectory; trainable from scratch or from pretrained flow-matching models.
- Exact trajectory likelihood enables self-distillation for few-step sampling.
- Achieves high-quality text-to-image results in as few as four steps without giving up the likelihood framework.
Summary
Diffusion-based models decompose sampling into many small Gaussian denoising steps — an assumption that breaks down when generation is compressed to a few coarse transitions. Existing few-step methods address this through distillation, consistency training, or adversarial objectives, but sacrifice the likelihood framework in the process.
Normalizing Trajectory Models (NTM) instead models each reverse step as an expressive conditional normalizing flow with exact likelihood training. Architecturally, NTM combines shallow invertible blocks within each step with a deep parallel predictor across the trajectory, forming an end-to-end network that can be trained from scratch or initialized from pretrained flow-matching models.
The exact trajectory likelihood further enables self-distillation: a lightweight denoiser trained on the model's own scores produces high-quality samples within four steps. On text-to-image benchmarks, NTM matches or exceeds strong image generation baselines with only four-step sampling while uniquely preserving exact likelihood over the generation trajectory.