NTM: Replacing Gaussian Assumptions with Normalizing Flows, Four-Step Image Generation Retains Exact Likelihood
In a paper accepted at NeurIPS 2026, the Apple Machine Learning Research team proposes Normalizing Trajectory Models (NTM), which replaces the single Gaussian distribution of each denoising step in diffusion models with conditional normalizing flows. This allows four-step sampling to match or surpass strong baselines while retaining exact likelihood over the generation trajectory, thereby supporting training strategies dependent on likelihood such as self-distillation.
Diffusion models decompose sampling into many small Gaussian denoising steps; this approximation holds when there are ample steps, but when generation is compressed into a few coarse transitions, the distribution at each step no longer approximates a Gaussian. Existing few-step methods—distillation, consistency training, adversarial objectives—adapt to this reality by bypassing step-by-step denoising, at the cost of abandoning the exact likelihood framework, which limits training flexibility and the ability to directly model the generative distribution.
In a paper accepted at NeurIPS 2026, the Apple Machine Learning Research team proposes Normalizing Trajectory Models (NTM). The core idea is to no longer assume that each denoising step is a single Gaussian, but instead to model each reverse step as an expressive conditional normalizing flow. Architecturally, NTM uses shallow invertible transformation blocks within each step to handle local distribution transformations, while using a deep parallel predictor across the trajectory to capture global structure, forming an end-to-end trainable network that can be trained from scratch or initialized from pretrained flow-matching models.
Exact trajectory likelihood brings a direct consequence: the score function induced by the model itself can be used for self-distillation—a lightweight denoiser trained on the score signal generated by the NTM itself can produce high-quality samples in four steps. The authors report that on text-to-image benchmarks, NTM matches or outperforms strong image generation baselines in four-step sampling, while uniquely retaining exact likelihood over the generation trajectory. It should be noted that these performance claims come from author-reported results in the paper's abstract; specific benchmark numbers, ablation studies, and failure modes are contained in the full text, which has not been read in its entirety here. NTM provides a concrete architectural solution to the speed-likelihood contradiction, but the magnitude of its actual advantages awaits confirmation through complete experimental evidence and independent reproduction.