Flow Map Learning Via Non-Gradient Vector Flow
Authors: Mark Goldstein
Anshuk Uppal
Raghav Singhal
Aahlad Manas Puli
Rajesh Ranganath
Venue: ICLR 2026 Poster (2026)
Keywords: flow map
flow matching
diffusion
density
ode
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: training flow maps from scratch with jvps without differentiating through nested model calls

Abstract: Diffusion and flow-based models benefit from simple regression losses, but inference (i.e, producing samples) incurs significant computational overhead because it requires integration. Consistency models address this overhead by directly learning the flow maps along the ODE trajectory, revealing a design space for the learning problem between one-step and many-step approaches. However, existing consistency training methods feature computational challenges such as requiring model inverses or backpropagation through iterated model calls, and do not always prove that the desired ODE flow map is a solution to the loss. We introduce SGFlow, an approach for learning flow maps that bypasses explicit invertibility constraints and expensive differentiation through model iteration. SGFlow trains a model to compute both the ODE solutions and the implied velocity from scratch by following non-conservative dynamics with a stationary point at the desired flow map. On the CIFAR image benchmark, SGFlow attains a favorable relationship of FID to step count, relative to flow matching, MeanFlow, and several other flow map learning methods.

Published as a conference paper at ICLR 2026
FLOW MAP LEARNING VIA NONGRADIENT VECTOR FLOW
Mark Goldstein' ; A'nshuk Uppal?, Raghav Singhal', Aahlad Aahlad Puli", & Rajesh Ranganath"
1 Center for Computational Mathematics, Flatiron Institute.
DDepartment of Applied Mathematics and Computer Science, Technical University of Denmark. Courant Institute, New York University.
ABSTRACT
Diffusion and flow-based models benefit from simple regression losses, but infer- ence (i.e, producing samples) incurs significant computational overhead because it requires integration. Consistency models address this overhead by directly learning the flow maps along the ODE trajectory, revealing a design space for the learning problem between one-step and many-step approaches. However, existing consis- tency training methods feature computational challenges such as requiring model inverses or backpropagation through iterated model calls, and do not always prove that the desired ODE flow map a solution the loss. We introduce SGFlow, an approach for learning flow maps that bypasses explicit invertibility constraints and expensive differentiation through model iteration. SGFlow trains a model compute both the ODE solutions and the implied velocity from scratch by following non-conservative dynamics with a stationary point at the desired flow map. On the CIFAR image benchmark, SGFlow attains favorable relationship of FID to step count, relative to flow matching, MeanFlow, and several other flow map learning methods.
INTRODUCTION
Diffusion and flow models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2020; Kingma al., 2021; Albergo and -aindenEijnden, 2022; Singhal al., 2023; Pandey and Mandt, 2023; Bartosh et al., 2024; Singhal et al., 2024; Albergo et al., 2023; Lipman et al., 2022; Liu et al., 2022) have improved generation in domains such as proteins (Abramson et al., 2024) and images (Peebles and Xie, 2023; Esser et al., 2024). Sampling from these models typically requires numerically integrating an ordinary or stochastic differential equation. Numerical integration requires multiple forward passes of neural network, leading to increased sampling latency and cost.
To ameliorate this generation cost by changing the training, recent approaches for consistency modeling and map matching (Song al., 2023; Song and Dhariwal, 2023; Kim al., 2023; Lu and Song. 2024; Boffi et al., 2024; 2025) aim to lear direct mappings from noise to intermediate or final data points along trajectories defined by probability flow ODEs, thereby avoiding costly integration. However, the methods have their respective complexities. For example, flow map matching requires model invertibility, while consistency models need either map in one step or introduce extra steps that leave the target ODE trajectory.
We introduce SGFlow (for StopGrad Flow), an approach that builds on flows and map matching methods and
Has true flow map as a unique stationary point
Does not restrict the class of neural networks used (e.g., to invertible functions)
Does require auxiliary losses involving invertibility or adversarial optimization
Does not require optimizing through nested calls to the model
Allows for generation along the ODE trajectory with any number of steps
"mgolésteing n@ rom.nsatttti

---

Published as a conference paper at ICLR 2026
Methods
Multistep
Regression
Required Prove Optimum
Nesting
Consusency Distiliaioon (Bong 2023)
Consistency Triining (Song et ol., 2023)
x
x
Consiseency Trieeciory Models (Kim 2023)
/
x
/
L-MM (Both et al., 2024)
/
/
/
x
LSD (Bofh 2025)
/
8
x
/
/
ESD (Boft et al. 2025)
/
8
/
x
/
x
PSD (Both ol., 2025)
/
/
/
x
/
/
MesnFlow (Geng 2025)
/
X
8
SGPlow (this work)
Table 1: Comparison to prior works. We categorize consistency modeling (flow map learning) techniques and our proposed SGFlow method according to: (1) ability adjust sampling steps post-training, (2) whether they follow the PF-ODE (Song et al., 2020), (3) whether they allow simulation-free training, (4) whether their objectives use regression, (5) whether model inversion required during training, (6) whether the true flow map is proven to be optimal or stationary, and (7) whether differentiation passes through nested model calls See Section for details.
Existing methods for learning flow maps fall into few categories in terms of their challenges; all challenges relate to the idea that a flow map is characterized by certain derivative properties and that losses minimize squared error to make these properties hold. Flow map matching and related methods rely on a fundamental relationship between invertible mappings and ordinary differential equations (ODEs). This relationship typically requires explicitly computing both the forward map (the model being trained) and its inverse during training, complicating training, and requires explicitly materializing large derivative matrix in the forward pass. Other methods, such as MeanFlow (Geng et al., 2025), do not explicitly enforce the model inverse identities, but do not prove that their loss is minimized at the true flow map or that the optimization stationary at that solution.
SGFlow avoids the complexity of tracking a model and its inverse, as well as materializing derivatives, by exploiting an alternate identity involving only Jacobian-vector products (JVPs) without inverse functions. This identity allows us to formulate the objective purely in terms of the forward map, without needing explicit access to its inverse. Since solutions ODEs naturally produce invertible mappings, the SGFlow objective implicitly encourages invertibility without explicitly enforcing it. Thus, optimality, SGFlow yields a continuously differentiable function that precisely integrates the velocity field, directly generating the desired data distribution. We summarize the trade-offs among recent methods in Table and in Section 5.
Experimentally, for basic training setup using the same common architecture, we ask how flow matching. McanFlow, SGFlow, Lagrangian Map Matching, Eulerian Map Matching, and Progressive Map Matching compare in moderate dimensions (CIFAR-10) on unconditional metries (FID) when decreasing the number of sampling steps. We additionally compare GPU memory usage for these methods.
BACKGROUND
Stochastic interpolants (Lipman et al., 2022; Albergo al., 2023), and more generally most diffusion and flow methods, hereafter just. flows pose generative modeling as transport of a simple base density target density. Interpolants tackle the problem follows. For E [0, 1]:
01 = and oo = 0. Commonly, a, = 1 t and oe = t. 2. Define X, = XxXo + o,X1 for base density X, ~ qu and data X1 ~ q1 (or vice versa).
1. Choose (a,, .τ.) where
3. Learn produce new samples along the trajectory of densities.
For function f. let f. 4 fc. Thus X, ái + 0,X1. It follows that X, has density qe satisfying:
0. = - (z)),(((,
xl,z))== E[X X. =]]
(I)
where is called the velocity. The PDE in Equation (1) is derived in the above works. To accomplish step three, one starts by making the observation that a density satisfies Equation (1) if and only if it is the density of the solution to the probability flow ODE dx edl integrated forward from
2

---

Published as a conference paper at ICLR 2026
Xu ~ qu or in reverse from X1 91 (Albergo and Hunden-Eijndens 2024). One then proceeds by first approximating u using the following (simulation-free loss:
(ve) EE | (t, Xx) - (o xx Nomo
which has minimizer = v and then solving da = vodt.
Background on Consistency Methods. Sampling from flows requires integration where each integration step evaluates a neural network to modeling a score, velocity, or similar. Therefore, knowing integrals of or similar quantities directly could, in principle, speed up sampling. The goal of consistency and map matching methods is to learn map along the trajectory implied by the optimal v. We review an example here, with others described in Section 5. Song et al. (2023); Song and Dhariwal (2023) seek to learn a mapping ÿ that maps interpolant samples X, q: to Xo. the t t= 0 solution to dr = edt when starting at X; (note that Xo usually differs from the independent endpoint sample Xo used to draw X, under the interpolant). The loss measures the distance between modeled outputs when evaluated at two different nearby points. Let sglğ) indicate stopgrad (i.e., bookkeeping term as a constant when computing loss gradients). Then the loss is:
Consistency( (g) == [xx) [x (di Xx), sglil At, At, X(-::))
The sample Xi- used in the target should optimally come from integrating the true velocity or an approximation ty small step At from Xt, where uy either comes from a pretrained diffusion model, or (vg.9) g) are derived from one another. In practice, further approximations are used to compute It-At. Approximations are introduced because X.- is not simply defined by drawing a second interpolant sample at a smaller noise level, but instead corresponds to integrating the velocity one step from XG; the velocity is unknown and thus may come from pre-trained model, which increases training cost and may not be good approximation in the first place.
It is challenging to directly learn solutions in just one step. While this allows multistep sampling, the re-noising step necessarily takes the trajectory off the probability- ODE and the resulting updates no longer correspond to integrating the PF-ODE solution. Highlighting the issue with methods featuring time argument, Kim et al. (2023) note that this CM multistep sampling approach "cxhibits degrading sample quality with increasing NFE, lacking a clear trade-off between computational budget (NFE) and sample fidelity". In practice, a number of training-time or inference-time changes made to this setup to try break apart the problem into somewhere between step and the hundreds of steps used by diffusions (Song et al., 2023; Lu and Song, 2024; Kim et al., 2023; Boffi al., 2024; Sabour al., 2025; Geng al., 2025; Zhou al., 2025). We discuss the various solutions and their trade-offs in Section 5.
3 METHOD
We present SGFlow, a method for learning to solve the probability flow ODE without adversarial training, without model inverse during training, without representing explicit derivative matrices, and without costly simulations from pretrained models. SGFlow trains a model to compute both the ODE solutions and the implied velocity from scratch by following non-conservative dynamics.
Consider a two-time map f that for t < u brings X₃ up to xu by solving the probability flow ODE dæ vdt. Such an f that integrates v can be defined as follows:
f(t, u, x) = x + J. ((s, = x + J." v(s, f(t, s, x))ds
Differentiating the recursive form on the RHS w.r.t. t using the total (material) derivative yields:
a.f + (a.f)u(t, 2²) O, fu, u, x
(5)
This is uniquely solved at the true flow map f. We can square the left-hand side for parameterized fo and take an expectation over X,
L = Ex. [1 + (0 (x,,
(6)
The true map I is the unique minimizer of this loss. Using v(l,,x) EEXX can expand,
L = =xx, [10.10 - (an.o) - (1(,..o (
(7)
3

---

Published as a conference paper at ICLR 2026
We then use the parameterization fo (u -t)ff(t, u.) for underlying model fo. The parameterization yields two useful properties:
time derivative: (tfolt,t, t, fo),
Jacobian: (t =1
Using these properties and evaluating at u, we see that the minimization of eq. (7) reduces to flow matching where fo (t, t)) is trained match the velocity:
L Ex, Exx -Xel1%)
(8)
which reveals that for the true f, we have that
- a.f(t,t, = E[X,|XX =
(9)
This motivates replacing the unknown in eq. (7) with stopgrad (,. The stopgrad is used under the principle that since the original did not provide gradient updates for f. neither should term that approximates it. The SGFlow method follows the negative gradient with respect to 0 of:
Lug =
(0,fo) X,12 ((0a..))
sg sg(fol(axxx)))1 (10)
where sg() means stopgrad() and fo(t, u. = x + (u - fa u,.). The expectation is taken over X₁ sampled by drawing data X1, noise Xo. and computing X, = axXo + o,X1 and X, do dX
In practice, the PDE must hold for all pairs t I u. Let gtt, u) be a joint distribution with support over t ts and with positive probability and Leg (gt(, sg) We now connect L and Lag formally. Theorem shows that optimization of and Lgg stop at the same solutions.
Take expectations over time and define L = E,, w) [Z]
Theorem 1. Let qt, u) be a joint distribution with support over t S u and with positive probability on t = u. Let the family include functions f that are continuously differentiable in all arguments. Let X, = a,Xp + o,X, and X. = á,Xo + o,X,, Define f(t, u, := x + (u - t) u, x)). Take expectations
((X)) Let sg stand for sop-gradient Define L = REg(t, [T.] and Lig = Eygt, Then fo is a stationary point of Lig w.r.t F if and only if f is a stationary point of L w.r.t. F.
This is shown in Section A.5.
The main point of the theorem is to establish that Leg has the same set of solutions as L despite not having access to v. The intuition is that, despite the stopgrad, when { = u, Lag tries match the velocity. We show that Lug is inaccurate, so the optimization continues moving and does not become stuck at functions that distill an incorrect velocity. As this match improves so does the match between the parameter updates from fay and C * u. The main reason this works is that f(t, r) appears in other terms outside the stopgrad, and those terms tell it where go. This is crucial and not all stopgrad optimizations benefit from this property.
Intuition.
stationary point when this velocity estimate is
Computation. Both terms in Lug can be computed as expected squared norms of Jacobian-vector products (JVPs), which use forward-mode autodifferentiation avoid explicitly materializing Jaco- bians, saving memory. Using Py Torch notation,
JVP[S. (t, .)) (a, c)) = (0.f) a (0.l) (0.f)
for (0.f, 0, f..f evaluated at (t, u, x)) For the first loss term, = 1, 0 and X, and for the second loss term, 0, and X, sg x,)] X, sgif(t, t,)] Though we have two distinct JVPs, we can split the batch and randomly assign either pair of (a, c) values to each batch element.
4

---

Published as a conference paper at ICLR 2026
Nongradient Flow Following the update rules of Lag does not correspond following the gradients of any one scalar objective J (section B.1). This is because the optimization dynamics are in general non-conservative. The stopgrad structure breaks the symmetry required for the updates to be the gradient single scalar function. In this sense, SGFlow is formally a (two-player) game rather than standard gradient descent one potential function, alheit trivial one where the main player controls all parameters except those in the stopgrad, and the stopgrad player keeps a virtual copy of the parameters that simply equal the first players parameters. Put another way, the optimization dynamics if taken in the limit of small step size correspond to non-conservative non non-gradient vector flow.
4
EXPERIMENTS
4.1 IMAGE MODELING ON CIFAR-10
Architecture. We modify the existing diffusion U-Net from Dhariwal and Nichol (2021) by embedding both and with the usual Fourier embeddings and then concatenate on input to a small MLP that maps the two times to a hidden representation for use in the network. We use 128 channels and channel multipliers set to (1,2,2,2) with attent ion set to (False, False, True, False).
Training Settings. We use dropout 0.1. We do not condition on the class label We use at 11-1 t and o, = 6, with noise at Xo and data at Xi. We train for 200,000 steps at learning rate 2e-4.
Losses. We benchmark based on Flow Matching (Lipman et al., 2022) as the reference method for sampling in the O(100) step regime. For few-step sampling, we benchmark with Meanflow (Geng et al., 2025), which StopGrad's s all model derivatives in the loss to avoid backpropagation through differentiation. We also benchmark other recently proposed methods from Boffi al. (2024): the Lagrangian loss, Eulerian loss, and Progressive loss (discussed in Section 5).
Method
10 steps
50 steps
100 steps
theory
Flow Matching
24.87
3.53
3.05
yes
Lagrange
248.76
230.43
221.22
yes
Euler
77.19
66.99
38.95
yes
Progressive
337.36
235.20
206.18
yes
Meanflow
37.32
4.54
4.23
SGFlow
12.26
2.88
2.81
yes
Table 2: FID scores versus sampling steps on CIFAR-10 computed from 50,000 EMA samples after 200,000 training steps. The "theory" column means whether the stationary points of the optimization have been proven to exist if and only if the function is the flow map that integrates the ODE.
Results. We report the Frechet Inception Distance (FID) (Heusel et al., 2017) in Table 2. We find that SGFlow produces better FID than MeanFlow each choice of sampling steps for the given (rather common for CIFAR) training configuration.
4.2 MEMORY USAGE
In the last column of Table 1, we note whether a method requires differentiation through an iterated model call. As discussed in Section 5, for a flow map model fa = x + (u - L). the Lagrangian loss involves nested evaluation fo (u, u, (o, u, r)), with a similar nesting for the Progressive loss. The Eulerian loss requires computing V ft,, u, x)) folt, t, Xt), which entails product-rule expansion. Here we empirically compare the peak GPU memory usage during the backward pass for different losses, holding the architecture and data size fixed (the U-Net with batch of CIFAR images).
Results. We demonstrate the peak GPU memory usage during backward pass in Table 3. As expected, the Lagrange, Euler, and Progressive losses are the most memory-intensive, reflecting the need to backpropagate through nested model evaluations or product-rule terms. On the other end, MeanFlow, which computes Jacobian vector products but detaches the full JVP, has the lowest memory usage, though the effect of this detachment optimization dynamics is not analyzed in that work. SGFlow falls between these extremes, striking a balance between memory usage, preserving the optimum, and benefiting from optimizing through the model derivatives in the loss.
5

---

Published as a conference paper at ICLR 2026
Flow Matching
MeanFlow
SGFlow
Lagrange
Euler
Progressive
16.8 Gb
14.2 Gb
43.2 Gb
69.8 Gb
69.8b
54.3 Gb
Table 3: Peak GPU memory usage during backward pass (in Gb). Values reflect the maximum allocated memory measured across the training step's backward pass (i.e., during gradient computa- tion). SGFlow strikes balance in GPU memory usage, preserves the true flow map an optimum, and optimizes through all model derivatives without detaching.
5 RELATED WORK
Sampling from continuous-time generative models such as diffusion and flow models requires numerical integration. Each integration step requires forward pass of neural network, leading to computational costs and slow sampling. Current approaches to address this cost can be broadly categorized into two types (1) distilling pretrained diffusion or flow model into a few-step solver (Salimans and Ho, 2022; Kim et al., 2023; Liu et al., 2023), and (2) learning a few-step solver (Zhou et al., 2025). Some approaches in this area allow for distillation as well as training from scratch (Song et al., 2023; Boffi et al., 2024; Boffi and Vanden-Eijnden, 2023).
Consistency Models (CMs) (Song et al., 2023; Song and Dhariwal, 2023; Lu and Song, 2024) learn a one-step map from noise to data, either by distilling a pretrained model by learning from scratch. Distillation requires sampling trajectories from the teacher model. To allow for more steps after either training approach, CMs iteratively re-noise the one-step solution back to successively smaller time under the interpolant and then denoise, but this can take the solver off the probability flow.
Consistency trajectory models (CTMs) (Kim et al., 2023) extend CMs to learn two-time maps using a combination of consistency and adversarial objectives, which requires training an additional discriminator model (Goodfellow et al., 2014). CTM and Gameflow both target the same mathematical object, the probability-flow ODE flow map (i.e., the integral of the ODE), but they learn this map through different means. CTM learns the map by distilling a teacher solver, and the losses for teacher and student involve several nested model evaluations (with data at Zo. for 0 < s S u S t S 1, the teacher integrates from to u. then jumps from s, then from s to 0; and the student jumps from t to s and then to 0). The objective depends on chosen feature-space distance and, in practice, includes DSM and GAN terms that further influence the optimum. Consequently. the CTM loss is sensitive to the quality of the ODE discretization used by the teacher (in practice CTM finds the need to use 2nd order solver during training) and necessitates the presence of the GAN.
Inductive Moment Matching (IMM) (Zhou et al., 2025) learns few-step model via implicit generative model trained with MMD (Smola et al., 2006; Gretton et al., 2012), with the MMD which is estimated biasedly within subsets of data. In practice, the authors must use time-weightee schedules and specific curriculum/indnetiv procedure to stabilize optimization. While IMM produces high-quality image samples, it solves the problem of marginally sampling the data distribution rather than sampling along a probability flow, where the latter is the task studied in this work.
MeanFlow (Geng et al., 2025) derives a bVsed objective for flow maps. They train f to bring xu. down to X, via the parameterization fe = x + (u t). and train via the following loss:
Ameanllow == (t,, u, Xt) - a + (u- (sg 0...fllt...x. iu + fg(0..olt......))
(11)
Applying the stopgrad sg to all model derivatives improves efficiency, but there are no differentiated loss terms that encourage the model derivatives a and az to move toward the true flow map derivatives. This contrasts SGFlow where sg,fo(t, t, x)) is used in place of E[X, Xe), but where another term in the loss trains these two quantities to match. Finally, between equations (10, 11) in Geng et al. (2025), v is replaced with X; where it appears quadratically, thereby pulling an expectation through a square and missing a resulting trace covariance term (Boffi et al., 2025).
Flow Map Matching (Boffi et al., 2024) learns a two-time flow map. This allows for mapping along the probability flow in either direction, without adversarial training. Their Lagrangian Flow Map Matching loss requires only time derivatives, but an additional invertibility loss encouraging invertibility via time-swapping so that f(t, u, fu, t, x)) se x. While straightforward to compute, gradient steps require evaluating the model and its inverse at each step of training.

---

Published as a conference paper at ICLR 2026
Boffi et al. (2025) optimize velocity matching along with one of the three flow map functionals. For the parameterization fa := (u f). the first one is:
LSD:
clagrange = 1.1 u, Xx) So(u, u, So(t, x.)II1
1,
LSD comes from the condition (t, u, v(u, v(u. ft, X.)) and uses f t, -)(,,)) f(t,t,)) = v(t, )). It variant of the Lagrangian loss from Boffi et al. (2024). The
next is:
ESD: Li == Fll u, X,) + VaJolt, X,) X.) (.))l
ESD comes from condition (1, x) + (I, (t), = 0 where in the loss, is replaced with f(t, 2)). The last one, for intermediate time m, is:
PSD:
CProgress L
== Flllso(t, u, Xx) So(m, u, Solt, m, x.))
PSD comes from the composition property: f(l, x) [(m, u, [(l, m, x)) for in E (4, u). SGFlow, LSD, ESD, and PSD all aim to enforce similar set of flow map properties. Clagrange and Crougess must optimize through nested model evaluations, doubling the computational graph for backprop. Couler optimizes through V=fo fo. which causes mode autodiiferreniaation to invoke the full product rule term, V V fel /V VoV=fal fo + ffl Vefe. To alleviate these challenges and improve optimization, Boffi et al. (2025) suggest specific StopGrad options for LSD,, ESD, and PSD to improve optimization. It would be interesting to more deeply understand why these stopgrads help given their empirical success for flow map losses.
Distillation methods. A complementary line of work approaches flow map learning by distilling the outputs of pretrained flow matching velocity models into few-step solvers. Specifically the unknown v() in the flow map identities is taken to be pretrained network. By contrast, we emphasize training from scratch, avoiding dependence on a teacher model and ensuring that all components of the flow map are learned end-to-end. That said, distillation can be attractive in practice when pretrained model is already trusted (when up corresponds to the endpoint distributions and the chosen o, and ot, when the objective is weaker-for example, marginally sample from the approximated data distribution without explicitly solving the probability flow ODE).
6 DISCUSSION AND LIMITATIONS
Identifying functions through PDEs Given PDE solved by sought-after mapping f. featuring combination of terms such as 0. Ōzf that should be set to 0, which terms should be parameterized by the model and which should be approximated as part of ground-truth loss target? If the equations can be rewritten in several ways, which yield easier or more challenging objectives? Answering this is applicable to improving training objectives for generative models as well more generally solving PDEs with machine learning.
Invertibility The loss targets an invertible function at optimum. To simplify training, we explicitly give up knowing the inverse, meaning that we only learn maps in one direction. Luckily, this is the usual scenario for generative modeling. For likelihoods, one can still substitute - f for vo in the probability flow ODE (Song et al., 2021; Boffi and Vanden-Eijnden 2023). Thus this method can be seen from the perspective of training a normalizing flow (Tabak and Vanden- -ijnden, 2010; Tabak and Turner, 2013; Rezende and Mohamed, 2015; Papamakarios et al., 2021) without requiring the invertible architecture or inverse-dependent loss.
Architectures. SGFlow, Flow Map Matching, Simplified Consistency Models, and MeanFlow all specify models whose time-derivative equal the target of diffusion model training. but directly adapt architectures meant for diffusion models themselves. Example architectures used in these works are the UNet from Dhariwal and Nichol (2021), the diffusion transformer from Peebles and Xie (2023); Ma et al. (2024), and the EDM architecture from Karras et al. (2022; 2024). These architectures may thus be suboptimal for the problem at hand, precisely because the target of interest is defined as function often computed in many diffusion model forward passes (an integral). In this work, compute limitations did not allow for the thorough exploration of architectures, but the authors believe that rethinking architectures is a convincing direction to improve the quality and training-effie -efiinngffficincyy of learned flow maps.
7

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This work was partly supported by the NIH/NHLBI Award RO1HL148248, NSF Award 1922658 NRT-HDR: FUTURE Foundations, Translation, and Responsibility for Data Science, NSF CAREER Award 2145542, NSF Award 2404476, ONR 900014-23-1-2633 Optum, and Apple. This work was also supported by IITP with grant funded by the MSIT of the Republic of Korea in connection with the Global AI Frontier Lab International Collaborative Research. Mark Goldstein would like to thank the Simons Foundation and Flatiron Institute for funding and compute resources that supported this project. Anshuk Uppal would like to thank the Centre for Basic Machine Learning Research in Life Sciences, and Technical University Denmark.
8

---

Published as a conference paper at ICLR 2026
REFERENCES
Abramson, J., Adler, J., Dunger. J., Evans, R., Green, T., Pritzel, A., Ronneberger, O., Willmore.
L.. Ballard, A. J., Bambrick, J.. et al. (2024), Accurate structure prediction of biomolccular interactions with alphafold 3. Nature, 630(8016):493-00,
Albergo, M. S., Boffi, N. M., and Vanden- E. (2023). Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797
Albergo, M. S. and Vanden-Fijnden, H. (2022). Building normalizing flows with stochastic inter- polants. arXiv preprint arXiv:2209.155: 1557I.
Albergo, M. S. and Vanden-Eijnden, E. (2024). Learning sample better. Journal of Statistical Mechanics: Theory and Experiment, 2024(10):004114.
Bartosh, G., Vetrov, D. P., and Andersson Naesseth, C. (2024). Neural flow diffusion models: Learnable forward process for improved diffusion modelling. Advances in Neural Information Processing Systems, 37:73952-73985.
Boffi, N. M., Albergo, M. S., and Vanden-Fijnden, E. (2024). Flow map matching. arXiv preprint 24XY::2466007077.
Boffi, N. M., Albergo, M. S., and Vanden-Eijnden, E. (2025). How to build a consistency model: Learning flow maps via self self-distillation. arXiv preprint arXiv:2505.18825
Boffi, N. M. and Vanden-Eijndews E. (2023). Probability flow solution of the fokker-planck equation. Machine Learning: Science and Technology, 4(3):035012.
Dhariwal, P. and Nichol, A. (2021). Diffusion models beat gans image synthesis. Advances in neural information processing systems, 34:8780-8794.
Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F., et al. (2024). Scaling rectified flow transformers for high-resolution image synthesis, 2024. URL https:/arriii.or//ab22440...0000 030.002263.000... 2.
Geng, Z., Deng, M., Bai, X., Kolter, J. Z., and He, K. (2025). Mean flows for one-step generative modeling. arXiv preprint 13X4
Goodfellow, I. J., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. (2014). Generative adversarial nets. Advances neural information processing
systems, 27.
Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., and Smola, A. (2012). A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. (2017). Gans trained by a two time-scale update rule converge local nash equilibrium. Advances in neural information processing systems, 30.
Ho, J., Jain, A., and Abbeel, P. (2020). Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851.
Karras, T., Aittala, M., Aila, T., and Laine, S. (2022). Elucidating the design space of diffusion- based generative models. Advances in neural information processing systems, 35:26565-26577
Karras, T., Aittala, M., Lehtinen, J., Hellsten, J., Aila, T., and Laine, S. (2024). Analyzing and improving the training dynamics of diffusion models. In Proceedings the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 24174-24184. 24184.
Kim, D., Lai, C.-H., Liao, W.-H., Murata, N., Takida, Y., Uesaka, T., He, Y., Mitsufuji, Y., and Ermon, S. (2023). Consistency trajectory models: Learning probability flow ode trajectory of diffusion. arXiv preprint arXiv:2310.02279
Kingma, D. P., Salimans, T., Poole, B., and Ho, J. (2021). Variational diffusion models. arXiv preprint ar 1rXiv:2177.006
Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Le, M. (2022). Flow matching for generative modeling. arXiv preprint 10.00.2210.02747.

---

Published as a conference paper at ICLR 2026
Liu, X., Gong, C., and Liu, Q. (2022). Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0300 030003
Liu, X., Zhang, X., Ma, J., Peng, J., et al. (2023). Instaflow: One step is enough for high-quality diffusion diffusion-based based text-to-image generation. In The Twelfth International Conference on Learning Representations.
Lu, C. and Song, Y. (2024). Simplifying, stabilizing and scaling continuous-time consistency models. arXiv preprint ar Xiv:2410.11081.
Ma, N., Goldstein, M., Albergo. M. S., Boffi, N. M., Hunden-Eijndent E., and Xie, S. (2024). Sit: Exploring flow and tiffed generative models with scalable interpolant transformers. In European Conference on Computer Vision, pages 23-40. Springer.
Pandey, K. and Mandt, S. (2023). A complete recipe for diffusion generative models. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4261-4272
Papamakarios, G., Nalisnick, E.,, Rezende, D. J., Mohamed, S., and Lakshminarayaaaa, B. (2021). Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22(57):1-664.
Peebles, W. and Xie, S. (2023). Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pages 4195-4205.
Rezende, D. and Mohamed, S. (2015). Variational inference with normalizing flows. In International conference on machine learning. pages 1530-1538. PMLR.
Sabour, A., Fidler, S., and Kreis, K. (2025). Align your flow: Scaling continuous-time flow map distillation. arXiv preprint arXiv:2506. 14603.
Salimans, T. and Ho, J. (2022). Progressive distillation for fast sampling of diffusion models. arXiv preprint 2xx::2202..0512.
Singhal, R., Goldstein. M., and Rangamath, R. (2023). Where to diffuse, how to diffuse, and how to get back: Automated learning for multivariate diffusions. arXiv preprint arXiv:2302.07261.
Singhal, R., Goldstein, M., and Ranganath, R. (2024). What's the score? automated denoising score matching for nonlinear diffusions. In International Conference on Machine Learning, pages 45734 45758. PMLR.
Smola, J., Gretton, A., and Borgwardt, K. (2006). Maximum mean discrepancy. In 13th
international conference, ICONIP, volume 6.
Sohk-Dickstein, J., Weiss, E.,, Maheswaranathan, N., and Ganguli, S. (2015). Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning. pages 2256-2265, pmlr.
Song, Y. and Dhariwal, P. (2023). Improved techniques for training consistency models. arXiv preprint arXiv:2310. 14189.
Song, Y., Dhariwal, P., Chen, M., and Sutskever, I. (2023). Consistency models. arXiv preprint 263.4.2303.01469
Song, Y., Durkan, C., Murray, I., and Ermon, S. (2021). Maximum likelihood training of score-based diffusion models. Advances in neural information processing systems, 34:1415-1428,
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. (2020). Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:20/ 13456.
Tabak, E. G. and Turner, C. V. (2013). A family of nonparametric density estimation algorithms. Communications on Pure and Applied Mathematics, 1652):145-164.
Tabak, E. G. and Vanden-Eijnden E. (2010). Density estimation by dual ascent of the log-likelihood. Commun. Math. Sci. (1) 217- 233, March 2010.
arXiv preprint Zhou, L., Ermon, S., and Song, J. (2025). Inductive moment matching. a5xiv:2503.07565.
10

---

Published as a conference paper at ICLR 2026
A PROOFS FOR STATIONARY POINTS
A. I
FIRST VARIATION DEFINITIONS
We consider scalar-valued loss functions L: F R that map a function / E F a real value.
Define the tangent space (F) at f. This space contains functions h E (T))) such that there exists
curve indexed by scalar such that for each fe E F, and we have that fo - f and (4l (= h.
The first variation 8L of such a functional L evaluated at f E F in direction h E (F(F) defined as:
8f]f:h] h] = d Llf chil cm
(12)
We then have that f is a stationary point w.r.t. F if 8[f; = for all (F(F).
A.2 STOPGRAD FOR FUNCTIONALS
We define the stopgrad symbol sg for functional as follows. Let be functional that maps two functions f.g to real value. Let [[ be functional that is written in terms of O with symbol as [[ = Off. sg[f]], then we evaluate the following two quantities as follows
[[ Olf.f] f] scIf: h] = 80 f: h.
(13)
(14)
That is, the functional evaluates as usual but in first variation, we do not perturb terms in sg. This corresponds to the stopgrad detach() operation used in machine learning code with autodifferentia- tion.

---

Published as a conference paper at ICLR 2026
A.3 FIRST VARIATION OF ORIGINAL LOss FUNCTIONAL
Our functional C[f acts on functions f. According to the definitions in Section A.I, we need to compute 8C[f: /|] (1t [ff + ehl). The functional is:
[ I11(0./((.....) (0..) (u..ann. f=x (w-
T=[f] (1(a.. (s, (dr ( [
A[o] [i T2l
[ = ((((t.)).g(x)).gtx))) [l]]]
Lets define the path fo by replacing f with fc := f + eh. Then:
x (u x=
f)fee = x + (u - (u= + eh) - I + (u-t)f + e(u - 1)h
(15)
Then
d c + - =[
Lde AIf.]
= =E( li T + +el]
d de
Talā + chi]
(16)
dc
dc
We first compute this derivative and then evaluate it at So
SzS. 0, x (u- (u- +el 17n]
(17)
de + 0 [(u- t)f} + ed: (1r t)/h (u (u-t)a,f - f + e (1 - t)0,h - h)
(18)
(19)
(u - (t(a) a.hh ( thh
(20)
and
0zfe O2 x + (u - t)f e(u - t)h = I (u t)oz(f + ch)
(21)
and
d de
0.
de [u - 1)(8.) + +0,h) - (j + tti]
(22)
de (u t)a. + + - t)0ch
ch
(23)
ce(u - t)0,h eh
(24)
de
(u 1)0ch
(25)
and
d
dzfe
I + (u - t) + +hi]
(26)
de
de
d
d (u - 1)0.
I +
de d (u- 1)0
(27)
de
de
(u t)0ph
(28)
For the first term,
Ti ehl 110rSe + (O.
Diffcrentiating
-- - + ehi = =hh]=
d.
= + (OnSt)
@ + (e)
(29)
de
de
= =2 + (Oatd)aa))
((u (u=- t)ich - h + (u 1)..h
(30)
12

---

Published as a conference paper at ICLR 2026
T thh)
(u- () + (f + eh)
(u- 1)ou(f +1
(31)
(u t)0rh - h (u- 1)
(32)
So
d de For the second term,
= -2((u-t)(0) + +zff()) - + t,2) _ ((u-)) ( - n) (33)
Tilf + chil c=0
Tzlf + chi = 110 - =)]
(34)
and
Tzlf + chj = 200uff(ts - v)
d
@afe(Es wi]
(35)
de
de
= =2(0nsc (itt - w)
d
(O,Se)(br - v)
(36)
de
= =2[(1 [ + (u - 1)0,( + eli) (it, - -0] (u (m( (1 -w) v)
(37)
So
d de
T2l chil ==0
=
(u - 1) (ite (v w)]
(u ()) w)
(38)
So combining
d de c=0
-
A
= =2 22(1u- ((lu-t](p +)(O.f + 0, fi:) - f f+2)) (lu a
(39)
ws] -2 ( (u - +)0,0) (Es - (t)x
(u-t)0,h(te -v)
(40)
At ( = u, this simplifies
Alfa] 1 = u] de
= -2 F)
(41)
c=0
which is the first variation for regression that makes f equal to |iit] Summarizing,
-t) + (u (10.f++0fft.) -f f + 21) ) (uu-))
() ((u()(()) Oghite) - 1))
sCJf h] =
(42)
- -2 + (w - (1)-)))))( (ite - T)][(v -1)a, - w)]
(43)
13

---

Published as a conference paper at ICLR 2026
A.4 LEMMA: VELOCITY MATCHES A STATIONARY POINT OF ORIGINAL FUNCTIONAL
Lemma 1. Let f be a stationary point of C. Assume that f is bounded. Assume that fo, E C in arguments (t, x) and that all expectations of terms featured in the integrand of C (i.e., v. S.O.f. O.f, i.S....) are finite. Then we have that f (1,1,))== =((,,) where the velocity x(),
Proof. We proceed by contradiction. By the premise, we are at a stationary point fr. Let f == I (u (u-t)f-, f) By the definition of stationary point in section A.1, we have that 8C[f"; h] = 0 for all admissible h. Suppose for the sake of contradiction that at this stationary point, the velocity does
not match, meaning
(t, -.f((t)tt))(((t())))))))
ct,, ))
(44)
suppose for contradiction
The proof proceeds by picking a direction for which the first variation is nonzero, providing a contradiction to being at stationary point. The contradiction (the direction for which the first variation is nonzero) is constructed to arise from assuming that the velocity does not match, meaning that by contradiction the velocity does not match. Specific care is taken ensure that this direction admissible, in this case meaning it is a continuous function.
We name a sequence of functions 9, such there exists 1)" such that 9y is continuous but yields the nonzero variation when chosen as a direction. To establish this is existence under continuity, the dominated convergence theorem is used.
Let us define g(t, u, .) 1[t ul (f ut, v(t, and evaluate it at = so that g t. (1,,,, - v((, We then define the soft indicator I,,( (4, u) that goes to [[ = u] as v - 0 and define it as:
1 (t,,,) u) 1[m 0/2(1 = + exp(- ,2 u)2) + 1[m 1]t m]
(45)
Using the soft indicator, we define a sequence of functions 9,, so that as 17 + 0 we will have pointwise convergence of (t, t, x) + g(t, t, x)) which also means that 9r, (t, u, x) for * goes to O. We pick
(t, u, x) L,, u) (F"(t) (t,,, v(t,
(46)
Pointwise convergence of g in cta. We first establish pointwise convergence of In to for all arguments (t, u, x) as 77 + from the right.
Vn z 0,
lim (+)++
(t, u, x) = go(t, u,
(47)
To show this, for any n > 0, use continuity of ga (f, x) in á (product of function without times the soft indicator which continuous). Then to establish for v = 0. we consider two cases = and # u. For equality:
lim 9n (1, x) --
(48)
lim 7-0+
Ln t) -((t)))
(49)
[I[ > 0)2(1
1+ ( + 1/77 ui] (f(t),t)) - o(t,)))
lim
(50)
--+
[In > 0j1 + 1/7 (f(,,,) -((x))))
lim
(51)
+-+
lim
1/7 2 0 (t,,, t, (F
- =(t
(52)
+-+
-
(53)
14

---

Published as a conference paper at ICLR 2026
which equals I,= z), establishing continuity. Now for t + Vō 0 we must name n(6) such that 9n(5) such that (gn((8 - gol < ô i.c. (9n(6) 0| < 6. Assume If vl < k uniformly in all input values t.x.
lim Ir (l, u, x)) 7-00+
1/7 > 0/2(1
1
+ -11 = u]] (f((),))) - o(l,2))
= -lim 7-0+
1 1+ expp 7 (t u)²)
lim 7-0
1/7 > 012(1
1 exp(- # (t - u)2) (tt,,,,)))
1
Since are finding a 5 and n66) that (91(8) 0| 8 which means [9 < 8, this just means we can set 6 to an upper bound on the term we are limiting: let the indicator take on as when it is 0 are done.
8 = 2(11 1+exp(...)
k
(54)
8
1
=2(1
(55)
k
1 cxxxp(
1/2 (t u)))
8
1
=1 1
(56)
2k
exp(
1
(t - 2)2)
7/
1
1
(57)
2k
1 + exp(
7/ (1 u)
2k 8
(58)
2k
2h
1 + exp(
7 (1 - 2)
2k- - 8
1
(59)
2k
1 1+ cxpl
(t 11-w) x)2)
17t
2k
1
+ exp(
(t - 2)²)
(60)
2k- - 8
172
2k
1
1 exp(
(1 n)
(61)
2k-
1/2
2k
2k 8
1
= exp(
(1 - u)
(62)
2k- 8
2k 8
172
8
1
exp(
(t n)3)
(63)
2
1
1
log
(t - u)²
(64)
2k
5
1/2
8
log 2k 8 1/2
(t u)²
(65)
log 2
1
(66)
(t u)
772
(t u)²
(67)
log 2k-6
(68)
Now note that the soft indicator is strictly < 1 and that 9, for fixed (t, u, x) is between kk and 0 or 0 and & depending on the sign of f* v, but never both. So its magnitude at most k. So
19,(5) gol ( 0) < 8 < k
(69)
This can help us ascertain that the above square root to solve for 17 will be well defined:
dck
2k 5 k
(70)
15

---

Published as a conference paper at ICLR 2026
(71)
2k 8
8
(72)
2k:
8
<1
(73)
2k
log 2k- log
(74)
log 2k- 8
0
(75)
meaning
(I - u²
(6))
(76)
log 2k-6
thus establishing convergence of 9n, + g as v) + 0+ for each (1, u, x) i.e. pointwise convergence.
Now recall the fist variation of L (section A.3) and consider it as a function of yy:
s(n) = " g.
E 2 -))a,-))a-))))) f- + (T (u -- (0)-)))00211111 (1-1)))0=))))
(77)
- t)) - In + -u (00gg)))
(78)
2 [1 (u (u--t) 1)0 (tr - w)] (1 10) -v) wn)
(79)
Pointwise convergence of integrand in eta. Collect the variables
(t, u, To, X1) and recall that 2, and it are functions of (zo, x1). Define (n(w)) as shorthand for the expectand so that s(n) = [w], Under the boundedness assumptions and noting that Φ, only polynomially combines 9, with (0.f*, a.fr, 0,f f*, u, o,.,...,, similar reasoning used to show + g can also be used to establish that 6, + pointwise.
Establish upper envelope. In addition to pointwise convergence of Φn - as n + 0 from the right, need an upper envelope G(w)). Beyond the assumptions, the only thing needed show that an upper envelope exists is to control the term (u I, u)]. The derivative of the soft indicator appears since 0.g, 0, (4,g9) (0.I,)g + I,,0gg At 0 have (, u) 1[t u].
so (u 1) (t, u) vanishes identically: it is for + u because In is constant, and for t = because of the t) prefactor. For 17 > define:
(t u²²
exp(-2)
((2)
d((2)
(80)
1
1+ exex)(22)
[ + exp( --)
Note that for > 0,
(,,,(1, 2 2(( u)
(t u)² u2)
o'(
(81)
172
and so
((u - (t, u)l = 2(u - 2(u t)
u)2 0 (t (t=u)²
= 420 S 420 (2) $ Co < oo (82)
n72
172
=20
Because r(2) 420 (2) is continuous and satisfies r0) 0 and r(z) 0 as co, it attains
a finite maximum. This bound is independent of 77. Thus every term containing (u - t)0 is uniformly bounded in v by product of constant (from the bound and L, E [0, 1)). The other quantities in & are bounded by assumption. Thus such bounding envelope G(w) exists.
Using dominated convergence. First,
lim s(7) lim [[0,] = lim p(t = w)El I u] + plt < u)EE +++
(83)
--00
9-0+
16

---

Published as a conference paper at ICLR 2026
By the pointwise convergence of 6, - 6 as 17 - 0+ and by the envelope, we can compute the limit of the first and second terms separately. Expanding the first term (with p = u):
lim p(t = u) tt= u] = lim p(t = u)BB 2
+ (1+
n n-00+
17 --0
(u- t)0rgn In + (u - ())
- 2 2[( + (u - -t)0 ((br
(u- 1t) 0) (tb -2)) ( w/]
2( - for + a2)) ( (- (-9n) 1t lt= u
=
lim p(t = u)lB 4-
= lim p(t = w)BE 2 jr + Elús mil) ( (-9n) |e
--0++
= lim p(t = u)B 2 - jr + u) )( ( (9gn) 11 = 1
9-01
lim 2 j + u) [ (m)) | w
= p(1 = =u)E
7-0
= p(t = u T)
lim w-
2|| jr + + t t= u
This term is greater than zero by assumption that the velocity does not match at the stationary point and assumption on the positive probabiity on p(t u) 0.
Expanding the second term (with p(t < u)):
lim p(t < (), It u] 2-00
lim p(t u)B 2 (u ()) f (I+ (u (u) 4-0+
(u- 1)0rgn In + (u- - 1)
2 2[(1 ( + (u - -t)aaa (0x, - w)] (u (1 o)) <u
p(t u)EE
lim 2 (u- - f ( 7-000
(u t)0tgn In + (u - ()(
2 2 (u (t)oa.. (tx -
(u (u- [
u
There's no rj in the first term in each of the two dot products that make up the expectand, so we can focus on the second term in the dot products, where t < u For the second term in the first dot product:
lim ((u- t)0.gn - In + (u- (1) -
lim ((u t)0..n. In(t, (F1,,))((,),)))) +(u - (),,, (1,(t,, (F
o(t,,)))11
n-00
lim
(u - t))0ggn
+++
lim (u - ()0, u) o(t.x))
lim (u- - ()(())))) + 1,0gg
-
17

---

Published as a conference paper at ICLR 2026
- lim (u - ()(0,1,) u) -f x(), 1-00
(f v(t, x)) x)) t) lim (t, u) =0
The last equality holds because the function and its time derivative both go zero.
This means the first product in the expecation goes to zero. By a similar argument the second term goes zero as well.
Putting it all together
= lim s(n) = lim E[ø,] lim p(t = L t t= u| + p(t < u) t < u] > 71+
4
1-
Resultingly,
Eno > Ost. Vnn" s.t. 0 <n < |s(n") - L
(84)
If we pick e = 0.57. then
|s(n") L| < 5L =
5L 55 >
(85)
so s(1'") 0. But this contradicts being a stationary point. cannot be that (t, + v(t.-) Therefore the velocity must match.
18

---

Published as a conference paper at ICLR 2026
A.5 THEOREM
We present proof about the functional stationary points of the GameFlow functional. We use the definitions of first variation and stationary point from section A.I and the definition of functional stopgrad from section A.2.
Theorem. Let qt,, u) be a joint distribution with support over S. u and with positive probability on t = u. Let the family F include functions f that are continuously differentiable in all arguments. Let I, = a,70 + 0 and i, = 0,20 + T1 Evaluate f at f(t, x) + (u t) z). Take expectations over qt, u) and ((10)q(Z1) Let sg stand for nop-ggaddent.. Define:
ci [1 (kc..))))))) + (0xf) (1.a,,
(* E. ((.aa)) (02 (1m.,))) (1(0=) ((.0a)) + (el(deflaanaal
Then fr is stationary point of t w.r.t F and only if f is stationary point of C w.r.t. F.
Proof. Case 1: If f" is stationary point of C, then f is stationary point of L
Since fr is stationary point, 8Cff"
by section A.4, we have that (t, = (1,,1, = -E[ir[01] where f = I
(u ft
L L
Since =
= C. then [c =8[ff =
Case 2: If f" is not a stationary point L, then f" not stationary point L
Since f is not stationary point of L. then 3h that is admissible (continuous) such that oL[ h]] # 0. Then,
SC h = B 1 u].
1st + u] u]...
(86)
LHS
RHS-L
RHS-R
If the LHS is nonzero, then one of RHS-L or RHS-R is nonzero. Consider both cases.
case The RHS-R is nonzero and RHS-L zero. RHS-L being zero means that the velocity matches, which means that RHS-R R has the same first variation between L and L So they must coincide regarding stationary points.
case 2b: RHS-I is nonzero and RHS-R is either zero nonzero.
case2b-i 8 [f*; h] + directly holds. This is all we are trying to ensure anyway, so we are done in this case.
case2b-ii Define the soft indicator, In:
(t, u) = 1nn > 0j((1 1 1 (1 w)1)) 1/n = 0]1 = w).
(87)
Define the direction, hy:
(,, u, x)) [t, u) u)h(t, z).
(88)
h, is continuously differentiable for any 17 > O. This is true cause it's a product of two functions that are each continuously differentiable (h is assumed continuously differentiable). Recall the mapping s(n) from Section A.4, that maps 7/ to 8C[;; foz))- Under the conditions of dominated convergence established in Section A.4, we know that 3n > 0 such that hu is a continuously differentiable function for which SCIf*; h. 0. This must mean that the velocity is not matched. But we know that if the velocity does not match, we are not at a stationary point of La either, since Lag and coincide penalizing velocity matching on f u.
19

---

Published as a conference paper at ICLR 2026
B OTHER USEFUL DERIVATIONS AND RESULTS
B.1 GRADIENT UPDATES ARE NOT OPTIMIZATION OF ONE SCALAR OBJECTIVE VIA GRADIENTS
We show here that there exists data distribution and model such that the Gameflow updates are not the gradient of any single scalar objective. We illustrate this by considering an simple ID setting. The key point is that, even in this restricted case, the update field induced by the stopgrad operator has zero curl and therefore cannot be written as the gradient of any scalar objective (0). The form to be differentiated is:
(g) Ex, | u, Xx)+ (O, X,)) [1
(89)
Ex [|(0=fo(t, (X, stopgrad (ft,, (x))))
for ft, u) + (u ftt u, 2). Let us work in ID and fix values of X; and X, d, constant, To accomplish this, we can choose Xi freely and set Xo = Xi - d, so that the interpolation satisfies both X, = x and X, = d. In this case the expectations in (89) collapse to evaluation at this point (equivalently, think of us approximating with Monte Carlo sample). Now, consider the parameters 0 = (01, 02) and single time pair (t, u) such that at the position X; I.
01, x, O, @gJe(t, x)) = 01: f0(1,4,x)) 02.
(90)
(For a sufficiently expressive model, such local values can realized; we only need existence of such a configuration.). Plugging these into (89) and dropping the expectation (single point), the Gameflow functional is:
Lug (0,d)² (01 (d stopgrad(02)))
(91)
Let V denote differentiation with the stopgrad applied to 0g in the second term. Differentiating (91) with respect to 02 yields
Vg, oo Lug 20,d² 2 20, 02)²
(92)
20, ( (d -02)²
(93)
20 d8 (dl 2d00 + 02)
(94)
20 (2d02 0%)
(95)
20, (2d 02).
(96)
Here the stopgrad on O2 only affects the second term and does not change the first term. For 02, all occurrences appear inside a stopgrad and the first term does not depend on 02, hence
VorLug O.
(97)
Thus the update field induced by SGFlow in this simple example is g(0,. 02) (g) (01, 02), (01, 02)) (20: 22,00((2 02), 0).
(98)
this update were the gradient of some sealar objective [(1, 02), then the mixed partial derivatives would commute:
91 = 0o J.
92 = =Jo.J = 0o.91 = -b,92-
(99)
However,
0g1 (20,0((2d 02))
(2dd-))))) -20, (2d 82 -2))) 40,(d (d 82),
(100)
002
002
0g2 0,
(101)
00
and hence the curl of the update field is
0gi
0gz
40, 0z),
(102)
00,
00,
20

---

Published as a conference paper at ICLR 2026
which is non-zero for generic 9 (for example, whenever 81 and 02 + d). Therefore is smooth vector field with non-zero curl and cannor be written as the gradient of any scalar objective J(01, 82) in general. Put differently, constraint the relationship between the parameters of the model and the value of the chosen datapoint is necessary to ensure zero curl.
This ID example is specific instantiation of Gameflow (89) with simple model and single training point. It shows that, once introduce the stopgrad on fo (f, z). the resulting optimization dynamics are in general non-conservative. The stopgrad structure breaks the symmetry required for the updates be the gradient of single scalar function. In this sense, SGFlow is formally (two-player) game rather than standard gradient descent on one potential function. Or one prefers, nongradient vector flow.
B.2 DIFFERENTIATING W.R.T. T VERSUS
Picking one of u just switches the sign of some terms in the loss, which ultimately only affects whether the learned map goes upward in time or downwards. The PDE can be obtained by differentiating the flow map ft,, u, x)) with respect to either endpoint. If X, solves X. = v.), then keeping fixed and differentiating t gives condition that evaluates at and
ft u. x) + (f, u, x)) xtt, 0,
(103)
while keeping t fixed and differentiating in gives a condition that evaluates at u and
ft, a.)::
(tf u, x) + (O,f(t, z)) u, f(t, u. 0.
(104)
Along the ODE trajectory dX- ds one has the identity
O x) 0. u, v,
so choosing to differentiate in versus u only changes the sign convention and which endpoint is held fixed.
21


Reviews:
[{"review_id": "sRg6CgVJ2w", "reviewer": "Reviewer_PeHa", "summary": "This paper proposes GameFlow, an objective for training flow map models for fast sampling. The method simultaneously trains the model to predict both the flow map and its underlying velocity, using a self-correcting stop-gradient target. This approach avoids costly model inversions and large Jacobians. The method is supported by a theoretical proof of correctness and empirical results on cifar-10.", "questions": "1. Can you provide a theoretical comparison between GameFlow and Consistency Trajectory Models[1]?\n2. What specifically causes the 3× memory overhead versus MeanFlow when both methods use JVPs?\n\n[1] Kim, Dongjun, et al. \"Consistency trajectory models: Learning probability flow ode trajectory of diffusion.\" arXiv preprint arXiv:2310.02279 (2023).", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. Theorem 1 rigorously proves that the true flow map is a stationary point of the GameFlow objective. The comprehensive proofs are technically solid.\n2. Provides reasoning and numerical evidence (Figure 1) that MeanFlow does not preserve the true flow map as an optimum.\n3. The paper is well-written and clearly motivated. It systematically identifies and addresses key challenges in the field.", "weaknesses": "1. Only cifar-10 is evaluated with no higher-resolution datasets. According to line 308, the cifar-10 experiments are conditional, yet MeanFlow at 100 steps achieves FID 4.58, which is worse than the 2.92 reported in Table 3 of the MeanFlow paper[1] (Unconditional cifar-10, NFE=1). This discrepancy needs clarification.\n2. Missing comparisons with consistency model variants (Consistency Models[2], simplified Consistency Models[3], etc.).\n3. Missing training curves (loss, FID vs. steps) to verify training stability and convergence behavior.\n4. What is the empirical impact of the corrective term in the loss? An ablation study is needed.\n\n[1] Geng, Zhengyang, et al. \"Mean flows for one-step generative modeling.\" arXiv preprint arXiv:2505.13447 (2025).\n\n[2] Song, Yang, et al. \"Consistency models.\" (2023).\n\n[3] Lu, Cheng, and Yang Song. \"Simplifying, stabilizing and scaling continuous-time consistency models.\" arXiv preprint arXiv:2410.11081 (2024).", "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": "3VYwpB2P3M", "reviewer": "Reviewer_mKce", "summary": "This paper proposes a method, GameFlow, to learn the flow function of the PF-ODE. Unlike other flow learning alternatives, GameFlow's optimum is proven to be the true flow map, and GameFlow does not necessitate computationally expensive or restrictive operations like invertible neural networks or nested differentiation through JVPs. To this end, GameFlow leverages an objective involving only the reverse flow and the true velocity field, the latter being equivalently approximated by the stop gradient of the optimized network itself. Experiments on toy data and CIFAR-10 against other flow models illustrate the soundness and benefits of GameFlow in terms of generative performance and memory footprint.", "questions": "The contributions of this paper appear to be valuable for the community but are hindered by presentation issues and a questionable experimental setting. Therefore, I do not recommend acceptance in the current state of the submission. Still, depending on the authors' response to the above weaknesses, I am willing to change my assessment.\n\nMost importantly, I suggest the authors to address the following points.\n1. Remove from the method name and framing as well as the paper title the notion of game.\n2. Provide further intuition on the rationale behind Theorem 1.\n3. Explain the experimental difference with other papers.\n4. Clarify the advantage of the method compared to consistency models.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. GameFlow is well motivated to address the shortcomings of other flow models. Although this is not new, the paper highlights and addresses the theoretical issues of MeanFlow. Further, GameFlow alleviates the computational and design burden of otherwise sound alternatives. To my knowledge, GameFlow relies on novel derivations in the field. For these reasons, I believe that GameFlow has the potential to influence future works in this research area.\n2. The proposed method seems, up to my limited assessment of the appendix, sound. Empirical results do confirm its advantages within the limits of the chosen experimental setting (cf. weaknesses).\n3. The paper is overall clearly written, except for a few issues described below.", "weaknesses": "### Questions on main claims\n\n1. The interpretation of GameFlow as a game is far-fetched. It solely relies on the presence of a stop gradient in the objective, making the second player trivial. Besides, while this does change the optimization landscape, this is the case of many other methods in or outside the domain (such as many the papers cited in the submission).\n2. Better intuition and presentation for the main result (Theorem 1) would be appreciated. Without reviewing the appendix in detail, it is challenging the understand the rationale behind the result -- although I understand this is a challenging task. My main question regards the fact that repeated arguments (l. 215 in the main paper, l. 1375 and 1406 in the appendix) rely on the loss values to conclude on the loss gradients, which is usually inconclusive when a stop gradient operator is involved.\n\n### Questionable experimental setting\n\n3. GameFlow is tested on limited benchmarks. While this is understandable depending on compute constraints and not eliminatory, it would have beneficiated from a more extensive empirical study involving standard datasets like ImageNet.\n4. More importantly, the FID values on CIFAR-10 are unusually high (compared to e.g. MeanFlow and other papers in the literature). This raises doubts on the realism of the experimental setting, and questions the assertion that the training configuration is \"common\" (l. 320).\n5. Reproducibility is limited as only part of the codebase is available within the submission PDF itself, and the paper shares few details. The promise to release the source code is duly noted.\n\n### Advantage w.r.t. consistency models\n\n6. The paper mentions consistency models as related methods but ignores them when discussing the method's advantages. To my understanding and based on Table 1, consistency models are sound and scalable; the only difference is that they do not model the entire flow (i.e. they are not multistep as stated in the paper). The implications of this difference in practice are unclear in the paper, especially as consistency models are left out of the experimental section.\n\n### Other issues\n\n7. Two minor claims need to be further clarified.\n    - The benefits of using split batch JVP are not explained or shown experimentally.\n    - The authors state the that gradients of stop-grad' loss \"are not the gradients of any one mathematical objective\", but there is no proof for this. It may be that these gradients do correspond to a mathematical objective.\n8. There are a number of formatting issues and typos (some of them spotted below). I would advise the authors to proofread their manuscript.\n    - Table 1 would be more readable with symbols such as ✓ and ✗.\n    - Parentheses are missing for the reference l. 75.\n    - \"covvariance\" should be \"covariance\" (l. 176).\n    - \"as\" should be \"is\" (l. 195).\n    - The second $\\hat{f}$ should be $\\tilde{f}$ l. 215, following l. 185.\n    - Figure 1 should be integrated using vector graphics.\n    - There is an extra period l. 303 and an extra space l. 345.", "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": "iKMMelHHXn", "reviewer": "Reviewer_oX9E", "summary": "The paper derives a new loss for training Flow Maps models, i.e. models that can approximate ODEs in 1 step. They derive their objective function from first principles and prove that it shares same minimizers than the ideal loss. Then, they apply their loss to toy data (multivariate Gaussian) and image generation benchmarks. In toy data setting, the demonstrate that Mean Flows objective can converge to incorrect optimum.", "questions": "* Going from Equation (5) to Equation (6) is done by differentiating with regards to t. Why differentiating by t and not by u, as done in Mean Flows? Could you compare both approaches? \n\n* You propose replacing the velocity field in (6) by sg(f). Actually you could also rely on a pre-trained neural network to approximate the velocity field. This could give another option to train your model in a distillation mode. Have you considered this?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Principled derivation from first principles of a novel training objective for learning Flow Maps. This is an important results, given the attention received by Flow Maps by the research community, and given their importance to reduce inference cost of diffusion models. \n\n* Proved that the tractable training objective has same optimum than the target objective. \n\n* Interesting analysis on the Multivariate Gaussian setting, showing that Mean Flows does not converge to the correct optimum.", "weaknesses": "* Experimental results seem weak and far from results in the literature. How can the authors explain this?  Indeed, 1-step models are now around 2/3 of FID on CIFAR10 (see for example \"Consistency Models Made Easy\" from Geng et al.). Comparing to a FID of 29 reached on 10-step sampling, this is very low. \n\n* Weight decay is rarely used in diffusion/consistency models. Why do the authors choose to use it? Can it explain lower FID than in the rest of the literature?", "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": "J6k4bMtXDM", "reviewer": "Reviewer_NL1h", "summary": "The paper proposes GameFlow, an approach for learning flow maps that bypasses explicit invertibility constraints\nand expensive differentiation through model iteration. It introduces a new training objective of flow-based models with theoretical guarantee of optimality in certain sense.\nNumerical experiments are conducted to verify the results.", "questions": "Please see weakness part.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper is theoretically grounded and proves that the proposed loss (7) is consistent with the true loss function in the sense that they share same functional stationary points. This is a non-trivial improvement upon MeanFlow. And the theories are also verified by experiments, where the solution to MeanFlow is not the ground truth in contrast to GameFlow. In addition, GameFlow outperforms MeanFlow and other methods on real-data experiments like CIFAR-10 generation.", "weaknesses": "1. Although the training objective is theoretically sound, it involves much higher computational cost than MeanFlow due to the backpropagation of JVP, especially when scaling to large models. This is the most severe issue of the proposed method and is inevitable in order to make the objective mathematically correct. Could the authors conduct larger scale experiments if possible to see the scaling ability of the proposed method?\n\n2. I don't see any connection between the proposed framework with game and the interpretation as two-player game is quite far-fetched. The authors may need to choose a more approriate title.", "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)