Normalizing Trajectory Models: Restoring Exact Likelihood to Few-Step Generation

The Core Tension in Few-Step Diffusion Generation

The ability of diffusion models to generate high-quality samples rests on a fundamental assumption: decomposing the sampling process into a large number of tiny Gaussian denoising steps. Only when each perturbation is sufficiently small does the Gaussian approximation hold. But when we push for faster generation and compress the entire sampling process into just a handful of coarse-grained transition steps, this assumption starts to break down.

The problem is that coarse-grained transitions require single steps to span much larger distributional shifts — territory where simple Gaussian denoising can no longer accurately characterize the true reverse process. This is the central tension in current few-step generation research: the trade-off between speed and fidelity.

The Cost of Existing Few-Step Approaches

The research community has explored several mainstream approaches to tackle few-step generation, each with its own trade-offs:

• Distillation: A pretrained multi-step teacher model guides a few-step student model. This works reasonably well, but it depends on a teacher model and lengthens the training pipeline.
• Consistency Training: The model is forced to remain self-consistent along its trajectory, enabling few-step or even single-step sampling.
• Adversarial Objectives: Adversarial training is introduced to improve generation quality.

All of these approaches share a common sacrifice: in compressing the number of steps, they abandon the likelihood framework. The computable, optimizable probability density properties that diffusion models originally possess are given up — stripping the model of precise likelihood evaluation and diminishing the theoretical interpretability of the generation process.

---

Distillation in few-step diffusion models takes several concrete forms, with Progressive Distillation and Consistency Distillation being two representative approaches. Progressive Distillation iteratively compresses two teacher steps into one student step, progressively halving the number of sampling steps. Consistency models define a "consistency function" that maps any point on the same denoising trajectory to the same initial state, enabling single-step or few-step sampling. The fundamental problem shared by these methods is that their training objectives are essentially regression-based (matching teacher outputs or specific points on the trajectory) rather than maximizing the log-likelihood of the data. As a result, the generation process no longer corresponds to any interpretable probabilistic model — making it difficult to apply these models to tasks requiring explicit probabilities, such as density estimation or uncertainty quantification, while also weakening the theoretical analysis of the generation process.

---

NTM: Rebuilding Each Step with Normalizing Flows

To address these challenges, researchers have proposed Normalizing Trajectory Models (NTM). The core idea is to fundamentally redesign each individual step of the reverse process.

Modeling Each Step as a Conditional Normalizing Flow

Rather than simply assuming each reverse step is a Gaussian denoising operation, NTM models every step as an expressive conditional normalizing flow. The key advantage of normalizing flows lies in their invertibility and exact likelihood computation — through a series of invertible transformations, the model can precisely track changes in probability density.

This means that even when generation is compressed to just a few coarse-grained transitions, each step retains sufficient expressive power to characterize complex distributional shifts, without relying on Gaussian approximations. More importantly, the entire model can be optimized under exact likelihood training — precisely the valuable property that existing few-step methods abandon.

---

Normalizing flows are an important branch of generative models. The core idea is to map a simple known distribution (such as a standard Gaussian) to a complex data distribution through a series of invertible and differentiable transformations. Because each transformation is invertible, the model can use the change-of-variables formula to compute the log-likelihood of data points exactly: $\log p(x) = \log p(z) + \sum_i \log |\det J_i|$, where $J_i$ is the Jacobian matrix of each transformation. This stands in sharp contrast to VAEs (which require a variational lower bound approximation) and GANs (which cannot compute likelihoods at all). The main challenge with normalizing flows is that ensuring invertibility places strong constraints on the network architecture (e.g., RealNVP's affine coupling layers, Glow's 1×1 invertible convolutions), which often limits the expressive power of individual transformations while significantly increasing parameter count and computational cost. NTM addresses this scalability bottleneck — which arises when normalizing flows are applied globally in high dimensions — by restricting flow usage to "shallow" structures that operate only within individual diffusion steps.

---

Combining Shallow Invertible Blocks with a Deep Parallel Structure

On the architectural side, NTM adopts a layered strategy: using shallow invertible blocks within each step, combined with a deep parallel structure.

The elegance of this design lies in balancing two competing requirements: the shallow invertible blocks ensure invertibility and exact likelihood computability at each step, while the deep parallel component provides sufficient modeling capacity to capture complex data features. In other words, NTM confines the "invertibility constraint" to lightweight local structures, while delegating the primary representation learning to the deep network — achieving both theoretical rigor and practical expressiveness.

Significance and Limitations

NTM's value lies in its attempt to bridge two long-incompatible goals: few-step generation and exact likelihood. Compared to distillation, consistency training, and adversarial training, NTM preserves the completeness of probabilistic density modeling — which has clear implications for applications that require density estimation, anomaly detection, or downstream probabilistic inference.

That said, based on currently available information, this remains an early-stage research contribution. The scalability of normalizing flows on high-dimensional data, training stability, and the computational overhead introduced by invertible architectures are long-standing challenges for flow-based methods. Whether NTM can match the sample quality of mainstream diffusion models on high-dimensional generation tasks like image synthesis remains to be validated with more complete experimental results.

For researchers following the frontier of generative modeling, NTM offers a new direction worth watching: rather than abandoning the likelihood framework when compressing step counts, it fundamentally redesigns how each step is modeled. This "return to likelihood" philosophy may provide a valuable new reference point for future research into efficient generative models.