Normalizing Trajectory Models

Abstract: 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 score produces high-quality samples in four steps. On text-to-image benchmarks, NTM matches or outperforms strong image generation baselines in just four sampling steps while uniquely retaining exact likelihood over the generative trajectory.

• The paper introduces NTM, which models each reverse step via conditional normalizing flows with exact log-likelihood, ensuring accurate multimodal generation.
• The architecture combines a shallow, invertible transporter and a deep trajectory-level predictor to efficiently map noisy images to a latent space.
• Experimental results demonstrate that NTM achieves a superior balance between inference speed and image fidelity while maintaining robust calibration and mode coverage.

Normalizing Trajectory Models: Exact Likelihood Few-Step Image Generation via Conditional Normalizing Flows

Introduction and Motivation

The Normalizing Trajectory Models (NTM) framework addresses a central inefficiency in diffusion-based generative models: the reliance on numerous small-step Gaussian denoising transitions, which become a bottleneck for inference speed and model flexibility. As each step is compressed, the Gaussian approximation used for the reverse process becomes inaccurate, particularly for multimodal or heavy-tailed posteriors. Prior advances, such as step reduction via distillation or consistency training, have boosted efficiency at the cost of tractable likelihood—sacrificing important benefits like mode coverage, calibration, and principled uncertainty quantification.

NTM directly models each reverse step as an expressive conditional normalizing flow, trained with exact log-likelihood. This represents a conceptual advance over previous few-step approaches, reconciling rapid sampling, tractable density modeling, and stable training.

Methodology

Architecture

NTM decomposes each reverse conditional $p(x_s \mid x_t)$ along the denoising trajectory into two components:

• Transporter ($f_{\mathcal{T}}$): A shallow, invertible spatial normalizing flow (built from TarFlow-style autoregressive Transformer blocks), mapping each noisy image $x_t$ and its earlier counterpart $x_s$ into a latent representation (u-space).
• Predictor ($f_{\mathcal{P}}$): A deep, parallel (non-autoregressive) Transformer, operating jointly over all trajectory timesteps, predicts Gaussian coupling parameters $(\mu, \sigma)$ that parameterize a per-step conditional distribution in u-space.

The joint model forms a normalizing flow from real image space to a diagonal Gaussian prior, yielding tractable likelihood via the change-of-variables formula. Sampling proceeds stepwise from noise to data: at each step, the predictor and transporter are sequentially applied, propagating samples through only $T=4$ steps in the most efficient configuration.

Figure 1: Text-to-image generation by NTM with 4 denoising steps; models are shown both trained from scratch and fine-tuned from flow matching checkpoints.

Training Pipelines

From Scratch

NTM can be trained end-to-end from randomly initialized weights, leveraging stochastic trajectories sampled via a Markovian forward process with analytically tractable marginal distributions. The architecture allows efficient training over batches where different elements sample varying trajectory lengths (e.g., $T\!=\!4,8,16$), making the model inherently step-adaptive without retraining.

Finetuning from Pretrained Models

NTM also provides a principled recipe for initializing from pretrained flow-matching (FM) or diffusion model backbones. By setting the transporter to identity and predictor to the pretrained Gaussian posterior, NTM retains initial generation quality. A crucial auxiliary mean-alignment loss stabilizes the predictor to the pretrained solution, preventing catastrophic drift and enabling successful non-Gaussian modeling.

Figure 2: Impact of auxiliary loss for mean alignment during finetuning; lacking this stabilizer leads to divergence and loss of fidelity.

Exact Trajectory Likelihood and Score-Based Denoising

A distinctive feature is that NTM's explicit trajectory modeling enables computation of joint score functions over all trajectory steps. These gradients—computed efficiently due to NTM's tractable likelihood—enable powerful trajectory-level denoising corrections, coupling refinement across timesteps via the exact forward covariance. This procedure is then distilled into a lightweight learned denoiser, amortizing the refinement into a single forward pass, yielding substantial inference acceleration.

Experimental Evaluation

NTM’s effectiveness is established across multiple axes:

• Text-to-image benchmarks (GenEval, DPG-Bench): NTM achieves GenEval 0.82 with only 4 steps (from scratch, $256^{2}$) and 0.76 upon finetuning ($512^{2}$), surpassing prior normalizing flow models like STARFlow (0.56 with 256 AR steps) and matching or exceeding advanced diffusion baselines at comparable step counts. Attribute binding, compositionality, and spatial reasoning are strong, although the finest sub-tasks (e.g., attribute-position binding) remain challenging at higher resolutions.
• Class-conditional ImageNet: NTM delivers FID 2.80 (16 steps) and 3.83 (4 steps), closely matching STARFlow (2.67, 256 steps), all within the exact NLL framework.
• Ablations: Key findings include the necessity of multi-step trajectories for high fidelity—NTM fails at $f_{\mathcal{T}}$0 due to limited transporter capacity (see Figure 1). For finetuning, the auxiliary loss is essential for stability; models with only the NLL objective rapidly diverge.
• Qualitative Results: NTM-generated images demonstrate compositional generalization and high adherence to textual prompts at moderate and high resolutions.

  Figure 3: Denoising trajectory comparison; NTM effectively maintains image structure at $f_{\mathcal{T}}$1, unlike flow matching, which suffers severe degradation at low step counts.

  

  Figure 4: Impact of trajectory scheduling—longer $f_{\mathcal{T}}$2 improves fidelity at the cost of slower inference, with $f_{\mathcal{T}}$3 providing the best practical trade-off.

  

  Figure 5: Diagrammatic comparison of TARFlows, NTMs, and Diffusion Models, highlighting the difference in conditional/distributional structure across spatial and temporal indices.

  

  Figure 6: Additional text-to-image generation examples for both NTM-from-scratch and NTM-finetuned, demonstrating compositional control and visual quality.

Theoretical and Practical Implications

NTM provides a robust middle ground between classic normalizing flows (where all expressiveness is concentrated in a deep spatial AR flow) and ODE-based flow matching/diffusion (where per-step transitions are simple Gaussians). The architectural insight is to distribute expressiveness across lightweight spatial flows (per step) and a deep trajectory-level predictor, which leverages the conditional structure for effective multimodal modeling while retaining step-parallelism.

• Exact Normalizing Flow Per Step: Each $f_{\mathcal{T}}$4 is modeled as a full normalizing flow, not merely a Gaussian regression, enabling NTM to recover multimodality and heavy tails in the reverse process.
• Tractable Joint Likelihood: The model remains fully probabilistic, unlike GAN- or consistency-distilled few-step approaches, which gives unique support for likelihood-based calibration, sample ranking, and anomaly detection.
• Adaptivity and Scalability: By decoupling the number of refinement steps from spatial AR depth, NTM allows for natural trade-offs between generation speed and quality, with efficient amortization via the learned denoiser.
• Limitations at $f_{\mathcal{T}}$5: The failure of single-step NTM generation demonstrates the inherent trade-off between per-step transporter capacity and the number of steps required to cover complex image distributions.

Future Directions

Several avenues emerge from the NTM paradigm:

• Distributional Post-Training: Combining the exact-likelihood NTM with post-hoc adversarial or perceptual losses could further close the remaining gap with leading few-step diffusion samplers, particularly at higher resolutions.
• Architectural Innovations for $f_{\mathcal{T}}$6: Realizing practical, exact-likelihood single-step models requires new transporter designs (e.g., deep AR blocks with stepwise capacity allocation or adaptive layering).
• Broader Modality Integration: The NTM framework is structurally well-suited for extension to video (as in STARFlow-V), cross-modal synthesis, and unified generative modeling pipelines.
• Quantitative Understanding of Expressiveness: Future work may formalize the expressiveness advantages achieved by distributing flow capacity across the trajectory versus the spatial domain.

Conclusion

Normalizing Trajectory Models introduce a powerful and flexible approach to few-step generative modeling—combining the step efficiency of consistency/cascade approaches with the tractable, stable density modeling and mode coverage of normalizing flows. By coupling shallow per-step invertible flows with a deep, joint trajectory-level predictor, NTM achieves significant quality gains with drastically reduced sampling steps, especially relative to earlier flow-based approaches. The ability to train from scratch or seamlessly transfer from pretrained foundations further enhances practical applicability. Open research questions around single-step expressiveness and further post-training enhancements offer fertile ground for future progress in tractable, scalable generative modeling.