PaperlayerView exact revision on arXiv

Source-linked AI summary

Normalizing Trajectory Models

Jiatao Gu, Tianrong Chen, Ying Shen, David Berthelot, Shuangfei Zhai, Josh Susskind

arXiv:2605.08078v2cs.CVcs.LG

TL;DR

Few-step generation makes diffusion’s Gaussian reverse-step approximation inadequate, while existing compression methods sacrifice exact likelihood. NTM replaces each reverse step with a conditional normalizing flow, matching or outperforming strong baselines in four-step text-to-image generation while retaining exact trajectory likelihood.

  • Problem

    Few-step generation makes diffusion’s Gaussian reverse-step approximation inaccurate because large transitions can have complex, multimodal reverse conditionals.

  • Method

    NTM models each reverse conditional with an invertible transporter and Gaussian predictor, enabling exact likelihood training and score-based denoising distilled into four-step generation.

  • Results

    On text-to-image benchmarks, NTM matches or outperforms strong diffusion baselines and prior normalizing-flow models using only four sampling steps.

  • Takeaways & Limitations

    NTM combines few-step image generation with exact likelihood over the generative trajectory, a property absent from the compared few-step approaches.

  • Takeaways & Limitations

    NTM’s single-step configuration produces severely degraded outputs, leaving exact-likelihood generation with fewer steps an open challenge.

Abstract

from arXiv · show

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.

1 Introduction

Normalizing Trajectory Models replace Gaussian reverse transitions with expressive conditional normalizing flows, retaining exact trajectory likelihood for few-step generation. The framework also supports pretrained-model initialization and score-based denoising across timesteps.

  • Motivation: Diffusion models rely on many small Gaussian denoising steps, an approximation that becomes problematic as generation uses fewer coarse transitions.The Gaussian approximation is accurate when step sizes are small.
  • Motivation: Distillation and consistency methods reduce sampling steps but provide no tractable density over the generative trajectory, while DDGAN uses adversarial training.Existing efficiency approaches therefore sacrifice the likelihood framework.
  • Method: NTM models the non-Gaussian reverse conditional p(x_s | x_t) with an invertible transporter and Gaussian predictor, yielding exact log-likelihood through change of variables.The transporter preserves dimensionality and invertibility while making the reverse conditional simpler to model.
  • Training and initialization: NTM can train from scratch on stochastic forward trajectories or initialize from pretrained diffusion or flow-matching models while preserving pretrained quality at initialization.Initialization sets the transporter to identity; the contribution summary also specifies zero-initialized scale correction.
  • Trajectory denoising: Exact trajectory likelihood enables score-based denoising that jointly denoises all timesteps by exploiting the noisy Markov forward trajectory.The gradient of the NTM loss provides a joint score for trajectory denoising.

2 Preliminaries

Flow matching and diffusion define related forward marginals but use different reverse-generation approximations: an ODE velocity field versus Gaussian transitions. These approximations become inadequate in few-step generation because large intervals induce complex, multimodal reverse mappings, motivating trajectory modeling with expressive normalizing flows.

  • Flow matching and diffusion: Flow matching interpolates clean data with Gaussian noise and generates samples by integrating a learned velocity-field ODE from noise to data.The interpolation is xt = (1 − t)x0 + tϵ, with ϵ ∼ N(0, I).
  • Flow matching and diffusion: Diffusion models can share flow-matching marginals but use a stochastic forward process whose discretized reverse transitions are Gaussian kernels.The reverse kernel is pθ(xs | xt) = N(µθ(xt, t, s), σ2(t, s)I).
  • Few-step limitation: Few-step generation fails because flow matching assumes locally linear velocity and diffusion assumes Gaussian reverse conditionals across increasingly large time intervals.With many steps, these approximations are accurate; with few steps, each transition must capture a more complex mapping.
  • Trajectory construction: A Markovian forward trajectory preserves the prescribed marginals at every timestep and yields a tractable joint distribution with Gaussian reverse conditionals when conditioned on clean data.Sequential transitions produce a correlated path from near-clean to near-noise.
  • Gaussian approximation: The marginal reverse conditional is a mixture of Gaussians over possible clean images, so a single Gaussian becomes severely inaccurate when few steps span large intervals.The conditioned posterior p(xs | xt, x0) is Gaussian, but integrating over complex and potentially multimodal p(x0 | xt) produces the marginal mixture.
  • Normalizing flows: Normalizing flows provide invertible data-to-latent mappings with exact likelihoods, while STARFlow’s deep-shallow architecture forms NTM’s architectural foundation.Autoregressive and coupling flows provide tractable Jacobians; STARFlow combines a deep block with lightweight refinement blocks.

3 Normalizing Trajectory Models

Normalizing Trajectory Models model each denoising conditional with an expressive normalizing flow, converting representation learning into exact trajectory-likelihood optimization. The framework combines latent-space prediction, parallel trajectory architecture, flexible training schedules, pretrained initialization, and optional denoising acceleration.

  • Core formulation: NTM models each conditional p(xs | xt) as a normalizing flow with exact log-likelihood, supporting training from scratch, pretrained initialization, and learned-denoiser acceleration.The framework explicitly targets the full conditional distribution at every denoising step.
  • Latent-space model: NTM predicts in latent u-space: fT maps xs and xt to us and ut, while fP generates ûs from ut, Gaussian noise z, and optional conditioning y.The predictor uses a Gaussian latent conditional, while the transporter supplies the expressive mapping to and from image space.
  • Core formulation: An invertible, same-dimensional transporter fT turns the representation objective into exact log-likelihood optimization through the change-of-variables formula.Unlike a compressive encoder, fT preserves dimensionality and contributes its Jacobian determinant to likelihood training.
  • Architecture: The deep Transformer predictor uses non-causal full-attention coupling across trajectory steps while processing spatial positions in parallel.It predicts µP(ut, t, s, y) and σP(ut, t, s, y) for each denoising step.
  • Training and sampling: NTM supports end-to-end or pair-wise training, and sampling sequentially applies the predictor before inverting the transporter to produce the final x-space sample.A single model can support different step counts, including predefined choices such as {4, 8, 16}, without retraining.
  • Acceleration: A lightweight denoiser gϕ can amortize trajectory score-based self-refinement into a single forward pass, reducing inference costs beyond standard sampling.Standard NTM sampling otherwise requires sequential predictor steps, autoregressive decoding, and test-time backpropagation for trajectory denoising.

4 Experiments

Experiments show that NTM delivers competitive text-to-image and ImageNet generation with four denoising steps while retaining exact-likelihood training. Ablations identify the transporter, auxiliary loss, and learned denoiser as key to quality, stability, and speed.

  • Implementation: NTM uses four denoising steps, a 24-layer full-attention predictor, and two four-layer TarFlow-style transporter blocks, trained with batch size 1024 on 64 H100 GPUs.Models use bfloat16, AdamW, FSDP, and 10% CFG dropout; finetuning applies residual parameterization with λ=2.5 MSE auxiliary loss.
  • Text-to-image generation: 0.82 GenEval and 79.64 DPG-Bench show that four-step NTM trained from scratch at 256×256 outperforms STARFlow and matches stronger diffusion baselines.STARFlow obtains 0.56 GenEval with 256 autoregressive steps, while NTM uses only four steps.
  • Class-conditional ImageNet: 2.80 FID with 16 steps and 3.83 FID with 4 steps make NTM competitive with STARFlow’s 2.67 FID using 256 autoregressive steps on ImageNet 256×256.The comparison uses only exact NLL training, without adversarial or perceptual distribution-level losses.
  • Finetuning: 0.76 GenEval and 83.38 DPG-Bench demonstrate that finetuned NTM scales to 512×512 through pretrained initialization, while position and attribute binding remain challenging.The finetuned model preserves pretrained FLUX quality and prompt adherence in four steps, with high-resolution detail, text rendering, and varied styles.
  • Denoising strategies: The learned denoiser provides an approximately 9× speedup while preserving score-refinement fidelity, with LPIPS 0.121 from a single forward pass.It amortizes iterative trajectory score denoising into one forward pass.
  • Ablations: T=4 offers the best finetuning quality–speed trade-off, the transporter prevents severe blur at four steps, and the auxiliary loss prevents divergence and catastrophic forgetting.The transporter recovers 50-step quality in four steps; mean alignment anchors the predictor to the pretrained velocity field.

5 Discussion

NTM occupies a middle ground between single-step normalizing flows and multi-step flow matching by modeling each denoising transition as a conditional flow. Its discussion identifies T=4–8 as a practical sweet spot, while T=1 exposes capacity and training limitations.

  • NTM’s Position: NTM models each intermediate conditional p(x_s | x_t) along a T-step denoising trajectory as a normalizing flow.This places NTM between STARFlow’s single-step autoregressive flow and flow matching’s many-step ODE integration.
  • Architectural Tradeoff: NTM distributes capacity across denoising steps, pairing a shallow transporter with a deep trajectory-level predictor.Each transporter uses 2 blocks × 4 layers, handling local non-Gaussian residuals while the predictor reasons across timestep levels.
  • Single-Step Limitation: T=1 produces severely degraded outputs because the shallow transporter must capture the data distribution’s entire non-Gaussian structure.In this setting, the predictor reduces to a single-step Gaussian coupling, creating a fundamental capacity constraint rather than a training issue.
  • Single-Step Limitation: T=1 finetuning also changes training dynamics because the mean-alignment auxiliary loss was designed for a multi-step denoising trajectory.Collapsing the trajectory to one step undermines the loss’s intended anchoring role.
  • Operating Regime: T=4–8 is NTM’s sweet spot, balancing distributed non-Gaussian modeling with efficient cross-timestep reasoning in parallel.The architecture supports a spectrum of deeper transporters with fewer steps or shallower transporters with more steps, trading sequential computation against generation quality.

6 Related Work

Prior work addresses scalable image flows, non-Gaussian diffusion transitions, few-step distillation, score-based refinement, and trajectory-level mappings. NTM combines these directions through conditional flow modeling and trajectory-level denoising while retaining exact likelihood.

  • Normalizing flows for image generation: Normalizing flows learn invertible mappings with exact log-likelihood, while TarFlow improves their scalability to high-resolution image generation.Conditioning on x_t makes each per-step flow simpler than the full marginal and reduces the number of required blocks.
  • Non-Gaussian reverse processes: DDGAN models non-Gaussian diffusion conditionals with stepwise discriminators, but lacks tractable density and is difficult to scale.The passage also identifies mode-seeking behavior as a limitation of GAN-based reverse processes.
  • Few-step generation and distillation: Few-step methods use progressive distillation, consistency mappings, or distribution-matching objectives to reduce sampling steps while improving generation quality.Examples include progressive distillation, consistency models, DMD, and latent consistency models.
  • Score-based denoising and refinement: TarFlow and STARFlow apply normalizing-flow likelihood gradients as score-based denoisers for independent per-sample refinement.STARFlow extends TarFlow’s noisy refinement procedure to latent-space generation.
  • Score-based denoising and refinement: NTM generalizes independent sample refinement to trajectory-level denoising across the generation path.This distinction follows the passage’s contrast between prior per-sample denoising and NTM’s trajectory-level formulation.
  • Trajectory-level modeling and flow maps: Consistency models, FlowMaps, and mean flows model relationships across multiple trajectory points through endpoint projections, arbitrary-time mappings, or one-step generators.These methods differ in whether they target the clean endpoint, arbitrary pairs of probability-flow ODE times, or a one-step flow-matching generator.

7 Conclusion

NTM models each reverse conditional as a normalizing flow, enabling exact log-likelihood training and support for scratch training or pretrained-model finetuning. Its trajectory likelihood enables score-based denoising distilled into a four-step sampler that outperforms prior normalizing flow models and matches strong diffusion baselines on text-to-image benchmarks while retaining exact likelihood.

  • Framework: NTM models each reverse conditional as a normalizing flow via an invertible transporter and Gaussian predictor, yielding exact log-likelihood training.The framework combines an invertible transporter with a Gaussian predictor for each reverse conditional.
  • Training: NTM supports training from scratch and finetuning from pretrained models.Both initialization modes are explicitly supported.
  • Distillation: Its trajectory likelihood enables score-based denoising that can be distilled into a four-step sampler.The denoising procedure uses the model's trajectory likelihood and produces a four-step sampler.
  • Results: On text-to-image benchmarks, NTM significantly outperforms prior normalizing flow models and matches strong diffusion baselines with only 4 steps while retaining exact likelihood.Exact likelihood over the generative trajectory is retained alongside competitive four-step sampling.

A Theoretical Analysis … B Algorithm Pseudocode

NTM’s theoretical analysis establishes each reverse transition as an exact-likelihood conditional normalizing flow, decomposes it into Gaussian denoising and spatial transformation, and specifies training and sampling procedures. It also derives marginal-preserving forward transitions, reverse posterior coefficients, and covariance-weighted self-refinement.

  • A Theoretical Analysis: Each NTM reverse transition is an invertible conditional normalizing flow whose exact likelihood follows from the change-of-variables formula.The transition composes a spatial autoregressive transporter with a trajectory-conditioned affine predictor.
  • A.1 NTM as a Conditional Normalizing Flow: The transporter maps x-space to u-space through invertible autoregressive blocks, while the predictor applies diagonal affine coupling conditioned on the noisier representation.The transporter uses TarFlow-style causal autoregressive blocks with alternating scan directions.
  • A.2 Decomposition: Gaussian Denoising + Spatial Flow: Without the transporter, NTM reduces to diagonal Gaussian denoising and, with fixed scale, recovers the standard diffusion or flow-matching weighted MSE objective.The learned transporter instead induces non-Gaussian x-space distributions through nonlinear invertible warping and Jacobian reweighting.
  • A.3 Effect of the FM Auxiliary Loss: The FM auxiliary loss anchors NTM’s predictor mean to the pretrained solution, stabilizing u-space while allowing scale and transporter components to model non-Gaussian structure.The coefficient is annealed from full strength before decaying so NLL can further tune the mean.
  • A.4 Forward Transition Preserves Marginals: NTM’s forward transition preserves the prescribed marginals, yielding xt ∼ N((1−t)x0, t2I) from xs ∼ N((1−s)x0, s2I).The result follows because the transformed forward variable remains Gaussian with the target mean and variance.
  • A.5 Reverse Posterior Coefficients: The reverse posterior q(xs | xt, x0) has a Gaussian form with mean A(t, s)xt + B(t, s)x0 and covariance C(t, s)2I.For finetuning, x0 is replaced by the pretrained velocity model’s predicted clean sample.
  • A.6 Trajectory Covariance Matrix: The self-refinement update uses the trajectory covariance matrix to propagate correlated corrections across timesteps rather than correcting each timestep independently.The covariance couples gradients according to the noise correlation structure.
  • B Algorithm Pseudocode: Algorithm 1 trains NTM by sampling a noisy forward trajectory, applying transporter and predictor likelihood components, and reverses sampling through predictor steps followed by transporter inversion.Sampling optionally applies classifier-free guidance, self-refinement, and a learned denoiser before VAE decoding.

B.1 Training … C.2 Training Hyperparameters

NTM uses a parallel Transformer predictor and shallow autoregressive transporter blocks to model trajectory steps, with a pretrained flow-matching initialization available. Its implementation also specifies trajectory self-refinement and training hyperparameters through an algorithm and tabulated configurations.

  • C.1 Model Architecture: Architectural specifications for the NTM models are summarized in Table 4.The cited passage identifies Table 4 as the source of the experimental model specifications.
  • C.1 Model Architecture: The from-scratch predictor is a non-causal Transformer that processes all timestep levels in parallel and predicts per-step coupling parameters from u-space representations and text embeddings.Timestep conditioning uses additive sinusoidal embeddings.
  • C.1 Model Architecture: In the finetuned setting, the predictor wraps a pretrained FLUX.2 flow-matching backbone and feeds its last hidden states into a zero-initialized output projection.The projection outputs residual corrections, and zero initialization reproduces the pretrained Gaussian posterior exactly.
  • B.3 Trajectory Self-Refinement: Trajectory self-refinement computes the NTM trajectory likelihood gradient with respect to a generated trajectory and optionally applies percentile-based gradient clipping.The procedure uses a frozen NTM model.
  • B.3 Trajectory Self-Refinement: The self-refinement correction weights gradients by trajectory covariance, couples updates across timesteps, and then normalizes the result to the clean domain.The update is ˆx ←ˆx −S g, followed by ˆx ←ˆx / (1 −t).
  • C.1 Model Architecture: Each transporter block is a TarFlow-style causal autoregressive flow with 2 Transformer layers and alternating identity and flip permutations for spatial mixing.The flip permutations scan left-to-right and right-to-left.
  • C.1 Model Architecture: At t ≈1, transporter blocks are skipped as identity transforms because the input is nearly isotropic Gaussian and spatial autoregressive coupling is uninformative.This condition applies at the highest noise level.
  • C.2 Training Hyperparameters: Training hyperparameters are listed in Table 5.The supplied passage provides the table reference but no individual hyperparameter values.

C.3 Denoiser Architecture … D.2 DPG-Bench

The paper uses a lightweight Transformer denoiser, shifted timestep scheduling, and coupling-parameter guidance, then evaluates text-to-image alignment with GenEval and DPG-Bench.

  • C.3 Denoiser Architecture: The denoiser is a lightweight Transformer that maps the predictor output at the cleanest level to a clean image in one forward pass.Because the trajectory is Markov, the cleanest-level predictor output contains all information needed for deterministic prediction.
  • C.3 Denoiser Architecture: The denoiser uses 2D rotary spatial embeddings, full non-causal attention, text embeddings concatenated to the input sequence, and patch-space image output.These components define the denoiser’s spatial processing, conditioning, and output representation.
  • C.3 Denoiser Architecture: After NTM converges, its frozen trajectory score denoising generates targets for training the denoiser with MSE loss.The denoiser is trained only after the main NTM model has converged.
  • C.4 Timestep Schedule: NTM uses a shifted timestep schedule that adapts to input sequence length and samples tmin per example from a uniform interval for noise-level robustness.The schedule is formed from T denoising steps, shifted using a sequence-length-dependent parameter based on the number of patches.
  • C.5 Classifier-Free Guidance: At inference, logits-guided classifier-free guidance operates on Gaussian coupling parameters, adjusting both mean and scale rather than only the mean.The formulation applies (1+w) log pc −w log pu in u-space, with the guided parameters clipped to [0, 1].
  • D Evaluation Benchmarks: GenEval measures fine-grained compositional text-to-image generation across six categories, including objects, counting, colors, position, and color attribution.Its overall score averages object-detection accuracy across all six tasks, using four generated images per prompt.
  • D.1 GenEval: GenEval verifies specified objects, attributes, and relations with an object-detection model and reports average detection rate over four images per prompt.The overall score is the average accuracy across all six tasks.
  • D.2 DPG-Bench: DPG-Bench evaluates alignment to dense, paragraph-length prompts describing complex scenes, using BLIP-2 VQA questions over semantic prompt elements.Its categories cover attributes, entities, global coherence, other requirements, and relations, with the overall score averaged across L1 categories as a percentage.

D.3 Class-Conditional ImageNet

On class-conditional ImageNet 256×256, NTM is trained from scratch in the FAE latent space and achieves competitive FID with far fewer sequential steps than prior normalizing flows. It does so using only exact NLL training, without distribution-level losses.

  • Evaluation setup: NTM is evaluated from scratch on class-conditional ImageNet 256×256 in FAE latents with 16× spatial compression and 32-dimensional representations.The evaluation reports FID-50K across different generation step counts.
  • Results: 2.80 FID with 16 steps is comparable to STARFlow (FAE) at 2.67 FID, which requires 256 autoregressive steps.This demonstrates competitive results with dramatically fewer sequential generation steps.
  • Training objective: NTM uses only the exact NLL training objective, without adversarial, perceptual, or other distribution-level losses.The comparison emphasizes that competitive few-step performance is retained within the normalizing-flow likelihood framework.

E Additional Qualitative Results · F Broader Impact

NTM provides additional qualitative samples from models trained from scratch and fine-tuned from flow matching under identical prompts. Its efficient, high-quality few-step generation and exact likelihood offer benefits while raising concerns about misuse and motivating more accountable model development.

  • E Additional Qualitative Results: The additional samples demonstrate NTM’s qualitative image-generation behavior beyond the main reported results.The paper explicitly identifies Figure 10 as additional samples from trained NTM models.
  • E Additional Qualitative Results: Figure 10 presents additional examples from NTM models trained from scratch and fine-tuned from flow matching.The two model variants are shown side by side.
  • E Additional Qualitative Results: The qualitative comparison uses the same text prompts for both NTM training variants.This controls the prompts while contrasting scratch training with flow-matching fine-tuning.
  • F Broader Impact: NTM advances efficient image generation by combining high-quality few-step sampling with exact likelihood.This summarizes the paper’s stated contribution in the broader-impact discussion.
  • F Broader Impact: Improved generative models can support creative tools, data augmentation, and scientific visualization.These are the beneficial applications explicitly identified by the authors.
  • F Broader Impact: Generative models also raise concerns about misuse for producing misleading or harmful content.The authors frame these risks alongside the benefits of improved generative models.
  • F Broader Impact: Developing models with exact likelihood, rather than implicit or adversarial formulations, is presented as a step toward more responsible development.The supplied passage truncates the authors’ sentence after “more con,” so the claim is limited to the stated direction.
Loading 2605.08078v2…