Terminal Velocity Matching (2511.19797v1)

Abstract: We propose Terminal Velocity Matching (TVM), a generalization of flow matching that enables high-fidelity one- and few-step generative modeling. TVM models the transition between any two diffusion timesteps and regularizes its behavior at its terminal time rather than at the initial time. We prove that TVM provides an upper bound on the $2$-Wasserstein distance between data and model distributions when the model is Lipschitz continuous. However, since Diffusion Transformers lack this property, we introduce minimal architectural changes that achieve stable, single-stage training. To make TVM efficient in practice, we develop a fused attention kernel that supports backward passes on Jacobian-Vector Products, which scale well with transformer architectures. On ImageNet-256x256, TVM achieves 3.29 FID with a single function evaluation (NFE) and 1.99 FID with 4 NFEs. It similarly achieves 4.32 1-NFE FID and 2.94 4-NFE FID on ImageNet-512x512, representing state-of-the-art performance for one/few-step models from scratch.

• The paper presents Terminal Velocity Matching (TVM), enabling one/few-step high-fidelity image generation by regularizing trajectories at the terminal time.
• It achieves state-of-the-art FID scores (e.g., 3.29 FID for 1-NFE on ImageNet-256) by tightly bounding the 2-Wasserstein distance between generated and real data.
• Efficient training is realized through fused attention kernels and transformer normalization fixes, ensuring scalability and robust inference.

Terminal Velocity Matching: A Scalable Framework for One/Few-Step Generative Modeling

Introduction

The "Terminal Velocity Matching" (TVM) framework (2511.19797) generalizes Flow Matching (FM) for generative modeling, with a primary aim: achieving high-fidelity generation with extremely fast inference—namely, in one or few sampling steps—on high-dimensional datasets such as ImageNet. Unlike prior flow-based or diffusion-based generative models, which regularize the generative process at the initial time or require multi-stage/curriculum learning, TVM regularizes the learned trajectory at the terminal time. This shift yields explicit theoretical guarantees on the distributional closeness between data and generated samples, measured via the $2$-Wasserstein distance. The approach achieves state-of-the-art FID scores for single- and few-step generation on high-resolution ImageNet, including 3.29 FID for 1-NFE and 1.99 FID for 4-NFE on ImageNet-256, outperforming recent methods trained from scratch. The framework also presents a practical, scalable training and inference pipeline using a single neural network, efficient attention kernels for Jacobian-Vector Product (JVP) evaluation, and critically addresses architectural pathologies in common backbone classs (e.g., Diffusion Transformers).

Theoretical Framework

TVM builds upon the ODE interpretation of generative flows. For a neural displacement field $\mathbf{f}_\theta(x_t, t, s)$, it defines the mapping between two arbitrary times $t > s$, and regularizes the network such that the terminal derivative $\partial_s \mathbf{f}_\theta(x_t, t, s)$ matches the ground-truth velocity at the trajectory's endpoint. This terminal matching, rather than initial or midpoint matching, yields a pivotal property: the TVM objective tightly upper-bounds the $2$-Wasserstein distance between the model and ground-truth distributions under Lipschitz regularity of the parameterization. Specifically, TVM interpolates between one-step and multi-step sampling by parameterizing both the net displacement (direct jump) and instantaneous velocity (for continuous flows) in a single model. This duality enables deployment of generative models that are performant at both the few-step and ODE-solver regimes.

Figure 1: An illustration of Terminal Velocity Matching. Left: displacement map by integrating true velocity; right: TVM jump along ground-truth in one step. The one-step generation $x_0$ from $x_t$ is correct if the model’s terminal velocity matches true velocity at $s=0$.

The TVM loss couples two core terms: (1) a matching of the terminal velocity's gradient (JVP) to the velocity at the displaced point, and (2) a boundary matching term equivalent to FM at $t = s$. Theoretical analysis connects the expectation of this loss to Wasserstein distance, offering distribution-level guarantees absent in other single-stage, few-step generative schemes (such as MeanFlow, Consistency Models, or Shortcut).

Practical Implementation and Architectural Fixes

Direct TVM optimization requires higher-order derivatives: evaluating the time-derivative of model outputs with respect to the terminal parameter in the flow. To enable tractable training at ImageNet scale with large transformers, the framework contributes a fused Flash Attention kernel for efficient JVP computation with backprop, reducing both runtime and GPU memory (up to 65% improvement compared to naive PyTorch).

The authors identify that common backbone architectures, such as DiT [peebles2023scalable], lack Lipschitz continuity due to their use of LayerNorm and unbounded modulation, destabilizing TVM training. They propose simple, non-restrictive normalization modifications: substituting LayerNorm with (parameter-free) RMSNorm for QK normalization and modulation, and additional AdaLayerNorm normalization, resulting in smoothed activations at all layers and empirically stable training.

Additionally, the model is conditioned on classifier-free guidance (CFG) parameters during training, allowing for flexible control at sampling time and robust handling of different guidance scales by rescaling the network’s output by the CFG weight.

Empirical Results

Extensive experiments are conducted on ImageNet-256 and ImageNet-512, with a single-stage training protocol and no curriculum. The principal results include:

• On ImageNet-256, TVM achieves 3.29 FID for 1-NFE and 1.99 FID for 4-NFE, outperforming or matching previous methods such as MeanFlow and IMM both in the one-step and few-step regimes.
• On ImageNet-512, TVM achieves 4.32 FID (1-NFE) and 2.94 FID (4-NFE), demonstrating consistent scalability and outperforming prior works, including sCT and MeanFlow at the same model size.
• The approach outperforms distilled diffusion baselines (e.g., DiT with 500-NFE Steps) in only 4 steps, highlighting sample fidelity and informativeness is not compromised by aggressive acceleration.

  Figure 2: One-step samples from TVM on both ImageNet-512 and ImageNet-256, illustrating visual quality at high resolution.

Qualitatively, samples from both the 1-NFE and 4-NFE regime on ImageNet exhibit high semantic quality and diversity, without mode collapse or oversmoothing. Additional ablation reveals that TVM is robust to random CFG sampling, does not require training curricula, and performs optimally with a fixed guidance weight. The entire method remains scalable: it does not require multi-particle loss (as in IMM), nor any form of consistency distillation.

Comparative Analysis with Prior Methods

TVM provides a critical advancement over MeanFlow and Consistency Models. In MeanFlow, random CFG leads to highly unstable gradient norms due to the dependence of the JVP on the initial velocity, whereas TVM’s terminal perspective produces much smoother, stable gradients, directly observable in network training and FID evolution.

Figure 3: Left, MeanFlow demonstrates wide fluctuation in gradient norm under random CFG, while TVM’s gradient norm remains stable, reflecting improved training dynamics.

TVM’s theoretical grounding in Wasserstein distance is also unique. Unlike prior work (e.g., PID or shortcut models), which only permit one fixed time end-point, TVM supports arbitrary start/stop intervals, enabling seamless interpolation from one-step to multi-step generation at inference, all with a single learned model.

Implications and Future Prospects

TVM enables practical, scalable, and distributionally-principled one/few-step generative modeling. Its formulation provides a path to robust generation regimes (including high guidance, high-resolution, and limited compute), and its architectural interventions suggest future generic recipes for stabilizing and controlling flow-based deep generative models. The simple, single-sample training loss accommodates batch size one, enabling scaling to large models and dense data without the explosion in memory or compute suffered by multi-sample or multi-trajectory-based methods (as in IMM).

Given the robustness of TVM’s theory—the upper bound on $2$-Wasserstein between model and data—future lines of inquiry might generalize such terminal-matching schemes to more expressive path parameterizations, extend to modalities beyond images (e.g., video, audio), and integrate richer domain priors or advanced guidance schemes.

Conclusion

Terminal Velocity Matching offers a theoretically principled and practically efficient avenue for training one/few-step generative models, overcoming critical limitations in prior art. The framework’s guarantees on Wasserstein distance, its architectural fixes for transformer Lipschitz continuity, and its scalable, curriculum-free training pipeline collectively advance practical generative modeling for high-dimensional data. TVM’s empirical results validate its claims, delivering top-tier sample quality in minimal sampling steps and opening new directions for scalable, distributionally robust generative model design.