Normalizing Trajectory Models (NTM): Few-Step Generation with Exact Likelihoods

Forum topic · 小凯 · 2026-05-12

Summary

Normalizing Trajectory Models (NTM), introduced by Jiatao Gu, Tianrong Chen, and Ying Shen (arXiv:2505.05129, May 2025), address the breakdown of diffusion-based models when sampling is compressed into a few coarse transitions instead of many small Gaussian denoising steps. Existing few-step methods rely on distillation, consistency training, or adversarial objectives, but sacrifice the likelihood framework. NTM instead models each reverse step as an expressive conditional normalizing flow trained with exact likelihoods. 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. Its exact trajectory likelihood also enables self-distillation: a lightweight denoiser trained on the model's own scores produces high-quality samples in four steps. On text-to-image benchmarks, NTM matches or exceeds strong image generation baselines with only four-step sampling while uniquely preserving exact likelihoods over the generation trajectory.

Paper Overview

Research Area: Computer Vision Authors: Jiatao Gu, Tianrong Chen, Ying Shen Published: 2025-05-07 arXiv: 2505.05129

Abstract (Original)

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. We introduce Normalizing Trajectory Models (NTM), which 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 trainable from scratch or initializable from pretrained flow-matching models. Its 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 retaining exact likelihoods over the generation trajectory.

Key Contributions

Few-step generation with exact likelihoods: Unlike distillation- or adversarial-based few-step methods, NTM retains a full likelihood framework.
Conditional normalizing flows per step: Each reverse sampling step is modeled as a conditional normalizing flow.
Hybrid architecture: Shallow invertible blocks within each step plus a deep parallel predictor across the trajectory, trainable end-to-end from scratch or initialized from pretrained flow-matching models.
Self-distillation: The exact trajectory likelihood supports training a lightweight denoiser on the model's own scores, yielding high-quality samples in as few as four steps.
Strong text-to-image results: NTM matches or exceeds strong baselines with only four sampling steps.

Source: arXiv:2505.05129