PaperlayerView exact revision on arXiv

Source-linked AI summary

Terminal Velocity Matching

Linqi Zhou, Mathias Parger, Ayaan Haque, Jiaming Song

arXiv:2511.19797v3cs.LGcs.AIcs.CVstat.ML

TL;DR

High-quality image generation often requires many sampling steps, which makes generation computationally expensive, especially for high-dimensional data. TVM matches terminal trajectory velocities and provides a distributional guarantee under a Lipschitz assumption, with architectural changes for diffusion transformers. It reports strong one- and few-step ImageNet results from scratch.

  • Problem

    Diffusion and Flow Matching methods typically require many sampling steps for high-quality outputs, making generation costly, particularly for high-dimensional data.

    Source: §1 · p. 1

  • Method

    TVM matches trajectory velocities at terminal times and trains jointly with the zero-displacement Flow Matching boundary case.

    Source: §1 · p. 2, §3 · p. 4

  • Results

    TVM achieves state-of-the-art one-step results from scratch and surpasses baseline diffusion models with 4 NFEs.

    Source: §8 · p. 10

  • Takeaways & Limitations

    TVM’s objective upper-bounds the 2-Wasserstein distance under Lipschitz continuity, and minimal architectural changes address the lack of that property in diffusion transformers.

    Source: §1 · p. 2, §4 · p. 5

    Higher CFG improves 1-NFE performance but worsens 2-NFE performance, a trade-off the authors leave for future study.

    Source: §7.1 · p. 8

Abstract

from arXiv · show

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.

1 INTRODUCTION

High-quality generation often requires many sampling steps, making it costly for high-dimensional data. TVM matches trajectory velocities at their terminal time and provides a distribution-level bound, with architectural changes addressing a key training constraint.

1 INTRODUCTION

  • Prior single-stage trajectory methods lack explicit distribution-matching connections, while IMM provides a distribution-level guarantee but requires multiple particles per training step.The cited passage names consistency-based approaches and MeanFlow among methods learning integrated trajectories.Source: §1 · pp. 1–2
  • TVM matches trajectory derivatives at the terminal rather than initial time and its objective upper-bounds the 2-Wasserstein distance between data and model distributions.The guarantee is stated under Lipschitz continuity; the paper identifies that current diffusion transformers lack this property and proposes minimal architectural modifications.Source: §1 · p. 2
  • A fused attention kernel supports backward passes on Jacobian-Vector Products and achieves up to 65% speedup over standard PyTorch operations.The kernel supports the terminal-velocity computation; the passage also describes CFG scaling and random CFG training without curriculum learning or specialized loss modifications.Source: §1 · p. 2
  • 3.29 FID with 1 NFE on ImageNet-256×256 outperforms MeanFlow’s 3.43 FID, while 4-NFE TVM reaches 1.99 FID versus DiT’s 2.27 FID.On ImageNet-512×512, TVM also reaches 2.94 FID versus DiT’s 3.04 FID with 4 NFEs.Source: §1 · p. 2

2 PRELIMINARIES: FLOW MATCHING

Flow Matching defines a linear path between data and a prior, then learns its marginal velocity field for ODE-based transport. The corresponding flow map describes transitions between times along that field.

2 PRELIMINARIES: FLOW MATCHING

  • Flow Matching linearly interpolates data x_0 and prior x_1 as x_t=(1−t)x_0+tx_1, with conditional velocity v_t=x_1−x_0.A ground-truth marginal velocity field exists but is not known analytically, so a neural network approximates it.Source: §2 · p. 2
  • The learned velocity field transports samples from the prior toward the data distribution by solving an ODE.The paper notes that classical ODE integration methods such as Euler are used to produce samples.Source: §2 · p. 2
  • A flow map gives the displacement from a start time to an end time by integrating the ground-truth velocity along the intervening path.The cited definition describes transitions between times along the velocity field.Source: §2 · p. 2
  • Figure 2 contrasts integrating the true velocity to obtain a displacement map with a model path that jumps directly between points on the ground-truth path.At zero displacement, its terminal-velocity condition reduces to Flow Matching.Source: Figure 2 · p. 3

3 TERMINAL VELOCITY MATCHING

TVM learns transitions between timestep pairs by learning the ODE integral, supporting one- or few-step generation with inference-time scaling. Its terminal-velocity conditions relate learned displacement to instantaneous velocity and support a distribution-level guarantee.

3 TERMINAL VELOCITY MATCHING

  • A single two-time-conditioned network learns both the displacement map and instantaneous velocity, with scaling chosen to satisfy the integral boundary condition when t = s.TVM directly learns the ODE integral for transitions between timesteps.Source: §3 · p. 3
  • Differentiating the displacement integral with respect to its terminal time bypasses explicit ODE integration; zero terminal-velocity error implies exact displacement, and the error reduces to the Flow Matching loss as t approaches s.The terminal-velocity error upper bounds displacement error.Source: §3 · pp. 3–4
  • TVM jointly minimizes general terminal-velocity error and its zero-displacement boundary case, using the learned instantaneous velocity as a proxy to guide displacement learning.The boundary case is the Flow Matching loss.Source: §3 · p. 4
  • Under Lipschitz continuity of the learned velocity field and mild regularity conditions, a weighted integral of the TVM objective upper-bounds the 2-Wasserstein distance between data and model distributions.The weighting depends on the Lipschitz constants.Source: §3 · p. 4
  • For practical training, TVM randomly samples both timesteps and uses an EMA-based, stop-gradient estimate of its per-time objective.When t = s, the objective reduces to Flow Matching.Source: §3 · p. 4
  • For class-conditional generation, the network conditions on class and CFG weight, and its instantaneous velocity is regressed against a conditional–unconditional velocity combination.The objective’s minimizer coincides with the ground-truth CFG velocity.Source: §3 · p. 5

4 PRACTICAL CHALLENGES

TVM’s practical implementation addresses Lipschitz control, efficient JVP computation, and training stability. The paper also reports using a separate time distribution for the FM loss.

Semi-Lipschitz control

  • RMSNorm-based changes and Lipschitz initialization keep activations smooth, while partially controlling Lipschitzness is reported as sufficient for empirical success.The changes include RMSNorm as QK-Norm, replacing LayerNorm, and normalizing AdaLN modulation parameters.Source: §4 · p. 5, Figure 4 · p. 5

Flash Attention JVP with backward pass

  • A fused Flash Attention kernel computes JVPs and supports backward propagation through JVP results, using less memory than naïve PyTorch attention.The training objective differentiates the map with respect to s, including through the JVP term.Source: §4 · p. 6

Optimizer parameter change

  • β2 = 0.95 reduces terminal velocity error fluctuations compared with the default AdamW β2 = 0.999.The change addresses loss fluctuations associated with higher-order gradients through JVP.Source: §4 · p. 6, Figure 5 · p. 6

Different time distribution for FM loss

  • A separate distribution for sampling s in the FM loss is empirically helpful and transfers a successful FM training time distribution to TVM.The method uses a separate s′ sample for the FM loss term.Source: §4 · p. 6

5 CONNECTION TO PRIOR WORKS

TVM connects to prior methods through its treatment of trajectory derivatives and timestep transitions. The section also notes an optimization instability in a naïve combination of PID and FM loss.

MeanFlow

  • MeanFlow differentiates with respect to t, whereas TVM enforces a differential condition with respect to s; the distribution-divergence relationship for MeanFlow remains elusive when using v_t.The paper states that using v_t in MeanFlow introduces training variance and can cause gradient fluctuations, especially with random CFG.Source: §5 · p. 6

Physics Informed Distillation

  • PID is the special case t = 1 in TVM’s two-time formulation and therefore supports only one-step samples.TVM introduces a starting time t in addition to PID’s terminal time s.Source: §5 · p. 6

PID and FM loss

  • A naïve combination of PID and FM loss suffers from optimization instability, while a continuous distribution on t is necessary for empirical success..Source: §5 · p. 7

6 RELATED WORKS

Diffusion and flow-matching methods use ODE-based sampling, while one- and few-step methods seek faster generation. TVM is distinguished by terminal-time regularization and explicit 2-Wasserstein bounds.

Diffusion and Flow Matching

  • Diffusion and flow matching ultimately solve ODEs with neural networks during sampling.Flow matching generalizes diffusion models to arbitrary priors with simplified training.Source: §6 · p. 7

One-Step and Few-Step Models from Scratch

  • Consistency models cannot transition between arbitrary timesteps without additional noise, while IMM requires multiple particles and MeanFlow and Flow Map Matching lack distributional guarantees.These are among the one-step and few-step approaches discussed for addressing slow ODE simulation.Source: §6 · p. 7

TVM in relation to prior work

  • TVM regularizes path behavior at the terminal rather than initial time and provides explicit 2-Wasserstein bounds; unlike sCT and MeanFlow, it supports backward passes through JVP computation.The passage also notes that Boffi et al. (2025) offers a unifying perspective on training flow maps from scratch with similar theoretical insights.Source: §6 · p. 7

7 EXPERIMENTS

Experiments evaluate TVM’s image generation, training choices, ablations, and computation cost. TVM outperforms prior from-scratch methods in one-step ImageNet generation and achieves strong few-step results at both resolutions.

7.1 IMAGE GENERATION

  • On ImageNet-256×256, with CFG w = 2, TVM achieves 3.29 FID versus MeanFlow’s 3.43 at 1 NFE; with 4 NFEs, w = 1.75 exceeds 500-NFE diffusion baselines.On ImageNet-512×512, TVM outperforms sCT and MeanFlow at 1 and 2 NFEs, and with w = 2.25, matches the 500-NFE DiT-XL/2 baseline using 4 NFEs.Source: §7.1 · p. 8

7.2 DISCUSSION ON TRAINING ADVANTAGES

  • TVM uses a single sample for loss calculation, supports random CFG without collapse, and uses neither a training curriculum nor modified loss weighting.The authors attribute random-CFG stability to the JVP being calculated only with respect to s; they say the simpler design may improve scalability.Source: §7.2 · pp. 8–9

7.3 ABLATION STUDIES

  • Biasing t toward 1 and s toward 0 favors longer strides, but excessive bias worsens results; gap sampling outperforms trunc and clamp over longer training.The separately sampled gap* scheme generally performs better in 1-NFE FID.Source: §7.3 · pp. 9–10
  • Constant CFG outperforms random CFG, and constant w = 2 converges faster than the Flow Matching default w = 1.5.Random CFG can converge, but some guidance scales degrade in FID.Source: §7.2 · p. 9, §7.3 · p. 10
  • An EMA target rate around γ = 0.99 is the reported sweet spot; scaled parameterization is generally beneficial, while adding t = s training trades lower 1-NFE performance for a marginal 2-NFE gain.The authors report that γ = 0 lags alternatives with γ > 0, while γ = 0.9999 can cause instability.Source: §7.3 · p. 10

7.4 MEMORY AND RUNTIME ANALYSIS

  • TVM with Lipschitz control has higher runtime and memory use, mostly from architectural changes; detaching the JVP can reduce runtime while biasing learning gradients.The measurements average 10 training steps, exclude EMA update cost, and use no performance optimization such as torch.compile.Source: §7.4 · p. 10

8 CONCLUSION

TVM trains one- and few-step generative models from scratch by matching a flow trajectory’s terminal velocity rather than its initial velocity. The conclusion reports a 2-Wasserstein upper bound and discusses Lipschitz continuity, while the network parameterization imposes a boundary condition on γ(t, s).

8 CONCLUSION

  • TVM matches a flow trajectory’s terminal velocity instead of its initial velocity and provides an objective that upper-bounds the 2-Wasserstein distance up to a constant.The conclusion also identifies Diffusion Transformers’ lack of Lipschitz continuity as a concern for the theory.Source: §8 · p. 10
  • TVM is presented as a framework for training one- and few-step generative models from scratch.The conclusion describes its one-step result as state of the art and says it surpasses baseline diffusion models with four NFEs.Source: §8 · p. 10

A.1 GENERAL NETWORK PARAMETERIZATION

  • The general network parameterization uses γ(t, s) with γ(t, t) = 0 as a boundary condition.The passage also introduces a velocity condition but does not specify it in the supplied text.Source: §A.1 · p. 14

A.2 TERMINAL VELOCITY ERROR UPPER BOUNDS DISPLACEMENT ERROR

The appendix relates terminal velocity error to displacement error and connects the terminal-velocity objective to flow matching. Its main theorem links the model to 2-Wasserstein distance under Lipschitz and regularity assumptions.

A.2 TERMINAL VELOCITY ERROR UPPER BOUNDS DISPLACEMENT ERROR

  • Under mild regularity assumptions, Lemma 1 gives an inequality bounding displacement error by terminal velocity error.The proof assumes the displacement maps are Riemann-integrable and uses the triangle inequality and a regularity assumption.Source: §A.2 · p. 14

A.3 TERMINAL VELOCITY ERROR REDUCES TO FM

  • As the terminal time approaches the initial time, the terminal velocity error reduces to the original flow-matching loss, equivalent up to a constant to conditional Flow Matching loss.The passage states that the expected L2-norm limit exists assuming the norm is bounded.Source: §A.3 · p. 15

A.4 MAIN THEOREM

  • The main theorem connects the method to a 2-Wasserstein bound assuming Lipschitz-continuous uθ and additional mild regularity conditions.The bound uses a functional of the Lipschitz constants and a non-optimizable constant.Source: §A.4 · p. 15

A.5 REDUCTION TO FLOW MATCHING

  • When t = s, the TVM loss reduces to Flow Matching loss.Source: §A.5 · p. 17

A.6 DERIVATION FOR CLASS-CONDITIONAL TRAINING TARGET

  • The minimizer of the class-conditional training objective is the CFG velocity wu(xs, s, c) + (1 − w)u(xs, s).Source: §A.6 · p. 17

B.1 LIPSCHITZNESS OF RMSNORM

The appendix bounds the Jacobian norm of RMSNorm and describes normalizing DiT modulation parameters. It also details a Flash-style JVP attention backward pass and its multi-step implementation.

B.1 LIPSCHITZNESS OF RMSNORM

  • RMSNorm’s Jacobian has bounded Frobenius norm, and therefore bounded matrix norm, using the bound on its entries and RMS(x) ≥ ϵ.Source: §B.1 · p. 18

B.2 FULL DESCRIPTION OF NORMALIZATION OF MODULATION

  • The DiT layer’s six modulation parameters are each passed through RMSNorm to form the normalized layer’s inputs.The parameters are a1(t), b1(t), c1(t), a2(t), b2(t), and c2(t).Source: §B.2 · p. 18

C FLASH ATTENTION JVP WITH BACKWARD PASS

  • The Triton implementation computes JVP SDPA forward and backward kernels, using online softmax to avoid storing large intermediate attention matrices.The JVP backward computation is split into smaller kernels to reduce register spills.Source: C FLASH ATTENTION JVP WITH BACKWARD PASS · pp. 18–19

C.1 MULTI-STEP BACKWARD PASS

  • The backward pass is split into smaller operations with row-parallel dQ and dQ̇ gradients and column-parallel dK, dK̇, dV, and dV̇ gradients.The authors report that redundant but coalesced computation outperformed a single fused kernel using atomic operations.Source: §C.1 · p. 19

C.2 EVALUATION

The JVP attention-kernel evaluation reports reduced peak memory and up to 65% speedup against vanilla SDPA on an H100. The appendix also presents the training algorithm and brief MeanFlow and flow-map self-distillation material.

C.2 EVALUATION

  • The JVP attention implementation reduces peak memory and achieves up to 65% speedup over vanilla SDPA in the backward-pass evaluation.The comparison uses an NVIDIA H100 SXM 80GB and evaluates a single SDPA operation’s backward pass, combining primal and tangent gradients.Source: §C.2 · p. 20, Table 5 · p. 21

D TRAINING ALGORITHM

  • The appendix presents the training algorithm and Figure 11’s PyTorch-style pseudocode for calculating fθ(xt, t, s) and its derivative with one JVP pass.Source: D TRAINING ALGORITHM · p. 21

E.1 MEANFLOW

  • The supplied MeanFlow passage identifies the associated objective as the MeanFlow loss.Source: §E.1 · p. 22

E.2 FLOW MAP SELF-DISTILLATION

F ADDITIONAL EXPERIMENT DETAILS

The appendix records implementation choices and additional CFG-conditioned experiments. It reports a late-training FID tradeoff across guidance scales in w-conditioned flow matching.

F ADDITIONAL EXPERIMENT DETAILS

  • Table 6 collects overall training details, while its footnotes clarify the t=s setting percentage and distinguish evaluation EMA from training-target EMA.Source: F ADDITIONAL EXPERIMENT DETAILS · p. 22

F.1 ARCHITECTURE AND OPTIMIZATION

  • The setup uses the standard Stable Diffusion VAE with a different scale and shift, conditions on t − s and β = 1/w, and uses DiT’s default optimizer with BF16 precision.It sets β₂ = 0.95 and uses no learning-rate scheduler. During random CFG training, β is sampled uniformly from 1/wmax to 1/wmin, with w = 1/β; wmax can be set to a realistic finite value.Source: §F.1 · p. 23
  • The XL model uses 18 rather than 16 attention heads for kernel efficiency, with similar FID during training and a 1.25× speed advantage.With 1152 hidden units, 18 heads give 64-dimensional heads that fit the kernel's CUDA-block allocation; 16 heads give 72-dimensional heads.Source: §F.1 · p. 23

F.2 DETAILS ON RANDOM CFG WITH MEANFLOW

  • MeanFlow's random-CFG mixing construction subsumes regular CFG and MeanFlow CFG, but its observed training fluctuation may depend on settings and environment.The experiments randomly sample κ and w with finite boundaries and omit adjustment tricks such as AdamW changes or gradient clipping.Source: §F.2 · p. 23

F.3 CFG-CONDITIONED FLOW MATCHING

  • In w-conditioned flow matching, FID for w=1.5 decreases while FID for w=2 increases at later training steps.The passage says a similar tradeoff is observed in TVM.Source: Figure 12 · p. 23, §F.3 · p. 24

F.4 ADDITIONAL VISUAL SAMPLES

  • Figure 13 presents additional ImageNet-256×256 samples from the 1-NFE TVM model.Source: Figure 13 · p. 25

Related papers

Loading 2511.19797v3…