Terminal Velocity Matching
Authors: Linqi Zhou
Mathias Parger
Ayaan Haque
Jiaming Song
Venue: ICLR 2026 Poster (2026)
Keywords: one-step generative model from scratch
diffusion
flow matching
Primary area: generative models

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.

Published as a conference paper at ICLR 2026
TERMINAL VELOCITY MATCHING
Linqi Zhou Luma Al
Mathias Parger Luma Al
Ayaan Haque Luma Al
Jiaming Song Luma Al
velocity
model
X. + (,
(( 6)
TVM (Ours)
de
dt
(X.. (, =0)==(((a ((
MeanFlow
t,a)
Flow Matching
((x..t.) !) de
Figure 1: (Left) conceptual comparison of method to prior methods. TVM guides the one-step model via terminal velocity rather than initial velocity. (Right) 1-NFE samples on ImageNet at 256 and 512 resolution.
ABSTRACT
We propose Terminal Velocity Matching (TVM), generalization of flow match- ing that enables high-fidelity one- and few-step generative modeling. TVM mod- els the transition between any two diffusion timesteps and regularizes its behavior at 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 introduce minimal architectural changes that achieve stable,
lack this property, single-stage training. To make TVM efficient in practice, we develop a fused at- tention kernel that supports backward passes on Jacobian- Vector Products, which scale well with transformer architectures. On 2mageNet-256xx255 TVM achieves 3.29 FID with a single function evaluation (NFE) and 1.99 FID with NFEs. It similarly achieves 4.32 1-NFE FID and 2.94 4-NFE FID on ImageNet-512x65 representing state-of- oftte-of-thearr performance for one/few-stey step models from scratch."
INTRODUCTION
Can we build generative models that simultaneously deliver high-quality samples, fast inference, and scalability to high--dmensiooal data, all from a single training stage? This the central challenge that continues to drive research in generative models. While Diffusion Models booll-Dcckstein et al., 2015; Ho et al., 2020; Song et al., 2020) and Flow Matching (Liu et al., 2022; Lipman et al., 2022) have become the dominant paradigms for generating images (Rombach al., 2022; Podell al., 2023; Esser et al., 2024) and videos (OpenAI, 2024: Wan et al., 2025), they typically require many sampling steps (e.g., 50) produce high-quality outputs. This multi-step nature makes generation computationally expensive, especially for high-dimensional data like videos.
In pursuing a single-stage training for few-step inference, recent methods have focused on directly learning integrated trajectories rather than relying on ODE solvers. Consistency based approaches (CT (Song et al., 2023), CTM (Kim et al., 2023), sCT (Lu & Song. 2024)) and trajectory matching methods like MeanFlow (Geng et al., 2025) learn to predict match trajectory derivatives. How- ever, these methods lack explicit connections to distribution matching, a fundamental measure of generative model quality. While Inductive Moment Matching (IMM) (Zhou et al., 2025) addresses
For Text-to-Image results 10B+ scale, visit https::///// //t::///mmm abs a /bl og engi neeri ng/tvm

---

Published as a conference paper at ICLR 2026
this gap by providing distribution-level guarantees through Maximum Mean Discrepancy, it requires multiple particles per training step, limiting scalability.
We propose Terminal Velocity Matching (TVM), new framework for learning ground-truth tra- jectories of flow-based models in a single training stage. Instead of matching time derivatives at the initial time, TVM matches them at the terminal time of trajectories. This conceptually simple shift yields powerful theoretical guarantees. We prove that our training objective upper bounds the 2- 2-Wasserstein distance between data and model distributions. Unlike IMM, our method provides distribution-level guarantees without requiring multiple particles. Our analysis also reveals a critical architectural limitation: current diffusion transformers (Peebles & Xie, 2023) lack Lipschitz conti- nuity, which destabilizes TVM training. We address this with minimal architectural modifications, including RMSNorm-basee nK-nommalinttion and time embedding normalization.
To make TVM practical at scale, we develop an efficient Flash Attention kernel that supports back- ward passes on Jacobian- Vector Products (JVP), crucial for our terminal velocity computation. Our implementation achieves up to 65% speedup and significant memory reduction compared to stan- dard PyTorch operations. We introduce a scaled parameterization where the network output natu- rally scales with the CFG weight w, allowing the model to handle varying guidance strengths more effectively. During training, we randomly sample CFG weights and directly incorporate them into objective function with appropriate weighting (1/w²) to prevent gradient explosion. This ap- proach enables stable training across diverse guidance scales without requiring curriculum learning or specialized loss modifications, making TVM straightforward implement and scale.
TVM achieves state-of-the-arr results on mmageNet-256x256 with 3.29 FID in single-step genera- tion (outperforming MeanFlow's 's (Geng et al., 2025) with 3.43 FID) and matches/exceeds diffusion baselines with just 4 function evaluation steps (i.e., 1.99 FID for TVM vs. 2.27 FID for DiT). Sim- ilarly, our method surpasses diffusion baselines with 4-NFE on ImageNet-512x55) (i.e. 2.94 FID for TVM vs. 3.04 FID for DiT) while outperforming prior from-scrraee methods such as sCT (Lu & Song, 2024) and MeanFlow on single-step generation. Our method naturally interpolates between one-step and multi-step sampling without retraining, requires no training curriculum or loss modi- fications, and maintains stability with simple architectures. Our construction provides new insights into building scalable one/few-step generative models with distributional guarantees, demonstrating that principled theoretical design can lead to practical improvements in both training stability and generation quality.
PRELIMINARIES: FLOW MATCHING
For a given data distribution (x())) and prior distribution p1 (x1). Flow Matching (FM) (Lipman et al., 2022; Liu et al., 2022) constructs a time-augmented linear interpolation x; between data Xg E _R_ and prior X1 E F such that x, (1 t)x, + tx, For each path x; conditioned on a (Xo, x1) pair, there exists a conditional velocity V, = x1 - Xo for each x,. Under this definition, ground-truth velocity field RD x [0, 1] - R D marginal over all data and prior exists but is not known in analytical form. Therefore, neural network us ux(xe,)) () is used as approximation via loss
(0) = Ex. wlw.v...llluooee (xe, 1) - ve
(I)
for all / E [0, 1] and x, ~ (x,) where pe (x.) denotes the marginal distribution over all data and prior. It can be shown that the minimizer Omin implies uem.. (xx,t) = u(x,, E) which can be used during inference to transport prior to data distribution by solving an ODE d u(x.,t). f).
For each ground-truth u(x,, t). there exists corresponding displacement map s : R., 10. 1] x [0, 1] + Rº (i.e. flow map (Boffi et al., 2024)) from any start time E [o, 1] to an end time E [0, 1]. It is defined as the ODE integral following u(x,,r) for all r E [s, t), Le.
t) Xx 1." (u(x,)))d
Empirically, (x, is used with classical ODE integration techniques such the Fuler method to produce samples.
"See I. Lipman al. (2022); Albergo al. (2023) for general path constructions.
2

---

Published as a conference paper at ICLR 2026
d t.....s))
uix:
Fo(8,
X,
N, =xxx
u(xu, uidu uide
X
X, = Ne. + fo(xe,t,,))
X1
d
X1
ground-rrtth displcceeett
ground truth velodit
mee-stop modee displacement
onp-stop model terminal velocity
An illustration of Terminal Velocity Matching. Left shows the ground-truth displacement map by integrating the true velocity. Right shows our model path directly jumping between points the ground-truth path one step. In method. the one-step generation Xa from coincides with ground-trudh -tunt Xo if the terminal velocity model d [xx, s) coincides with ground velocity u(x., 8) for all s [O, 4] along the true flow path (see Eq. (7)). The terminal velocity condition is jointly satisfied with the boundary case when model displacement is 0, where matching ds f(x,, t, sam with uxxe, t) reduces to Flow Matching.
Figure 2:
dz
TERMINAL VELOCITY MATCHING
We propose Terminal Velocity Matching (TVM), a single-stage objective that directly learns the ODE integral in Eq. 2. By learning the transition between any two timesteps, TVM can generate high quality solutions in one step or few steps, while enjoying inference-time scaling.
Let f(x,, t, s) (x,, t, s) x, denote the displacement of the velocity field. We observe that it must satisfy the following two conditions:
I f(x,,t,) f." u(x,, r)dr
- f(x,, t,
s=t u(x,, t).
(3)
ds
The first condition is the definition of net displacement and the second condition is true by differenti- ating both sides of the first condition w.r.t. s evaluated s f. It explicitly relates the displacement map (with large time jump) to the marginal velocity field (with infinitesimal time jump), allowing us interpolate between one-step sampling and ODE-like infinite-step sampling.
One of our key insights is that we can use a single two-time conditioned neural network Fx L, learn both the one-step displacement sampler from t and the instantaneous velocity field. For simplicity, we let our model with learnable parameters 0 be
(x 1,
1) (xe, 8)
un (xx, () fx (,
F& (xt, (,
(4)
where the scaling (s - t) is chosen satisfy integral boundary condition when s² Condition ② can be easily enforced by FM loss (in Eq. (1)) and condition ! can be naively enforced via the displacement error
fo(x,,t,0) J. u(x,, r r)dr
Caispa (0) Ex,
(S)
Once the above error is minimized to zero, one can obtain one-step samples by calling x + fx(x,, I, O) for any x₁ (x(x;) at [0, 1). However, this objective is infeasible because it re- quires ODE integration for each starting point We address this challenge by proposing simple sufficient condition to the network that bypasses explicit training-time ODE simulation.
Terminal Velocity Condition. Explicit integration can be bypassed via differentiating w.r.t. integral boundaries. For the ground-truth net displacement f(x,, 1, 8) condition D, differentiating w.r.t. gives rise to the following condition on terminal velocity, i.e.
d f(x,, s) u( (1 (xx s).
(6)
ds This condition is true for any ground-truie -ruth net displacement f, and show in Appendix A.2 that given t E [0, 1] and our parameterized map fa (x;, s).
d (x,, u 4(
2
C'assp (0)
Ex.
ds.
ds
This similar CTM (Kim al., 2023). See Appendix A.1 for conditions general scaling factors.
3

---

Published as a conference paper at ICLR 2026
This result shows that the terminal velocity error the right hand side upper bounds the displace- ment error, and so zero terminal velocity error implies that displacement from to 0 matches exactly. Moreover, it is easy to see that the terminal velocity error reduces to the marginal FM loss f + (see Appendix A.3). FM can thus be understood as matching a trajectory's 's terminal velocity when the net displacement is O. An illustration of framework is shown in Figure 2. Despite the sim- plicity and generality, in practice, fulfilling this condition is still difficult due to the requirement of and Fortunately, this issue can be effectively addressed using learned network as proxies.
Learned networks as proxies. Specifically, we propose the following approximation
(8) as proxies for the ground-truths. To properly guide the terminal velocity, ux(X,; needs to first approximate the ground-truth u(x,;8) 8) for any x, and s. Therefore, the proxy terminal velocity error can be jointly optimized with Flow Matching, which, the terminal velocity eπror when displacement is O. We use the term "Terminal Velocity Matching" for this joint minimization of general and boundary-case velocity error, where the objective is
u(,( t, s), = (x + fo(X,, t, s)
noted above, is a special boundary case of
d fx s) us (x; fx L, 8) ds
Che (T) Ex.,..V
uu (x, 5) Vo
(9)
satisties
satisfies
for each time t E [0, and E [0, E]. Intuitively, this objective leverages a single network to parameterize both the instantaneous velocity field and the displacement map, the former of which is learned from data to guide the learning of the latter. To provide further theoretical justification, in the following theorem, we formally establish a weighted integral of our objective as a proper upper bound on the 2- 2-Wasserstein distance between the data distribution po (xo) and our model distribution to #p pushforward from pe (x:) via our parameterized flow map.
Theorem (Connection the 2- Wasserstein distance). Given 1). let f. fr be the distribution pushforward from (x())) via fo (Xx,t, 0), and assume ug(-,s) is for all & E [0, t] with Lipschitz constants I(s),, with additional mild regularity conditions,
W & [
(0)ds CC,
(10)
where W2(:s-) is 2- Wasserstein distance, A[-] is a functional of L(-), and C is a non-optimizable constant.
Training objective. The theorem relates our per-time objective to distribution divergence. How- ever, for practicality, we avoid computation of the above weighting function and instead choose to randomly sample both and via distribution p(s, t) such that
(Tv Rr, [ (0))
(11)
where notably (TVM(0) reduces Flow Matching objective when (see Appendix A.5), In practice, we employ biased estimate of the above objective by using exponentially averaged (EMA) weights and stop- gro-ggradient for our proxy networks (Li al., 2023). The biased per-time objective (trm( is
d ||
2
Ex,,
lo (xx, ds
uo Axr + for, (Xe, l, 8), 8) 1eya + [uo(xs))) 8) Vs
(12)
where &sg and 0% are the stop-grad weight and stop-grad EMA weight and 1sze O when t and 1 otherwise to ensure the constraint to reduce to FM loss when ==
Classifier-free guidance (CFG). In the case of class-conditional generation. The ground-truth locity field is replaced by a linear combination of class-conditional velocity u(x,, r, c) and uncon- ditional velocity u(x,,r) r)) (Ho & Salimans, 2022), such that the new displacement map is
(x s, = =x, + S. [wu(x,, r, c) (1 w))(x,, r)) dr,
(13)
4

---

Published as a conference paper at ICLR 2026
where is the CFG weight, is class and Ø denotes empty label. To train with CFG, we ad- ditionally condition the network on w and c, and our class-cod class-conditiona. condetiodai map is fg(X,, &, 8, w) = (s - F) (x,, 8, c, w) where the additional w scale is chosen due to linear scaling in magnitude for marginal velocity w.r.t. w. The instantaneous velocity ug (x,, 8, C, w) is regressed against conditional velocity wvv, + w)n(x;;)) r) where we can approximate u(x,;"" r) with our network (Chen et al., 2025). The per-time and per-class Flow Matching term can be modified as
FPM (0) = Ex,.v.
[ s, c, w) - wV + (1 w) w)us (x.. S, Ø, 1))
(14)
where 0g denotes EMA weights. We show in Appendix A.6 that the minimizer of this objective coincides with the ground-truth CFG velocity in Eq. (13). Our class-conditional objective LivM (0) can be modified as
Ex.ic w
s,s,,, t, ds fo ||
d
ug: (xt + (Xe, t, 8, &, w,, 114s + Lop (0)
(15)
The weighting 1/w" to prevent exploding gradients because the magnitude of ground-truth ve- locity scales linearly with Final objective simply samples each of t, s, w under some distribution p(t, s)p(w) and computes the above loss in expectation. We randomly set c = Ø with some proba- Our training algorithm is shown in Algorithm 1.
bility (e.g. 10%) and for each Ø set
Sampling. Our construction can naturally interpolate between one-step and n-step sampling. See Figure 3 for PyTorch-style sampling code.
def sampl ing (net, X, ts = torch. torch. linspace 0, n+ for t,s in 1tp(ts(:),,ts111::
w):
x = x + (s-t) * net t, s, w) return
PRACTICAL CHALLENGES
Figure 3: PyTorch-style sampling code.
We note and address several challenges practically implement our objective.
Semi-Lipschitz control. Theorem 1 makes the crucial assumption that (x(x,; 8) is Lipschitz tinuous. However, modern transformers with scaled dot-product attention (SDPA) and LayerNorm (LN, Ba et al. (2016)) are not Lipschitz continuous (Kim al., 2021; Qi et al., 2023; Castin al., 2023). This issue similarly applies to diffusion transformers (DiT) (Peebles & Xie, 2023), Our insight is to make minimal and restriccit changes to the architecture for Lipschitz control.
As shown in Figure 4, the original DiT experiences training instability leading steep jump in network activations. As a solution, we adopt RMSNorm as QK-Norm, which coin-
AdaLN Act Norm
wo Lipschitz
w/ Lipschitz
sides with the proposed L2 QK-Norm (Qi al., 2023) with learnable scaling and is provably Lipschitz continuous. We also substitute all LN with RMSNorm (without learnable pa- rameters, denoted as RMSNorm (-)), whose Lipschitzness show in Appendix B. In addition, DiT introduces Adap- tive LayerNorm (AdaLN) where the output of RMSNorm is modulated by MIP outputs of time embeddings denoted as RMSNorm (ar) @ a()) + b(t) where is the input feature and a(t), b(t) are scale and shift respectively. However, the Lips- chitz constant of this layer depends on the magnitude of a(t)
100k
1524
2004
25dc
Training Steps
Figure 4: Activation norm of last time embedding layer. Same trends follow for all other layers.
which can grow unbounded and is subject instability. We therefore employ RMSNorm (.) again on all modulation parameters for
AdaL.N(r, t) RMSNorm (x) @ RMSNorm (a(1)) + RMSNorm (b(t))
(16)
Figure 4 also shows the activation with proposed changes. Activations stay smooth after our fixes. Finally, we follow Qi et al. (2023) and use Lipschitz initialization for all linear layers except for time embedding layers. Note that these modifications do not explicitly constrain the Lipschitz constants of all but the key layers where instability can arise. We find such partial control of the Lipschitzness is sufficient for empirical success.
5

---

Published as a conference paper at ICLR 2026
Flash Attention JVP with backward pass. The training objective involves the time derivative of our map fx t, s), which can be derived as
fo( 1, = Fo 2, 8) + (1
(17)
ds
where the last term involves differentiating through the network with Jacobian- bacobian-Vector Product (JVP). This poses significant challenge for transformers because automatic differentiation pack- ages, e.g. PyTorch, often do not efficiently handle JVP of SIPAA. Open-source Flash Attention (Dao ct al.. 2022) also has limited support for JVP. Crucially, different from prior works (Lu Song, 2024; Geng et al., 2025; Sabour et al., 2025), gradient is also propagated through the JVP term (x,,tx,tt,s To tackle these challenges, we propose an efficient Flash Attention kernel that (i) fuses JVP with forward pass, (ii) uses significantly less memory than naive PyTorch attention, and (iii) supports backward pass on JVP results. We detail the implementation in Appendix C.
Optimizer parameter change. Due to higher-order gradi- ent through JVP, our loss can be subject to fluctuation with the default AdamW 82 0.999. We take inspiration from language models (Touvron et al., 2023) for mitigation and use 82- 0.95 to speed up update of the gradient second moment. As show in Figure 5, the terminal velocity error fluctuates significantly less after B2 change.
Ba 0.999
3,00
Po = 0.95
3.90
5 1029
Mted
0,055
536
1806
2014
2528
Training Steps
Scaled parameterization. The ground-truth CFG velocity scales linearly in magnitude with w, so using neural net- works to directly predict the velocity may be suboptimal. We therefore additionally investigate simple scaled alternative as fx(x,, f, s, c, w) (s F) (Xx, 8, C, w) so that ux(x,, w) = FF (x., 8, w) which scales with w by design. We study the effect of this parameterization in experiments.
Figure 5: Smoother terminal velocity er- ror with B2 = 0.95.
Different time distribution for FM loss. We find it empirically helpful use separate distribution to sample different specifically for the FM loss (see Eq. (15)) is because we can directly transfer the proven successful time distribution from FM training for TVM. How this is used can be found in Algorithm and we ablate this decision in Section 7.3.
5 CONNECTION TO PRIOR WORKS
MeanFlow. MeanFlow (Geng et al., 2025) minimizes loss Ex,,,. F Fo t, -) Figil() where
Fig = u(x;,t) t) + (s - t) [u(x.,)) t). Fo, (X,, t, + F.Fo, (x),,,
(18)
This loss can be equivalently rewritten Ex,,t. [| || fo t, 8) + u(x;; t) [|2] where (x,, t, s)
(s-1)Fo (x,, I 8) and loss minimized if and only if d fo (x., This exhibits duality with our proposed method in that we enforce differential condition w.r.t. while MeanFlow differentiates w.r.t. / which requires u(x,,)) to be propagated through JVP. In prac- tice, u(x;, is replaced with v., which introduces additional variance during training and can cause fluctuation in gradient, especially under random CFG during training (see Section 7.2). Additionally, the relationship between the loss and distribution divergence remains elusive with the introduction of vt. In contrast, we show our loss upper bounds 2- Wasserstein distance, which provides unique insights on smoothness properties of network.
-u(x,,)) (see Appendix E. 1).
Physics Informed Distillation (PID). PID (Tee et al., 2024) as inspired by Physics Informed Neural Networks (Raissi al., 2019; Cuomo al., 2022) distills pretrained diffusion models u ux,x,, !) into
one-step samplers. It parameterizes the one-step net displacement as (x1, s) = - -1 (x),s) where X1 pi (x1) and trains via distillation loss
d
Ex,,
(x s) u + fo, (x1, 5),
(19)
ds
Our method generalizes the setting by introducing the starting time in addition to the terminal time Under this view, PID sets f. = 1 and can only generate one-step samples. We additionally show

---

Published as a conference paper at ICLR 2026
Figure One-step samples from TVM on both ImageNet-512x512 and ImageNet-256x2222
in Section 7.3 that naive combination of PID and FM loss suffers from optimization instability and a continuous distribution on is necessary for empirical success.
RELATED WORKS
Diffusion and Flow Matching. Diffusion models (Sohl-Dickstein al., 2015; Ho al., 2020; Song et al., 2020) learn generative models by reversing stochastic processes, while Flow Matching (Liu al., 2022; Lipman et al., 2022) generalizes this arbitrary priors with simplified training. Both approaches ultimately solve ODEs with neural networks during sampling.
One-Step and Few-Step Models from Scratch. To address slow inference from ODE simulation, recent methods aim for few-step generation in a single training stage. Consistency models (Song et al., 2023; Lu Song, 2024) parameterize networks to represent ODE integrals but cannot jump between arbitrary timesteps without injecting additional noise, which can limit multi-step perfor- mance. Two-time conditioned approaches enable arbitrary timestep transitions: IMM (Zhou et al., 2025) provides distribution consistency via Maximum Mean Discrepancy but requires multiple par- ticles; MeanFlow (Geng et al., 2025) and Flow Map Matching (Boffi et al., 2024) match trajec- tory derívatives but lack distributional guarantees. Other variants bypass differentiation via Monte Carlo (Liu & Yue, 2025) or combine distillation with FM (Frans et al., 2024).
Unlike these methods, TVM regularizes path behavior the terminal time rather than initial time and provides explicit 2- Wasserstein bounds. While sCT and MeanFlow only compute forward JVP, TVM uniquely supports backward passes through the JVP computation, These innovations drive both our theoretical insights and architectural improvements. Another recent work (Boffi et al., 2025) has proposed unifying perspective on training flow maps from scratch; encompasses many approaches such as Shortcut (Frans al., 2024), CM (Song al., 2023; Lu & Song, 2024), MeanFlow (Geng et al., 2025) and shares similar theoretical insights as TVM.
EXPERIMENTS
We investigate how well TVM can generate natural images (Section 7.1), discuss its advantages compared to previous methods (Section 7.2), ablate various practical choices (Section 7.3) and dis- cuss its computation cost (Section 7.4).
7.1 IMAGE GENERATION
1mageNet-256x2556 We present quantitative results in Table under FID (Heusel et al., 2017).
We adopt the default DiT-XL/2 architecture (Peebles & Xie, 2023) and inject - s as the second timestep, following IMM (Zhou al., 2025) and MeanFlow (Geng et al., 2025). We additionally employ our semi- Limi-hipschitz control techniques for training stability or we notice activation explo- sion as described in Figure 4, and train with constantly sampled CFG, i.e. models with

---

Published as a conference paper at ICLR 2026
(EB
(.)
NPR (U))
Dill
DiflsionFlew
ADM (Dharival & Nichol, 2021)
250x
10.96
554M
ADM (Dtaarie) & Nichol, 2021)
250x2
7.72
359M
LDM 4-G (Kombach et al., 2022)
250x2
3.00
400M
SimDiff ('loopeheem
202.
512x2
3,00
SCT-XL/2 (Ma al. 2024) (si 1.50
D/T-XL/2 (Peebles & Xie, 2025) (w 1.5
DII-X1/2 (Pechles & Xic. 2023) (u 1.26)
250x2
250xx
250x
2:27
2.15
5:22
675M
675M
0756
EDME XL (Kamas et al. 2024)
EDME L (Kamas at. 2024)
VIMM+ (Kingna Cor, (0244 20244
U (Bio 2l22)
250x2
512x2
63x2
6382
2,65
4,35
1.85
1.58
775M
501M
1.10
2B
OneFFww-Seep from Scratch
DXT XL/2 (Peclees & Xit. 20233 (or 1.25)
250x3
4,64
675M
iCT-XL2 (Song & Dhariwal, 2023)
34.24
675M
SIT-XE22 (Ma etal., 2024 (or 1.5)
DIT XL/2 Pechies & Xst, 20233) (ar 1.30
250x
250x2
3.04
2,62
675M
675M
20.3
Shonenn--X (Frans 2024)
10.60
OueFFos-Sepp from Seradd
IMM XL22
2025)
675M
SCT & Song, 2024)
5.15
778M
2x2
3.99
675M
4.85
775M
2x
2.51
673M
SCEX (1. Song. 2024)
4.33
1.18
MeanFlow XL/2 (Geng et al., 2025)
3.43
676M
3.73
2.93
576M
MennFiow (Geng 2025)
5.24
TVM-XL/2 (Ours) (w 2)
3.29
578M
3.17
2800
6786
TVM-XL/ (Ours) (t 2.50)
TVM-XL/2 (Ours) (w 1.75)
2.61
4.58
678M
678M
TVM-XL/2
1.99
2.94
Table I: FID results on 1mageNet-256x256.
Table 2: FID results on ImageNct-512x512
HeanFlow TVM
Newn
Crad Norm
45
TVM
H=255
TVM
coral aw2
10C Training Steps
Training Steps
28
10
Training Steps
Figure 7: (Left) McanFlow is subject to wide variation in gradient norm if CFG scales (i.c., A. and w) are randomly sampled under naive settings (sce Appendix F.2 for details). TVM shows much smoother gradient norm. (Middle) McanFlow's gradient norm strongly correlated with the fluctuation [|u(xe, t) TVM's []u, t)l is much more stable under the same CFG setting. (Right) Our method converges with random CFG at training time, although tradeoff exists between different CFG in FID. Constantly sampled CFG works best.
and w = 1.75 are two different models trained from scratch. We describe additional training de- tails in Appendix F. Our method achieves state-of-the-ar 1 N-NE FID among methods trained from scratch, outperforming MeanFlow and IMM. With CFG w 2, TVM can achieve noticeable im- provements over MeanFlow, e.g. 3.29 FID vs. 3.43 FID for I-NFE and similarly for 2-NFE. With 4 NFEs, 1.75 also exceeds 500-NFE diffusion baselines. We additionally show qualitative 1-NFE samples on the right of Figure 6.
ImageNet-512x512. We train with the same settings in FID scores in Table 2. We rerun MeanFlow using the same settings as in ImageNet-256x42 as our baseline in addition to sCT (Lu & Song, 2024) under similar model sizes. TVM again outperforms SCT and MeanFlow in I-NFE and 2-NFE regime. Notably, TVM-XL/2 outperforms SCT-XL with 1.1B parameters, highlighting TVM's more optimal use of model capacity in fitting the image dis- tribution. Moreover, with w = 2.25, TVM with 4-NFE can match 500-NFE DiT-XL/2 baseline in performance, further demonstrating the scalability of our algorithm to higher resolution.
and we show the
Intriguingly, for both datasets, TVM trained with higher CFG performs better on I NFE while worse on 2 NFEs. We believe this implies fundamental trade-off between different NFE quality and that the network is limited in capacity in fitting all NFEs well. We leave more detailed studies on this trf and any design improvements to future work.
7.2 DISCUSSION ON TRAINING ADVANTAGES
Single sample objective. Unlike IMM (Zhou et al., 2025) which uses more than 4 samples to cal- culate its loss, use a single sample to for loss calculation without losing a distribution-matee interpretation. This also allows the objective to be scaled to large models and high-dimensional datasets where batch size each GPU is constrained to be 1.
Training with random CFG. Our construction allows us to randomly sample CFG scale during training without collapse. We attribute this stability to our JVP being only caleulated w.r.t. which is invariant to starting position x, and time t. In contrast, CT (Song et al., 2023; Lu & Song. 2024) and Fean (Geng al., 2025) require velocity u(x,, 1) to be used in the JVP calculation. In
8

---

Published as a conference paper at ICLR 2026
trunc (Pe.)))-)4
FID
clamp (14, (2)
FID
gop (y.: (1., n.)
FID
475
trunc
( 0.4, LO), 0.4)..0)
4.59
2.0), 0.4, 1.0)
3.88
( L.), 0.4
5.12
clamp
L.).( 1.4)
4,00
(2.0. 1.0)( 0.4)
4.11
0.8, LO). 1..1.0)
3.72
gap
(.0), 1.0.
4.01
(2.0, 1.0).( 0.6, 1.0)
4.00
[ 1.4) 1.4)
3.95
(2.0, (-0.6,
7,88
(1.0,100) 1.0), 1.0)
3.66
(-1,0,12)
(-D44,
3.82
100
2014
2004
(1.0, (-0.1..))-
3,70
(.0), 0.4, 0.0,
3.83
1.4,,, (-04.110)
3,94
Training Steps
Table Ablation studies on different time sampling schemes, evaluated by I-NFE FID.
Figure 8: FID trend the sampling schemes.
p(w)
I-NFE
EMA rate
I-NFE
Scaled Parum.
I-NFE
2-NFE
% t=s
I-NFE
2-NFE
rand., 10 1.5
9.37
==
10.24
yes, w 2
3.72
3.35
0
3.72
3.35
rand., w
5.14
7 = 0.9
5.08
nO, 27 2
3.82
3.27
10%
3.91
3.18
const., w7 1.5
6.66
? 0.99
4.90
yes, w 1.5
6.04
4.60
3.88
2.97
const., W
4.81
0.999
6.04
no, 22
1.5
9.32
7.02
30%
3.97
3.07
(a) Random vs. constant CFG sampling evaluated target lig. at example w's.
(b) EMA of pseudo-
(c) With vs, without scaled pa- rameterization.
(d) Prob. for 1 = 8 dur- ing training.
Table 4: FID ablation on various samplinggpprramcteizzaiooo decisions.
the case of random CFG, this velocity can vary widely in magnitude which, if propagated through JVP, can cause wide fluctuation in gradient norm (see left two in Figure 7) and causes training instability. Our method, in comparison, enjoys much smoother gradient norm and u(x,, t) norm, and successfully converges even in the presence of random CFG. We note that random sampling of CFG is not optimal as some CFG scales experience degradation in FID during training, and constant CFG performs better in comparison. We postulate that the under-performance of random CFG is due to limited capacity of the network and the 1/w² factor that downweights high CFG. This phenoemenon is similarly observed CFG-conditioned FM training (see Appendix F.3) and we leave any improved design to future work.
No schedules and loss modification. We do not rely on training curriculum such as warmup sched- ules in sCT. For each CFG scale, we use the default CFG velocity for all s, while MeanFlow relies on additional hyperparameters to turn on CFG only when is within a predetermined range. We also strictly adhere to the simple C2 loss without any adaptive weighting as proposed by MeanFlow. We believe the simplicity in our design allows for more scalability.
7.3 ABLATION STUDIES
We ablate various implementation decisions and discuss insights from different parameter choices. Results are presented with XL/2 architecture trained for 200K steps with batch size 1024.
Time sampling. Similar to Flow Matching, different time sampling schemes can greatly affect performance. We explore 3 different kinds of sampling schemes.
Truncated sampling (trunc). Let (Me, ot), (Ha, oe) denote t being sampled from logit- normal distribution with mean and standard deviation (e), and beinsg sampled from trun- cated
distribution with parameters (Mg, o.)) such that s < t.
Clamped independent sampling (clamp) Let (Nt, (1))a)) (1s,0a) denote t. and s being inde- pendently sampled from logit-normal distributions with mean and standard deviation (14, ot) and (μs, J.), and set s if >
Truncated gap sampling (gap). Let (Hg) (Ha o.) denote the gap L-s being sampled from logit-normal distribution with mean and standard deviation (Hg. og), and s sampled from logit-normal with parameters (Ha o.) truncated at 1 - g. Then set t = + g.
In Table we show comparison within each sampling scheme and conclude that better results are ob- tained when is biased towards 1 and biased towards ( for the model to learn taking longer strides. However, biasing too much, e.g. f4 2.0, a; 2.0, leads to worse results. For gap, sampling t with lower mean is preferrable to higher mean. In Figure 8, we also observe t truncss performance degrades and cl amp plateaus faster than gap. Therefore gap wins over longer training horizons.

---

Published as a conference paper at ICLR 2026
All above sampling schemes follow the naive joint sampling of ts, We lastly explore separate time distribution for the FM loss term (see Section 4). We follow gap-sampler and denote the sampler gap* with parameters (/g.. )) (Ps, 0x) where (s. () is jointly sampled for the first loss term and (Ms, Tx) is used to construct new logit-normal distribution to independently sample s' for the FM loss. We find in Figure that gap* generally performs better in 1-NFE FID.
FID
gap (-0.8 1.0), (-0.4,
3.44
gap" (-0.8, (-0.4,004,,..) 1.0)
3.36
Figure 9: Sampler comparisons.
CFG sampling. As described in the previous section, due to limited capacity of the model, we observe tradeoff in performance when CFG is randomly sampled during training. This is reflected in Table 4a. We note that constant CFG always outperforms random CFG, and for constant CFG sampling we find w = 2 converging faster than the default w = 1.5 for Flow Matching.
EMA target rate The target EMA weight f* plays significant role in accelerating convergence of the model. Shown in Table 4b, non-EMA target, Le. ? 0, noticeably lags behind y > alternatives. However, too large of a y e.g. 0.9999, also causes instability because of the overly slow target update. A sweet spot exists around 7 = 0.99 which we use as default. Besides attribute its success to variance reduction because EMA's slower weight update implies much lower optimization noise. In addition, EMA is commonly used to evaluate diffusion models for its quality boost (Song et al., 2020), so being the optimization target also provides better learning signal to the model.
Scaled parameterization. In Table 4c, we find scaled parameterization is generally beneficial, but its benefit may vary depending on training/data settings. We therefore suggest testing different choices for different settings for best performance.
Inspired by MeanFlow (Geng al., 2025), we investigate whether setting 1 == 8 (when it reduces to pure FM training) is helpful for overall performance. We find that >> 0% actually degrades 1-NFE performance while it marginally improves 2-NFE performance. This trad- off persists throughout the training but we observe diminishing return as training goes on. Therefore, we do not find this practice helpful in general and leave out of our design space in general.
Probability for
7.4 MEMORY AND RUNTIME ANALYSIS
We analyze per-step runtime and Per-GPU memory consumption (averaged over 10 training steps without counting EMA up-
Runtime (s)
Memory (GB)
MeanFlow (w/ naive SDPA)
OOM
MeanFlow (w/ our kernel)
0.81
46.73
date cost) without any performance opti-
TVM (w/ improved DiT)
0.95
71.44
mization (e.g. torch. compile). Shown
TVM (w/ naive DIT)
0.86
59.53
on the right is a comparison with Mean- Flow using 256 batch size on 8-GPU H100 cluster on 2 Since JVP
TVM (w/ improved DIT, detach JVP)
0.69
55.71
Figure 10: Per-step time and per-GPU memory study.
with Flash Attention is not officially sup-
ported by PyTorch, the simplest way implement McanFlow is to naïve SDPA, which runs out of memory. MeanFlow with our kernel does not run OOM. TVM with Lipschitz control (Section 4) experiences higher runtime and memory mostly due to architectural change, since TVM with naive DiT is only marginally more expensive than MeanFlow with naive DIT. We note that much of the additional compute be compiled away via PyTorch. Additionally, if step time is concern, can simply detach the JVP which biases learning gradient but dramatically reduces runtime. We leave further efficiency optimization to future work.
CONCLUSION
We present Terminal Velocity Matching, framework for training one/few-step generative model from scratch. Different from prior works, we match the terminal velocity of flow trajectory instead of the initial velocity, and we show our objective can explicitly upper bound 2- Wassersseei distance up to a constant. Our proposed objective is conceptually simple and easy to implement, and our theory sheds light on flaws of current diffusion transformers for their lack Lipschitz continuity. TVM achieves state-of-the ow-toh---artttteoe art one-step result for model trained from scratch and surpasses base- line diffusion models with only 4 NFEs. We hope it can provide new insights into making scalable and performant one/few-step generative paradigms to come.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint (tt:::23330087797, 2023.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: 1607. 064000 2016.
Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A vit backbone for diffusion models. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 22669-22679, 2023,
Nicholas M Boffi, Michael S Albergo, and Eric Vanden-Ein Lunden-Eijnden. Flow map matching. arXiv preprint arXiv:2406.07507, 2, 2024.
Nicholas M Boffi, Michael Albergo, and Eric Vanden- How to build consistency model: Learning flow maps via self-distillation. arXiv preprint arXiv:2505.18825 2025.
Valérie Castin, Pierre Ablin, and Gabriel Peyré. wrXiv:2312.14820, 2023.
How smooth is attention? arXiv preprint
Huayu Chen, Kai Jiang, Kaiwen Zheng, Jianfei Chen, Hang Su, and Jun Zhu. Visual generation without guidance. arXiv preprint arXiv:2501.15420. 2025.
Salvatore Cuomo, Vincenzo Schiano Di Cola, Fabio Giampaolo, Gianluigi Rozza, Maziar Raissi, and Francesco Piccialli. Scientific machine learning through physics-informe informed neural networks: Where we are and what's 's next. Journal of Scientific Computing, 92(3):88, 2022.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in neural information processing systems, 35:16344-16359, 2022.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-8794, 2021.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference on Machine Learning, 2024.
Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel. One step diffusion via shortcut models. arXiv preprint arXiv:2410.12557, 2024.
Zhengyang Geng, Mingyang Deng, Xingjian Bai, Zico Kolter, and Kaiming He. Mean flows for one-step generative modeling. arXiv preprint 1724 2025.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho and Tim Salimans. 2rXiv:2207.12596 2022.
Classifier-free diffusion guidance.
arXiv preprint
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851. 2020.
Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. simple diffusion: End-to-end diffusion for high resolution images. In International Conference on Machine Learning, pp. 13213-13232. PMLR, 2023.
Tero Karras, Miik Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyz- ing and improving the training dynamics of diffusion models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 24174-24 184, 2024.

---

Published as a conference paper at ICLR 2026
Dongjun Kim, Chich-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning proba- bility flow ode trajectory of diffusion. arXiv preprint a:Xhv:2310.02279. 2023.
Hyunjik Kim, George Papamakarios, and Andriy Mnih. The lipschitz constant of self-attention. In International Conference on Machine Learning, pp. 5562-5571, PMLR, 2021.
Diederik Kingma and Ruiqi Gao. Understanding diffusion objectives as the elbo with simple data augmentation. Advances Neural Information Processing Systems, 36, 2024.
Lingxiao Li, Samuel Hurault, and Justin M Solomon. Self-consistent velocity matching of proba-
bility flows. Advances in Neural Information Processing Systems, 36:57038-57057, 2023.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Wenze Liu and Xiangyu Yue. Learning to integrate diffusion odes by averaging the derivatives. arXiv preprint arXiv:2505. 14502, 2025.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0300 00000,00003, 2022.
Cheng Lu and Yang Song. Simplifying, stabilizing and scaling continuous-timg consistency models. arXiv preprint arXiv:2410. 1081, 2024.
Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden-Eijnden, and Sain- ing Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. arXiv preprint 2:Xi::2401.087400 2024.
Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. CoRR, abs/1805.02867, 2018. URL http: / wwwpsiiii.rrr/aeoooo abs /1805. 02867.
OpenAl. Video generation models as world simulators. https: /pppappenaaaaao//aaa//eee/ hann/sora/,22/ 2024.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4195 4195-4205, 2023.
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhom, Jonas Müller, Joe Penna, and Robin Rombach. Sdxl: Improving latent diffusion models for high-resolution image synthesis. arXiv preprint arXiv:2307.01952. 2023.
Xianbiao Qi, Jianan Wang, Yihao Chen, Yukai Shi, and Lei Zhang. Lipsformer: Introducing lips- chitz continuity to vision transformers. arXiv preprint arXiv:2304.08850 2023.
Maziar Raissi, Paris Perdikaris, and George Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and putlern recognition, pp. 10684-10695, 2022.
Amirmojtaba Sabour, Sanja Fidler, and Karsten Kreis. Align your flow: Scaling continuous-time flow map distillation. arXiv preprint arXiv:2506. 14603, 2025.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing, pp. 2256-2265. PMLR, 2015.
Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. arXiv preprint arXiv:2310.14188 2023.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:201 134566, 2020.
12

---

Published as a conference paper at ICLR 2026
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01669, 2023.
Joshua Tian Jin Tee, Kang Zhang, Hee Suk Yoon, Dhananjaya Nagaraja Gowda, Chanwoo Kim, and Chang Yoo. Physics informed distillation for diffusion models. arXiv preprint arXiv:2411.0833 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Llama: Open and efficient foundation language models. arXiv preprint (tti:2302..39711, 2023.
Pascal Vincent. A connection between matching and denoising autoencoders. Neural compu- tation, 23(7):1611-1674 2011.
Team Wan, Ang Wang. Baole Ai, Bin Wen, Chaojie Mao, Chen-Wei Xie, Di Chen, Feiwu Yu, Haiming Zhao, Jianxiao Yang. et al. Wan: Open and advanced large-scale video generative models. arXiv preprint 25XX::25033.03144 2025.
Linqi Zhou, Stefano Ermon, and Jiaming Song. Inductive moment matching. In Forty-second International Conference on Machine Learning, 2025.
13

---

Published as a conference paper at ICLR 2026
A THEOREMS AND DERIVATIONS
A.1 GENERAL NETWORK PARAMETERIZATION
In general, we can parameterize our net displacement as
fg Xx, t, s) (tt,
(20)
for some 7(,,8) that satisfies r(t, 1) for boundary condition. And for the velocity condition,
let
d
ux(( ()
fx, L, aamt ox)FFo
(21)
ds
where = = s)(==
We derive
d ds [X,, 1, 80m below for clarity.
fg (x 1 8)
0,,(11,)) (F(x,)))((x(,,5)) +7((, (x) ==
(22)
ds
s=t
-,,(t, Fol t,t) + t(t, (x 5=
(23)
0,7((,)))== (x,,t)
(24)
(25)
A.2 TERMINAL VELOCITY ERROR UPPER BOUNDS DISPLACEMENT ERROR
Lemma 1. Under mild regularity assumptions, the following inequality holds,
C'assp (0) S { Exx. || a
[ - u 8,,
ds
(26)
where px(x.) is marginal distributions for initial points x..
Proof. We assume both displacement maps are Riemann-integrabee, then
fa(x,,t,0) - 1. u(x,, s) s)ds
Clispl (0) = (xim
(27)
1//0 d fo(xx t, s)ds - /' (((((xittss))) s)ds
(xix
(28)
ds
2
Exiepe(x))
(x(x,, t)
u(v(xx, s),
ds
(29)
ds
where (+) uses triangle inequality and regularity assumption.
A.3 TERMINAL VELOCITY ERROR REDUCES TO FM
Consider the terminal velocity error for each time s as
d tx - (x,, 8),
Ex,
(30)
ds
2
Expand the inner term
lo Tx( 1,8) Fo(xe, + (s -1),,F((x,,1,))
(31)
ds
14

---

Published as a conference paper at ICLR 2026
and for the inner norm term its limit exists as t
lim
fx(x,, t, u(u((x,, t, s).
(32)
(-
ds
lim
(o 4 + (s t)0, F Fo
(33)
1-5
Fo 8, 8
(34)
Thus, the limit of its expected L2 -o-morm exists (assuming this norm is bounded) and is equal to C2-norm of its limit, which
Ex,
[Fo(X,, - u(x., }]
(35)
and this is the original FM loss, which is equivalent (up constant) conditional Flow Matching loss used in practice in Eq. (1).
A.4 MAIN THEOREM
Theorem 1 (Connection to the 2- Wassersstei distance). Given [0,1], 1), let (P be the distribution pushforward from pc(x,) via fg (x,,t, 0), and assume uo(:, 8) Lipschizz-coniinen for all [0, (] with Lipschitz constants M(8), with additional mild regularity conditions,
WF po) S [ CTVM)) (0)ds + +C,
(10)
where W2(:.-) 2- W-ssserstein distance, {|-] is a functional of L(-), and C is non-optimizable constant.
Proof. Note that the ground-truth flow map &b is invertible and that (,,,(x,, t,0).0tt = X and (bb(( 0, t, Xy.
( to
(xx00)00((((x00000))) xxl1dxxo
(36)
/ po po(xo)llx (xx 0,4), (xx 0,4,,
(37)
/ (x)]lx + (x,, 0) - (xx,, t, t)
(38)
= / mm(xx))/ /.000 ds
fo(x,,t, t, - 1." u(xx, $ sdss
dxi
(39)
(0)
$ / px(x))) //" [ ds
2
fo(xx,,,))
ux) t, s), 8)
dsdx,
(40)
o(x....)
where (i) is due Wasserstein distance being the infimum of all couplings, and we choose par- ticular coupling of the two distribution by inverting with &b and remapping with respective flow maps. And (ti) is due Lemma Now, we inspect e(x,, t, s) specifically by noticing that E(x,,t, 1,
- fo (x t, u(bi(,,, s), s) + (x 8), u t, s),
+ 1g (x 1, 8), ug s), $]12
(41)
d
fg (x, t, (x, t, 8), 8) ds
+ (1 ux)(xx,, t, s)11
6(x()t.))
(lo s) - uo 2, 5), s)113
(42)
d
8(x,, + L(s)
fo (2xx,, t, u) ul uxbb(x,, t, u), u) du
(43)
ds
e(xeit,))
15

---

Published as a conference paper at ICLR 2026
where (i) is due triangle inequality and (it) is due to Lipschitz- coniciuuz- assumption. We further
notice that right-hand- sigtthand--sie contains a term that is the integral of the left-hand-side For simplicity, we hold x, and constant and let
y(s) - L. e(xe,t, t, uddu.
-
j(s) = -x(x,, t)
and we arrive at the following inequality,
(ss) - 6(x,,t, + ((s)((s))
(44)
-0(x) S i(s) + [())((8)
(45)
eli (())d5(xx,,,r)) 4,r) dr
d (())
(46)
(x,,, ( el7 L(a)du
(eli y())
(47)
f." elT (lo)dd8(xx,,))) S y015 eei ((w)duy(()))
(48)
elt eli y( -
ali (u)ddu((x t, r)dr
(49)
f eli (('L(u)u-S3 S2 8(x))u((x,,) rdr
u(s) -
(50)
u(s) < {"
ef." Lu)ddu(xx,, r)dr
(51)
Therefore, setting s 0 we have
L E(xx,t, u)du /," c L(u)du 8(x,, t, u)du
(52)
x(()(r)
where the left-hand side is the inner term of Eq. (40). Then,
Eq. (40) V / pi(x)) 1 A[T](()) 5(x,, 4, 4,8,,,
(53)
= 1. \[[[(s) Ex. [1 t)), 2
+ +luu(((xx,,,s) (,) uul s), s) ds
(54)
1. \[[(()) Ex. [|o((x),, uo((fxxx s)1
(xo
+ Ex,
(bxe,) t,,, - uu t, s), 8) 3)| ds
(55)
= 1.' {[4](s)
Ex. [|10(xts (0(x))),,) (x 8)
+ Ex, [| 8) u(x,, s) ds
(56)
where (a) can be rewritten as
(a) Ex..v.
[| s) -val(]] +Õ
(57)
where Ò is some non-optimizable constant (Lipman et al., 2022). This is also a classical result connecting score matching and denoising score matching (Vincent, 201 1).
16

---

Published as a conference paper at ICLR 2026
Now, after substitution, notice that our bound in Eq. (56) becomes
L. XLM + C
(58)
where some other constant, which completes the proof.
A.5 REDUCTION TO FL.W MATCHING
When show that LTv reduces to Flow Matching loss.
C. (0) Ex.xx...
d
(x,,,t, us (x, .). 8)
(x.. s) v.l
(59)
ds
s=t
[[gg( upl (x, 1) [xo(xe; vi
(60)
A.6 DERIVATION FOR CONDITINNN. TRAINING TARGET
In Eq. (14), we introduced the CFG training target
wv, + (1 w w)uby (x,, 8, Ø
We derive below that the minimizer of Eq. (14) is the CFG velocity wu(x,, 8, c) + (1 w)u(x,, s).
Proof. Consider the training objective (without weighting for simplicity)
Ex Nx,,w
uu (xx. s,
[wvv + (1 w))b (x.. 8, 2))
(61)
when 2, w 1, then it reduces
Ex.. ||| [X,: (x,, 2)) v. v
(62)
with the minimizer min satisfying u (xe, Ø) = s).
At minimum of the loss for other w and must satisfy
u Ummm (x,, 8, c) Ev, (VV, (1 w)) u (X,, 8, Ø) Xx 8, w|
(63)
wBE, [vo Xxy 8, cl + (1 w wu) (x Ø) wu(xx, 8, + (1 s)
(64)
(65)
B ADDITIONAL DETAILS ON PRACTICAL CHALLENGES
B. LIPSCHITZNESS OF RMSNORM
Recall the definition of RMSNorm, for input x e Id and a small constant é
RMSNorm
RMS(2)) where RMS(x)
1
Σ
24 +
(66)
And its Jacobian can be calculated as
RMSNorm(1.)
d
24 RMS(x)
(67)
dæy
da;
RMS (2r) 2T, / (RMS(2))
(68)
RMS(2) x
dej
(69)
RMS(2 (z) d RMS(r)3
17

---

Published as a conference paper at ICLR 2026
def get..f.and (net, xt, t, 5, C, w):
def model model wwwpperpp....... (x., t.,, s.): # we use t-s for second time condition return net t., (t. - 8.), c, w)
F, dFds funch.func.. wroper_,r (xt, s), C8. 1))
f.ts xt + (s -stt)
dfds (F (s
dFds)
return fits, dfds
Figure 11: PyTorch-tyyee JVP code.
Since matrix norm (largest singular value) a (A) of matrix A is upper bounded by its Frobenius norm, and RMS(x) 2 e, we have each element dæ,
RMSNorm (r;) in the Jacobian matrix bounded
via
d RMSNorm(Z.) dæ
dij RMS xz)
T d RMS 3
(70)
51
rv/Vd
2g/VV RMS(2))
(71)
RMS
RMS(x)
RMS (x)
1
(72)
(73)
Therefore, the Frobenius norm is bounded and hence the matrix norm.
B.2 FULL DESCRIPTION OF NORMALIZATION OF MODULATION
Note that there are 6 modulation parameters in total for each DiT layer, denoted as
(1).). by (t), ea (t), (t), (2()), (t) split(AdaLN Modulation(t). 6)
(74)
and we pass each of the above parameters through RMSNorm (+) to obtain
(t), 67 (t), c7 (t), (t), 65 (t), c (t)
(which can be done parallel) and the new normalized DiT layer is
G (t) ATTN(RMSNomm' (x) aj (f) + bj (1))
x C5 MLP(RMSNorm (xx) 02 (t) + by (t))
C FLASH ATTENTION JVP WITH BACKWARD PASS
transformer models, scaled dot-product attention (SDPA) often among the most, not the most, computationally expensive operations. The cost stems not only from its high FLOP requirements O(MN) in general, and O in the case of self-attention but also from the quadratic memory footprint of the query- key matrix multiplication.
Computing the Jacobian- Vector Product (JVP) of SDPA is even more demanding, typically requiring about three times the cost of the standard forward pass. Flash attention (Dao et al., 2022) fuses the matrix multiplication with an online sofumax operation (Milakov Gimelshein, 2018), thereby eliminating the need to store the intermediate QK matrix in GPU memory. Subsequent work has shown that JVP SDPA can also be implemented in a FlashAttention-stylee manner, where both, primal SDPA and JVP SDPA are computed jointly to avoid redundant computation (Lu Song, 2024).
Building on these ideas, we implement efficient JVP SDPA forward and backward kernels in Triton. We first take inspiration from open-souree implementations without backward support'. And the
'https: / thub. wwm/mommaaaoooooooooooo 845, 5./ sCM/b' ob/ma n/standa one_mul head_jvp_ test.py
18

---

Published as a conference paper at ICLR 2026
additional backward pass through the standard ("primal") SDPA is handled independently using the open-source implementation from (Dao et al., 2022). To obtain full gradients with respect to Q, K, and V. we combine the imput gradients from both backward passes.
Similar to standard SDPA, the JVP backward pass can leverage online softmax to avoid storing large intermediate matrices in GPU memory. However, the increased complexity of JVP SDPA requires additional optimizations to run efficiently on GPUs. Most notably, we found it crucial to split the backward computation into multiple smaller kernels reduce register spills caused by the large number of intermediate tensors.
Background. Recall the attention operation
QKT Vdk.
ATTN(Q,K, K, v) = V softmax
(75)
and let the query, key, and value blocks be denoted by Q RAxx K Rxd and V RRxd The tangent inputs denoted as Q. K, V. We denotes the sog-sun-exponettial normalization for the i-th row of the attention scores, short form for combining the softmax stabilization factor and the normalization.
V as the softmax scaling factor, and &
C.1 MULTI-STEP BACKWARD PASS
For best performance, we decided to split up the backward pass into multiple smaller operations with shared paths through the graph. Furthermore, the gradients dQ and dQ are computed in row-parallel order, while dK, dK, dV and dV are processed in column-paralle order. In our tests, redundant, coalesced computation of the large parts of the backward pass greatly outperformed a single, fused kernel relying atomic operations.
We split the operation into 6 steps: 1) preprocess shared intermediates, 2) process dK and first part of dK, 3) process d and first part of dQ, 4) process second part of dK, 5) process second part of dQ, 6) process dV and dV.
Step 1: Preprocess shared intermediates row- wwr-aarllll. In the first step, we preprocess two inter- mediate sums Σ₁ E R and Eg E I used in steps 2-5.
51==Epp Σ1 = Poj (dovT).
(76)
Σg Σ Pij ((1δVT), (aòrt) Nij
(77)
where
de
Ps = exp (aS.s - b(),
Siz= = Σo
QirKrr
No = a$i
(78)
4
r=1
amd
de
Š. = Σ +
p4 p:= EEp, Pvj (aSt.)
(79)
Step 2: process dk and dK, column-parallel.
[[K])==a>]] ((UOTT) -I1s) Pg Qu:
(80)
(aR));= = Σ ((
- E1i Psj Q
(81)
19

---

Published as a conference paper at ICLR 2026
Step 3: Process dQ and dQ1 row-parallel
(dQ1);; = Σ (aóvT), -En,s Psg
Kj;
(82)
(d)).: a Σ [((OTT)) - -11
Pij
Kjz
(83)
Step 4: Process dK column-parallel.
(dK),, = (dKi), + Σ a (-E))) Sig + En. 14
Po
(84)
+ [(ort)) + (dórt) (asis
Mi
E2.
Pis
Qi.
4.
Step 5: Process dQ row-parallel
(dQ).: (dQ1);;: Σ a (- Seg E1, H4
Ps
(85)
+ +((uOrT), + (dort), (asis
Hi
12, Pss Kg
1,
Step 6: Process dV and dv column-parallel.
(dV)s: = Σ (o)())
(86)
(dV);: = Σ Pij a$. (4:) (ao)... )
(87)
Caching softmax statistics. Like previous flash-attention implementations, statistics from the forward pass to speed up the backward pass, namely the log-sum-exp l, the sums I and μ. for each row of the output O. Thus, the total overhead of the cache only three values per row of Q.
cache softmax
C.2 EVALUATION
We built test bench to evaluate latency and peak memory consumption of our flash JVP SDPA kernels on different input shapes using an NVIDIA H100 SXM 80GB. Due to the lack of existing alternatives, we compare against vanilla SDPA. i.e. SDPA written as explicit math operations, which currently is the only way to train transformers in PyTorch with JVP enabled.
As our contribution focuses on the backward pass, we limit the latency and peak memory evalu- ation to the backward pass of single SDPA operation, combining both paths through the primal ("normal") and the tangent (JVP) gradients.
Results. Shown in Table 5, our implementation achieves significant reduction in peak memory consumption. Compared to the reference, we save memory not only by reducing the cached variables between forward and backward pass, but more importantly by avoiding to store N intermediate attention scores. At the same time, our implementation achieves a speedup of up to 65% compared to the reference.
20

---

Published as a conference paper at ICLR 2026
Latency [ms]
Peak Memory [MB]
H
S
ours
vanilla
ours
vanilla
128
1.31
1,51
64.69
64.80
1,024
1.38
1.54
69.52
94.02
1
4,096
1.96
1.53
86.06
508.1
8,192
3.98
4.33
108.1
1,816
1
16,384
10.06
16.11
152.3
7,024
32,768
40.24
63.85
240.5
27,808
24
128
1.40
1,55
80.55
83.17
24
1,024
1.42
2.03
196.4
784.4
24
4,096
15.13
24.52
593.5
10,721
24
8,192
58.70
96.93
1,123
42,115
24
16,384
238.4
2,182
24
32,768
958.6
4,300
Table 5: Performance comparison of our flash JVP kernels against vanilla SDPA kernels PyTorch. H and stand for number of heads in multi head attention and sequence length. Vanilla SDPA ran out of memory on NVIDIA H100 in the last two tests.
Algorithm TVM Training
Input: initialized model Sd, data po e) and prior pi (x1), time distribution p(l, s), guidance distribution p(w)
Initialize 8* 0,0 t 0
11 10 are EMA with rate Ar2
while model not converged do
Sample (Xp, c, x1) ~ Po c)p1 (x1)
Randomly drop with prob. 10%
Sample (f, s, w) ~ p(t, s)p(w)
// optionally sample ~ p(s') for the second loss term.
(1 x)x0 txi
X, f- [1 - 8)xo + 8X1
Il optionally set x + (1 - 8')x0 + 8'x1.
X1 Xo
I/ optionally set Vol (Tri (0)
Xy.
< optimizer step by minimizing (TVM(0) = Er,,, // optionally use s and xy for the second loss
// see Eq. (15)
term
0 + EMA update with rate X
0 <- EMA update with rate *
end while
Output: learned model f
D TRAINING ALGORITHM
We present the training algorithm in Algorithm 1. We additionally show PyTorch-style pseudo- d de (x,, together with one JVP pass.
code Figure for calculating (x,, t, s) and
E RELATION TO OTHER WORKS
E. MEANFLOW
Let (x,,t,, (ss-t] ( s,), we inspect
Lol 1,) + u(xx;!) dl
(88)
--Fo(xe,ts) t, + (s +) d F Fo t, s) u(x(,t)
(89)
dt
=-Fo(x,,t + t) Vxe Vx, Fo(x,, t,
0, (x(, + u(xe,t)
(90)
21

---

Published as a conference paper at ICLR 2026
ImageNet-256x256
ImageNet-6122512
Parameterization Setting
Architecture
DiT-XL/2
DiT-XL/2
DIT-XL/2
DFT-X1/2
Params (M)
678
678
678
678
2dd time conditioning
8
t-8 8
t- 8
t- 8
Hidden dim
1152
1152
1152
1152
Number heads
Main normalization
RMS Norm
RMS Norm RMS Norm
RMS Norm
RMS Norm
QK-Norm type
RMS Norm
RMS Norm
RMS Norm
Linear layer init
Spectral
Spectral
Spectral
Spectral
Time Embed init
Spectral
NO, 0.02)
[00. 0.02)
Spectral 300K
Training iter
300K
300K
300K
Training Setting
Optimizer
AdamW
A
Adam 10 8
A
Optimizer c
10 8
10 8
10 8
31
0.9
0.9
0.9
0.9
B2
0.95
0.95
0.95
0.95
Learning rate
0.0001
0.0001
0.0001
0.0001
Weight decay
0
0
Batch size
2048
2048
2048
2048
p(x, I)
gap* (-0.8, 1.0),
(-0.4, 1.0)
Scaled param.
yes
yes
no
yes
% f s
0%
0%
0%
0%
w
1.75
2.5
2.25
Target EMA rate
0.99
0.99
0.99
0.99
Eval EMA rate*
0.9999
0.9999
0.9999
0.9999
Label dropout
0.1
0.1
0.1
0.1
Table 6: Experimental settings for different architectures and datasets.
Therefore,
d fo (x t, u(xxx,t) t) dt
(91)
2
12
Po(x,,4,8) + (8 1)
u(xe, 1) Vx, Fo(x,,t, + O.F (x,,(x),4)
u(xx,() ()
(92)
Fig
which is the MeanFlow loss.
E.2 FLOW MAP SELF-DISTILLATION
Boffi et al. (2025) proposes three different self-distillation approaches for training flow maps from scratch.
F ADDITIONAL EXPERIMENT DETAILS
We present the overall training details in Table 6
'Except for zero-init layers in AdaLN-Zero.
All time embedding MI.P layers before input into DiT blocks.
This means the percentage of setting = 8, e.g. 0% means both loss terms exist at all times.
EMA used for final evaluation, separate from the EMA used for training target.
22

---

Published as a conference paper at ICLR 2026
F.1 ARCHITECTURE AND OPTIMIZATION
VAE. We follow Zhou et al. (2025) for the VAE setting, which uses the standard Stable Diffusion VAE (Rombach et al., 2022) but with different scale and shift. Please refer to the paper for details.
Architecture. All architecture decisions follow DiT (Peebles & Xie, 2023) except for the changes described in the main text. For our XL-sized model, we follow DiT-XL and use 1152 hidden size but use 18 heads instead of 16 heads. This is purely for efficiency reasons because 18 heads under 1152 total hidden size implies head dimension is 64, while the original 16 heads result in head dimension 72. Flash attention JVP's runtime is sensitive to redundancy in memory allocations. As 64 is power of 2 our kernel can fully allocate appropriately sized CUDA blocks, while 72 leaves significant chunks unused. We observe that the original 16-head decision is x 1.25 slower than the 18-head variant. In comparing FID of the two versions, we observe they perform similarly throughout training.
Following Zhou et al. (2025), we use - s as our second time condition into the architecture rather than directly injecting For injecting w, we follow Chen et al. (2025) and use / = 1/w our 1 and set w = 1/3. Note
condition, and random CFG is used training, we sample 3 ~ Lt( Wan wau
that Chen et al. (2025) uses 3 ~ Ld(0, 1) which amounts to Wmin = 1 and Wmax cc, but arbitrarily large w is never used practice so Wmax can be set to a realistic finite value.
Optimization. Besides setting 32 = 0.95, we follow the default optimizer used by DiT and optimize with 6F16 precision. We de use any learning rate scheduler.
F.2 DETAILS ON RANDOM CFG WITH MEANFLOW
In MeanFlow (Geng al., 2025), the authors introduce mixing scale k such that the field with guidance scale w is given by
v(x,, 1, c, w) wv, + (x, w) + (1 w (x) 4, w)
(93)
It specifies that the effective guidance scale is w/ π.- This is because since us c, w) v(x,,, t, c, w), rearranging LHS and dividing both sides by (1 x) gives
(1 k)v w) = wv. (1 k))L (xe, w) v(xi,,,c,w) w) (1
(94)
w
w
+ (1
uo (xxx w.
(95)
This constrains E [0, 1). However, in the case of random CFG, to make use of us (xt, 4, w)), we try the simple linear mixing (the default CFG reweighting)
vxx,, L c, + x) = wV, suu,(,, (,c, (1 w (x,,o(
(96)
where w and i are both randomly sampled with finite boundaries. In this case (x,x,, f, 1) & v(x,, L, e, w +*) and thus x is not constrained to be smaller than When w = O, it becomes regular CFG with network approximation of the CFG velocity, and when 0 it becomes MeanFlow CFG with v. approximation of the CFG velocity. This construction subsumes both implementation cases. In experiments,
21(1, Cmax) for some constant Cmax- However, we acknowledge that this observed training fluctuation may depend on exact training settings and environments, and may be fixable via empirical tricks such adjusting AdamW parameters or gradient clipping, etc. We present the training in the simplest settings without such tricks to best illustrate our point.
K (0, Cmax),
F.3 CFG-CONDITIONED FLOW MATCHING
As in our method, we similarly observe tradeoff in FID if FM is trained to condition on CFG scale w with randomly sampled during training (Chen al., 2025). During in- ference time, w is injected into the network so that the CFG velocity field can be approximated by single forward call. We inject w using positional embedding just like the diffu- sion time, and during training sample ~ U(00) 1) and
22
1=1.5
w=2
100k
150k
23
Training Steps
Figure 12: experiences tradeoff
w-conditioned FM training

---

Published as a conference paper at ICLR 2026
set w 1/8, following Chen et al. (2025). We show in Fig- ure 12 that as the model trains, the FID of w = 1.5 decreases but = 2 increases for later training steps. This tradeoff is similarly observed in our method as presented in the main text.
24

---

Published as a conference paper at ICLR 2026
F.4
ADDITIONAL VISUAL SAMPLES
Figure 13; Additional ImageNet-25665 256x256 samples from 1-NFE TVM model.
25


Reviews:
[{"review_id": "nzoQw91JIe", "reviewer": "Reviewer_32jy", "summary": "This paper proposes Terminal Velocity Matching (TVM), a single-stage objective that generalizes flow matching by learning ODE integrals between any diffusion timesteps via terminal velocity matching. TVM upper-bounds 2-Wasserstein distance under Lipschitz assumptions. On ImageNet-256×256, TVM achieves 3.30 FID in 1-NFE (SOTA) and a new few-step Pareto frontier (2.49 FID at 2-NFE), with natural n-step interpolation.", "questions": "1. How sensitive is performance to the time sampling distribution $p(t,s)$? The paper uses a biased gap-based scheme — would uniform sampling or curriculum hurt?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. TVM reframes the problem of learning long-horizon ODE jumps as a terminal velocity condition (Eq. 6–7), providing a clean theoretical link between displacement error and velocity matching.\n2. Theorem 1 establishes a distribution-level guarantee ($W_2$ upper bound) without requiring multiple particles (unlike IMM).\n3. Duality with MeanFlow is clearly articulated: The paper shows MeanFlow matches initial velocity while TVM matches terminal velocity (Appendix E.1), offering a compelling symmetry and explaining why TVM is more stable under random CFG.\n4. This paper introduces a custom Flash Attention + JVP kernel to accelerate.", "weaknesses": "1. While inference is fast, training cost (FLOPs, GPU-hours) vs. MeanFlow, sCT, or IMM is not reported.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2dFBYqOBsp", "reviewer": "Reviewer_4c17", "summary": "This paper introduces Terminal Velocity Matching (TVM), a new training objective for generative modeling that generalizes Flow Matching (FM) by matching velocity fields at terminal rather than initial points along probability flow trajectories. The authors prove that the TVM loss upper bounds the 2-Wasserstein distance between data and model distributions, providing a principled connection between training dynamics and optimal transport theory. They also propose practical architectural modifications—such as semi-Lipschitz normalization, FlashAttention-based Jacobian-vector products, and scaled parameterization—that stabilize training and enable few-step sampling. Empirically, TVM achieves state-of-the-art single- and few-step image generation performance on ImageNet-256 with a DiT backbone.", "questions": "1. Can the authors clarify why terminal velocity matching improves convergence over initial or midpoint matching?\n2. How tight is the Wasserstein upper bound empirically—has this been tested or approximated?\n3. Is the FlashAttention JVP essential for convergence, or mainly for speed?\n4. Could the method generalize to other architectures or modalities (e.g., U-Nets, text-to-image, or video)?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The theoretical formulation is elegant and well-motivated, linking Flow Matching to a Wasserstein upper bound through terminal velocity constraints.\n2. The method achieves excellent efficiency–quality trade-offs, outperforming existing one-step and few-step baselines (e.g., Consistency Models, MeanFlow) on ImageNet-256.\n3. The architectural refinements (semi-Lipschitz normalization, FlashAttention JVP) are practically valuable contributions that could generalize to other diffusion or flow-based frameworks.", "weaknesses": "1. While theoretically grounded, the intuition behind “terminal velocity” could be elaborated further—especially how it differs in practice from midpoint or integral matching.\n2. The scope of evaluation is limited to class-conditional ImageNet-256. Demonstrating robustness on higher-resolution or unconditional datasets (e.g., ImageNet-512, COCO) would strengthen generality claims.\n3. The paper relies on a single architecture (DiT-XL/2). It is unclear whether TVM’s benefits extend to U-Net–based or continuous-time flow models.\n4. The FlashAttention JVP optimization, though valuable, feels somewhat tangential to the conceptual contribution and could be relegated to an appendix.\n5. Despite the theoretical connection to Wasserstein distance, there is no empirical verification of this bound (e.g., via transport cost evaluation), leaving its practical tightness untested.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fRs09xlwKl", "reviewer": "Reviewer_ep7Z", "summary": "The authors propose Terminal Velocity Matching (TVM), a few-step generation framework that matches terminal-time velocities instead of initial-time ones. To make TVM scalable for transformer architectures, the paper introduces minimal architectural modifications (e.g., RMSNorm-based AdaLN, Lipschitz initialization) and an efficient Flash-Attention JVP kernel that supports backward passes through Jacobian–Vector Products (JVPs). Combined with a simple training recipe—without curriculum training schedules or adaptive weighting—TVM achieves strong results on ImageNet-256×256, setting a new Pareto frontier for few-step generative models (FID = 3.30 @ 1-NFE, 2.49 @ 2-NFE) while maintaining training stability.", "questions": "1. **Distributional guarantees vs. flow-map loss.**\n\nThe authors mention that Flow Map methods lack explicit distributional control, whereas TVM provides a 2-Wasserstein control. Could the authors elaborate on the difference? At a glance, the TVM loss resembles the Lagrangian distillation loss proposed in Flow Map—what key term enforces distributional regularity here?\n\n2. **Effect of removing forced FM loss.**\n\nTVM achieves superior 1-step generation performance even with 0% enforced Flow Matching samples. What intuition do the authors have for why completely removing the auxiliary FM supervision leads to better convergence?\n\n3. **EMA acceleration.**\n\nWhat is the authors’ intuition behind the observed acceleration of convergence when using EMA weights for the training objective? Does this act primarily as a variance-reduction mechanism?\n\n4. **Distillation potential.**\n\nHave the authors explored or considered using TVM as a distillation objective? Since TVM naturally supports terminal-condition regularization, it could serve as a stable distillation formulation for distilling multi-step teacher models into few-step generators.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. **Boundary condition at terminal time**\n\nBy enforcing velocity matching at the terminal rather than initial timestep, the method avoids evaluating JVPs involving guided velocities that often exhibit large norms and high variance during training. This could be more beneficial for stabilizing training when scaling to larger dimensionality, where the guided velocities often exhibit even larger norms and higher variance.\n\n2. **Simple and stable training recipe**\n\nThe paper achieves stable one-stage training with only minor modifications to model architecture and without requiring any curriculum schedule, auxiliary objectives, or adaptive weighting heuristics. These simplifications make the approach both reproducible and again appealing for scaling to large models or high-dimensional data.\n\n3. **Customized JVP kernel supporting efficient backpropagation**\n\nThe authors develop a fused Flash-Attention JVP kernel that supports backward passes on Jacobian–Vector Products within transformer blocks. This kernel significantly reduces memory usage and improves runtime efficiency during training. This engineering contribution is non-trivial and essential for making TVM practical at scale.", "weaknesses": "1. **Backpropagation through JVP.**\n\nUnlike prior continuous-time consistency models where JVP terms are detached from the gradient graph (sCM, MeanFlow, etc.), TVM explicitly backpropagates through the JVP term introducing additional computational cost. This could become prohibitive for large-scale models. Providing quantitative analysis (e.g., runtime, memory, or gradient-cost overhead relative to MeanFlow) would help stress the concern.\n\n2. **Insufficient ablations.**\n\nWhile the design choices are well-motivated, the paper would benefit from deeper empirical validation. For instance, ablations comparing training with vs. without adaptive weighting on the training loss, the convergence of the original DiT versus the Lipschitz-controlled variant, and MeanFlow vs. TVM under identical random classifier-free guidance (CFG) conditions. Reporting gradient-norm / JVP-norm comparisons between TVM and MeanFlow would further support the claim of reduced JVP instability. Including recent baselines such as FACM (Peng et al., 2025) would also improve completeness.\n\n3. **Need for deeper insight behind ablation trends,**\n\nThe empirical results raise several interesting but unexplained phenomena. For example, using any proportion of $t=s$ (pure flow-matching loss) degrades 1-step performance, and employing EMA targets accelerates convergence. A deeper discussion of why these trends arise would enrich the paper’s interpretability and practical guidance.\n\nPeng, Yansong, Kai Zhu, Yu Liu, Pingyu Wu, Hebei Li, Xiaoyan Sun, and Feng Wu. \"Flow-anchored consistency models.\" arXiv preprint arXiv:2507.03738 (2025).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DbrjvuLbW8", "reviewer": "Reviewer_EXtX", "summary": "The paper proposes a novel method for training one or few-step generative models by matching the terminal velocity. Through a clever rearrangement of the displacement equation between arbitrary points t and s along the ODE trajectory, the paper derives a novel objective which contains the instantaneous velocity at s and the generic Flow Matching loss. The formulation requires high order gradients through the JVP operation, but proposes improvements to the JVP computation to ensure scalability. The method further introduces improvements to DiT architecture and a way to include CFG conditioning. The method is tested on the common benchmark Imagenet 256x256 and achieves SOTA results for one and few-step generation.", "questions": "- 1) Maybe I missed it but I could not find the exact number of training iterations for the results in table 3, same for batch size.\n- 2) There also seem to be no info about the total training time. The method uses a forward pass with JVP plus a standard forward pass with the EMA model, and backward pass through JVP. Intuitively this should be much slower than MeanFlow or sCT, and it would be interesting to understand by how much.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The loss formulation is novel and well motivated, the paper introduces several relevant contribution such as DiT modification, scalable JVP computation, novel CFG training procedure. The results are strong and the method is scalable.", "weaknesses": "To me a weakness of the method it is hard to tell from the manuscript to what we can attribute the good performance. The overall results are great so there is no doubt that the proposed method is effective, but I wonder for example what results one could get from MeanFlow when using the improved DiT architecture. TVM is good but introduces some complications, such as using a EMA version of the model in the target, and having to backpropagate through JVP. While the latter is addressed, I wonder if just the improved architecture can work well on MeanFlow. I checked through the text and it doesn't specify whether meanflow is trained with standard DiT, so I am not sure what the retrained results from table 1 correspond to.\n\nA common benchmark for these models is CIFAR10 but that was not reported. While it is a relatively small dataset, it can provide useful insight on the performance of the model without CFG, as well as on the sensitivity of the model to specific hyperparameters.\n\nOverall, I think this method is promising, but seems like many details are left out, which makes it hard to fully understand the significance of the contribution (see questions). Also there is no open source code.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)