Keeping Likelihood in Few-Step Generation: A Look at Normalizing Trajectory Models

Few-step image generation has become a crowded research area, with distillation, consistency models, and adversarial objectives all competing to compress diffusion sampling from dozens of steps to a handful. Most of these methods deliver impressive speedups, but they share a quiet cost: they abandon the probabilistic foundation that made diffusion models attractive in the first place. Once you distill a model into a noise-to-data map, you can sample quickly, but you can no longer assign a tractable density to a sample. The paper Normalizing Trajectory Models (NTM) by Gu, Chen, Shen, Berthelot, Zhai, and Susskind tackles exactly this tension. It asks whether you can have both: four-step generation and an exact likelihood over the generative trajectory.

The answer the authors arrive at is yes, and the mechanism by which they get there is worth examining carefully, because it reframes the few-step problem as a representation learning problem in disguise.

The Gaussian Bottleneck

The standard diffusion or flow-matching setup decomposes generation into many small reverse steps. Each reverse step models the conditional p(x_s | x_t) as a Gaussian whose mean is predicted by a network. The reason this works is subtle but important. When the step size is small, the true reverse conditional really is close to Gaussian, because the transition only spans a thin slice of the diffusion trajectory. The single-Gaussian assumption is a good local approximation.

The trouble starts when you try to take fewer, larger steps. As the authors put it, when a transition must span a large interval, the true reverse conditional p(x_s | x_t) becomes a mixture of Gaussians that can be multimodal and heavy-tailed. A single Gaussian simply cannot represent this. This is the Gaussian bottleneck, and it is the root cause of quality degradation in naive few-step sampling. Flow matching has the analogous problem framed differently: it assumes the velocity field is locally linear within each step, which also breaks when steps are coarse.

NTM formalizes this with a stochastic trajectory framework. Rather than treating the reverse process as a deterministic ODE, the authors construct a Markovian forward trajectory over a fixed schedule 0 = t_0 < t_1 < ... < t_T = 1 that satisfies the marginal constraint q(x_t | x_0) at every step. Each forward transition has the clean form

x_t = α_{s,t} x_s + σ_{s,t} ε

with coefficients derived directly from the noise schedule. This makes the per-step distribution an explicit, well-defined modeling target rather than an approximation that quietly degrades.

The Transporter and the Predictor

The core idea of NTM is elegant. Instead of assuming the reverse conditional is already Gaussian, the model learns a latent space where it becomes Gaussian. This is accomplished with two components:

• An invertible transporter that maps each step into a latent representation. Crucially, this is not a compressive encoder. It preserves dimensionality and remains invertible, which is what allows the change-of-variables formula to deliver exact log-likelihood.
• A Gaussian predictor that operates in this learned latent space, where the reverse conditional has been straightened into something a Gaussian can model.

Architecturally, the design combines shallow invertible blocks within each step with a deep parallel predictor across the trajectory. The result is an end-to-end network that yields exact trajectory likelihood through the change-of-variables formula. The authors draw an explicit and interesting analogy here: the framework resembles a predictor-encoder architecture in the style of BYOL or I-JEPA, but the invertibility constraint converts it from a representation learner into a normalizing flow. This bridge between self-supervised representation learning and probabilistic generative modeling is one of the more conceptually rich aspects of the paper.

A practical strength is the initialization recipe. NTM can be trained from scratch using stochastic forward trajectories, or it can be initialized from any pretrained flow-matching model by setting the transporter to identity and the predictor to the pretrained Gaussian posterior, with zero-initialized scale correction. This means you do not pay a quality penalty at initialization; you start from a working flow-matching model and refine it toward expressive few-step behavior. For anyone with existing diffusion checkpoints, this lowers the cost of adoption considerably.

Self-Distillation From the Model's Own Score

The part I find most clever is how NTM exploits its own likelihood. Because the generated trajectory is a noisy sequence drawn from a Markov forward process, the gradient of the NTM loss provides a joint score that denoises all timesteps simultaneously by exploiting their correlations. In other words, the exact likelihood is not just a passive scoring tool. It actively provides a correction signal across the entire trajectory at once, using the known Markov covariance structure.

This signal can then be distilled into a lightweight learned denoiser that runs in a single forward pass. The notable detail is that this self-distillation requires no additional training data; the model distills from its own score. This sidesteps one of the common headaches of distillation pipelines, where you need a strong teacher and often a large synthetic dataset.

On the empirical side, the headline result is GenEval. NTM achieves 0.82 on GenEval with only 4 denoising steps when trained from scratch, compared to 0.56 for the prior normalizing flow model STARFlow, which required 256 autoregressive steps. That is both a large quality gain and a dramatic reduction in sampling cost relative to the closest likelihood-based competitor. On text-to-image benchmarks, NTM is reported to match or outperform strong few-step baselines in image quality and compositional accuracy, while uniquely retaining exact likelihood.

My Take

What makes NTM interesting is not the raw GenEval number, impressive as it is, but the reframing of few-step generation as a question of where the Gaussian assumption lives. Existing methods either accept the bad single-Gaussian approximation (naive few-step sampling), discard the probabilistic model entirely (distillation, consistency, GAN), or pay a high step count to keep it (STARFlow). NTM instead says: keep the Gaussian, but only after transporting the data into a space where the Gaussian is true. The invertibility constraint is what makes this honest, because it preserves the change-of-variables bookkeeping that a normalizing flow needs.

There is a deeper inference worth drawing. The predictor-encoder analogy suggests that the transporter is implicitly learning a representation in which the generative dynamics are simple. If that representation is genuinely good, one might expect it to be useful beyond generation, for example as features for downstream tasks. The paper does not explore this, but it is a natural direction. A generative model whose internal latent geometry linearizes the reverse process could be a strong representation learner as a byproduct.

The practical payoff, as I noted in the accompanying summary, is that you get a few-step generator that is simultaneously a density model. That combination is unusual. It enables outlier and anomaly detection on generated or real samples, calibrated confidence estimates, and principled model comparison via likelihood, all from the same network you use for fast sampling.

Limitations and Open Questions

Several caveats deserve honesty. The invertible transporter adds parameters and architectural complexity relative to a plain velocity network, and exact-likelihood normalizing flows have historically been harder to scale than their unconstrained counterparts. The paper reports strong results at 256x256 from scratch and 512x512 from finetuning, but the field's frontier sits at higher resolutions, and whether the invertibility constraint remains cheap there is unproven. The quality of NTM's likelihood as an actual anomaly detector is also asserted by implication rather than benchmarked in the abstract, and flow likelihoods are known to behave counterintuitively on out-of-distribution data.

The open questions I would most like answered are whether the learned latent space transfers to representation tasks, how the exact trajectory likelihood compares numerically against continuous flow likelihoods on held-out data, and how far the step count can be pushed below four before the multimodal reverse conditional overwhelms even an expressive flow. NTM does not close the few-step problem, but it reopens a door that the field had largely walked away from: the door where speed and tractable likelihood coexist.