Terminal Velocity Matching

AI summary

Overview

Research area: Generative modeling — specifically one-step and few-step image generation with diffusion/flow-matching methods.

Technical level: Advanced. The paper combines a Wasserstein-distance bound, Lipschitz-continuity analysis of transformer architectures, and custom CUDA/Flash-Attention kernel work alongside large-scale ImageNet benchmarks.

Scope: The paper introduces Terminal Velocity Matching (TVM), a single-stage training objective that regularizes flow trajectories at their terminal time rather than their initial time, yielding state-of-the-art one- and few-step image generation on ImageNet at 256×256 and 512×512.

What This Paper Is About

Diffusion models and flow matching produce high-quality images but need many sampling steps (the paper cites 50 as typical), which is expensive. A wave of recent work tries to train a model in a single stage that can generate good samples in one or a few function evaluations, but existing approaches such as consistency models, MeanFlow, and sCT either lack a formal connection to distribution matching or require multiple particles per training step (as with IMM).

TVM's central idea is a conceptual shift: instead of matching the time derivative of a trajectory at the initial time (as flow matching does), TVM matches it at the terminal time. This lets one network parameterize both the instantaneous velocity field and the large-jump displacement map, and it comes with a proof that the training objective upper bounds the 2-Wasserstein distance between the data distribution and the model distribution — provided the model's velocity field is Lipschitz continuous.

Key Contributions

1. A new single-stage training objective (TVM). The method models the transition between any two diffusion timesteps and regularizes behavior at the terminal time. It reduces exactly to standard Flow Matching when the start and end times coincide (t = s), so flow matching is recovered as a special boundary case.

2. A distribution-level theoretical guarantee. Theorem 1 shows that a weighted integral of the TVM objective upper bounds the squared 2-Wasserstein distance between the model's pushforward distribution and the data distribution, assuming the model's velocity field is Lipschitz continuous. Unlike IMM — which also gives distribution-level guarantees — TVM does not require multiple particles per training step.

3. Minimal architectural fixes for Lipschitz control. Because Diffusion Transformers (DiT) are not Lipschitz continuous, the authors add RMSNorm-based QK-normalization, replace LayerNorm with parameter-free RMSNorm, apply RMSNorm to the AdaLN modulation parameters, and use Lipschitz initialization for linear layers (except time-embedding layers). They call this partial Lipschitz control and find it sufficient in practice.

4. An efficient Flash Attention kernel supporting Jacobian-Vector Product (JVP) backward passes. The terminal-velocity term requires differentiating through the network via JVP. The authors build a kernel that fuses JVP with the forward pass, cuts memory relative to naive PyTorch attention, and — uniquely relative to sCT, MeanFlow, and related work — supports gradient propagation through the JVP term. They report up to 65% speedup.

5. Practical training recipes for CFG. The network is conditioned on the classifier-free guidance weight w, and the objective is reweighted by 1/w² to prevent exploding gradients. The authors note this enables stable training across guidance scales without curriculum learning or specialized loss modifications, and they introduce a scaled parameterization where the network output naturally scales with w.

Main Findings

• ImageNet 256×256, one-step: TVM achieves 3.29 FID with a single function evaluation (NFE), compared with 3.43 FID for MeanFlow at 1 NFE.

• ImageNet 256×256, four-step: TVM reaches 1.99 FID with 4 NFEs, versus 2.27 FID for the DiT-XL/2 baseline at w = 1.5 with 250×2 NFE.

• ImageNet 512×512: TVM achieves 4.32 FID at 1 NFE and 2.94 FID at 4 NFE, versus 3.04 FID for DiT at 4 NFE.

• Capacity efficiency: TVM-XL/2 is reported to outperform sCT-XL, which has 1.1B parameters, indicating more optimal use of model capacity for fitting the image distribution.

• Guidance trade-off: For both ImageNet resolutions, TVM trained with higher CFG weight performs better at 1 NFE but worse at 2 NFEs. The authors interpret this as a fundamental trade-off between quality at different NFE counts, possibly limited by network capacity to fit all NFE regimes well.

• Instability of vanilla DiT: The original DiT exhibits training instability, visible as a steep jump in the activation norm of the last time-embedding layer. The proposed RMSNorm-based changes keep activations smooth.

• Optimizer sensitivity: Because of higher-order gradients through the JVP, the loss fluctuates under the default AdamW β₂ = 0.999. Setting β₂ = 0.95 (following practice from language models) produces noticeably smoother terminal velocity error.

• Baselines in the comparison table (ImageNet 256×256): ADM at 10.96 FID (554M params, 250×2 NFE), LDM-4-G at 3.60 FID (400M), DiT-XL/2 at 3.22 FID with w = 1.25 (675M), DiT-XL/2 at 2.27 FID with w = 1.5 (675M), SiT-XL/2 at 2.15 FID with w = 1.5 (675M), iCT-XL/2 at 34.24 and 20.3 FID for 1 and 2 NFE (675M), Shortcut-XL/2 at 10.60 FID for 1 NFE (675M), IMM-XL/2 at 8.05 FID (1×2), 3.99 FID (2×2), and 2.51 FID (2×4) (675M), and MeanFlow-XL/2 at 3.43 FID (1 NFE) and 2.93 FID (2 NFE) with 676M parameters. The TVM-XL/2 rows of the table were truncated in the provided content, so their per-row figures are not reported here beyond the abstract and body text.

• Not reported: The provided content does not report dataset sizes, total training compute, wall-clock training time, or total parameter counts for the TVM models beyond the "XL/2" designation and the DiT-XL/2 architecture they build on. Text-to-image results at 10B+ scale are referenced only via an external blog link and are not quantified in the paper text.

Methodology in Plain English

Flow matching works by drawing a straight line between a data point and noise, training a network to predict the direction of that line (the velocity) at every point along the way, then integrating that predicted field with an ODE solver at sampling time. TVM keeps the same setup but asks a different question: rather than only asking "what is the direction right here, right now," it also asks a network to predict the entire jump from a starting time to an ending time in one shot.

The key mathematical trick is that the derivative of the jump function, evaluated at the terminal time, must equal the velocity field at that terminal point. So a single network can be asked to output both things at once: a displacement prediction and an instantaneous velocity. Training combines two losses — a displacement loss that enforces the terminal velocity condition, and a standard flow-matching loss that anchors the velocity field to real data. Crucially, the ground-truth quantities needed for the terminal condition are replaced with the network's own predictions, using exponential moving averages and stop-gradients to keep training stable. This removes the need to simulate ODEs during training, which would otherwise be infeasible.

Two engineering problems arise. First, the theory assumes the velocity field is Lipschitz continuous, but standard Diffusion Transformers are not — so the authors patch the architecture with RMSNorm in several strategic places to keep activations bounded. Second, computing the terminal-velocity term requires a Jacobian-Vector Product through the attention layers, which PyTorch and open-source Flash Attention handle poorly — so they write a fused kernel that supports the backward pass through it. Sampling is then straightforward: start at time t = 1 with noise, call the network once with (t, s) and step, and optionally repeat for more steps. The provided PyTorch-style sampling snippet shows this loop.

Why This Matters

Impact on research. TVM is one of the few from-scratch few-step methods with an explicit distribution-level guarantee (a bound on the 2-Wasserstein distance), rather than just a trajectory-matching loss. It also highlights that transformer architecture — specifically the absence of Lipschitz continuity — is a first-order obstacle to stable training of this class of objectives, which is a finding that generalizes beyond TVM itself. The demonstration that gradients can be propagated through the JVP in attention opens a path that prior work (sCT, MeanFlow) did not take.

Real-world applications:

• Real-time image generation and editing interfaces, where users expect near-instant response and cannot wait for 50–500 network evaluations.
• Video generation, which the paper explicitly identifies as a motivating case because multi-step sampling is especially expensive for high-dimensional data.
• On-device or edge inference, where compute budgets make multi-NFE diffusion impractical.
• Interactive creative tooling and prototyping where rapid iteration on generated outputs matters more than peak sample fidelity.

Industry relevance. The authors are affiliated with Luma AI, and the paper points to text-to-image results at 10B+ scale, indicating the method is intended for production-scale generative systems rather than only academic benchmarks. The Flash Attention kernel work is directly relevant to anyone training these models at scale, since it addresses a concrete memory and speed bottleneck.

Future Directions

• The NFE trade-off. The observed pattern that higher CFG weights help 1-NFE but hurt 2-NFE quality suggests a capacity limitation. The authors explicitly leave deeper study of this trade-off and design improvements to future work.
• Scaling beyond 512×512 and toward video. The paper motivates video as a key target but reports only ImageNet image results; scaling behavior on high-dimensional video data is not addressed in the provided content.
• Extending JVP kernel support. The custom Flash Attention JVP kernel is presented as filling a gap left by PyTorch and open-source Flash Attention. Broader integration into standard tooling is a natural next step.
• Full versus partial Lipschitz control. The authors deliberately apply only partial Lipschitz control and find it sufficient empirically. Whether stronger or more principled control yields further stability gains is left open.

Target Audience

This paper is best suited for machine learning researchers and engineers already familiar with diffusion models, flow matching, and ODE-based generative modeling — particularly those working on few-step or one-step generation, training stability of transformers, or high-performance attention kernels. Readers need comfort with Wasserstein distances and Lipschitz continuity to engage with the theory, though the core intuition (match the derivative at the end of the path, not the start) is accessible with background in flow matching. Practitioners focused on production-scale image or video generation will find the architectural and optimizer recipes directly useful.

Authors’ 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.