Flow map matching with stochastic interpolants: A mathematical framework for consistency models
Authors: Nicholas Matthew Boffi
Michael Samuel Albergo
Eric Vanden-Eijnden
Venue: Accepted by TMLR ()

Abstract: Generative models based on dynamical equations such as flows and diffusions offer exceptional sample quality, but require computationally expensive numerical integration during inference. The advent of consistency models has enabled efficient one-step or few-step generation, yet despite their practical success, a systematic understanding of their design has been hindered by the lack of a comprehensive theoretical framework. Here we introduce Flow Map Matching (FMM), a principled framework for learning the two-time flow map of an underlying dynamical generative model, thereby providing this missing mathematical foundation. Leveraging stochastic interpolants, we propose training objectives both for distillation from a pre-trained velocity field and for direct training of a flow map over an interpolant or a forward diffusion process. Theoretically, we show that FMM unifies and extends a broad class of existing approaches for fast sampling, including consistency models, consistency trajectory models, and progressive distillation. Experiments on CIFAR-10 and ImageNet-32 highlight that our approach can achieve sample quality comparable to flow matching while reducing generation time by a factor of 10-20.

Published in Transactions on Machine Learning Research (05/2025)
Flow map matching with stochastic interpolants: A mathematical framework for consistency models
Nicholas M. Boffi*
boffi@cimm..yn.ccmmmm
Courant Institute of Mathematical Sciences
Michael S. Albergo
albergo@myy...mm edu
Courant Institute of Mathematical Sciences
Eric Vanden-Eijnden
eve2@cims.nyu.coo
Courant Institute of Mathematical Sciences
Reviewed on OpenReview: https: // openreview. formm? id= cqDH0e6ak2
Abstract
Generative models based on dynamical equations such flows and diffusions offer exceptional sample quality, but require computationally expensive numerical integration during inference. The advent of consistency models has enabled efficient one-step or few-step generation, yet despite their practical success, systematic understanding of their design has been hindered by the lack of a comprehensive theoretical framework. Here we introduce Flow Map Matching (FMM), principled framework for learning the two-time flow map of underlying dynamical generative model, thereby providing this missing mathematical foundation. Leveraging stochastic interpolants, we propose training objectives both for distillation from a pre-trained velocity field and for direct training of a flow map over an interpolant or a forward diffusion process. Theoretically, we show that FMM unifies and extends a broad class existing approaches for fast sampling, including consistency models, consistency trajectory models, and progressive distillation. Experiments on CIFAR-10 and ImageNet-32 highlight that approach achieve sample quality comparable to flow matching while reducing generation time by factor of 10-20.
1 Introduction
In recent years, diffusion models (Song et al., 2020; Ho et al., 2020; Sohl-Dickstein et al., 2015; Song and Ermon, 2020a;b) have achieved state of the performance across diverse modalities, inchuding image (Dhariwal and Nichol, 2021; Rombach et al., 2022; Esser et al., 2024), audio (Popov al., 2021; Jeong et. al., 2021; Huang et al., 2022; Lu et al., 2022a), and video (Ho et al., 2022a;b; Blattmann et al., 2023; Wu et al., 2023). These models belong to a broader class of approaches including flow matching (Lipman et al., 2022), rectified flow (Liu et al., 2022a), and stochastic interpolants (Albergo and Canden-Eijnden, 2022; Albergo et al., 2023a), which construct a path in the space measures between a base and a target distribution by specifying an explicit mapping between samples from each. This construction yields a dynamical transport equation governing the evolution of the timo-dependent probability measure along the path. The generative modcling problem then reduces to learning velocity field that accomplishes this transport, which provides efficient. and stable training paradigm.
At sample generation time, models in this class generate data by smoothly transforming samples from the base into samples from the target through numerical integration of a differential equation. While effective, the number of integration steps required to produce high-quality samples incurs a cost that can limit real-time applications (Chi et al., 2024). Comparatively, one-step models such as generative adversarial
* Authors contributed equally; author ordering determined randomly.

---

Published in Transactions on Machine Learning Research (05/2025)
Xo.
X No.s
X,,1 Nx.r o Xo,s
L
X o xxr o Xo.s
XO.1
Figure 1: Overview of Flow Map Matching. Our approach learns the two-time flow map X.,e that transports the solution of an ordinary differential equation from time s to time t. Unlike methods that learn instantaneous velocity fields, this bidirectional map can be used to build an integrator with arbitrary discretization. The integrator is exact in theory, and -crucially- its number steps can be adjusted post-training to balance accuracy and computational efficiency. The flow map can be distilled from a known velocity field or learned directly, and supports arbitrary base distributions, as illustrated here with amage-lo-imagee translation.
networks (Goodfellow et al., 2014; 2020; Creswell et al., 2018) are notoriously difficult to train (Metz et al., 2017; Arjovsky et al., 2017), but can be orders of magnitude more efficient to sample, because they only require single network evaluation. As a result, there has been significant recent research effort focused on maintaining the stable training of diffusions while reducing the computational burden of inference (Karras et al., 2022). Towards this goal, we introduce Flow Map Matching a theoretical framework for learning the two-time flow map probability flow. Specifically, our primary contributions can be summarized as:
A mathematical framework. We characterize the flow map theoretically by identifying its fundamental mathematical properties, which we show immediately lead to distillation-based algorithms for estimating the flow map from a pre-trained velocity field. Using our characterization, we establish connections between our framework and recent approaches for learning consistency models (Song et al., 2023; Song and Dhariwal, 2023; Kim et al., 2024a) and for progressive distillation of a diffusion model (Salimans and Ho, 2022; Zheng et al., 2023).
Lagrangian vs Eulerian learning. We develop a novel Lagrangian objective for distilling a flow map from a pre-trained velocity field, which performs well across our experiments. This loss is complementary to a related Fulerian loss, which we prove the continuous- limit of consistency distillation (Song et al., 2023). We show that both the Lagrangian and Fulerian losses control the Wasserstein distance between the teacher and the student models.
Dircct training. We extend our Lagrangian loss to design novel direct training objective for flow maps that eliminates the need for a pre-trained velocity field. We also discuss how to modify the Eulerian loss to allow for direct training of the map, and identify some caveats with this procedure. Our approach elucidates how the standard training signal for flow matching can be adapted to capture the more global signal needed to train the flow map.
Numerical experiments. We validate our theoretical framework through experiments on CIFAR-10 and ImageNet 32 x 32, where we achieve high sample quality at significantly reduced computational cost. particular, we highlight that learning the two-time flow map enables post-training adjustment of sampling steps, allowing practitioners to systematically trade accuracy for computational efficiency.
2 Related Work
Dynamical transport of measure. Our approach is built upon the modern perspective of generative modeling based dynamical transport of measure. Grounded in the theory of optimal transport (Villani,
2

---

Published in Transactions on Machine Learning Research (05/2025)
2009; Benamou and Brenier, 2000; Santambrogio, 2015), these models originate at least with (Tabak and Banden--iinden, 2010; Tabak and Turner, 2013); but have been further developed by the machine learning community in recent years (Rezende and Mohamed, 2015; Dinh et al., 2017; Grathwohl et al., 2018; Chen et al., 2019). A breakthrough in this area originated with the appearance of score-based diffusion models (Song et al., 2020; Song and Ermon, 2020a;b), along with related denoising diffusion probabilistic models (Ho el al., 2020; Sohl-Dickstein et al., 2015). These methods generate samples by learning to time-reverse a stochastic differential equation with stationary density given by Gaussian. More recent approaches such as flow matching (Lipman et al., 2022), rectified flow (Liu et al., 2022a;b), and stochastic interpolants (Albergo and Vanden-Fijnden, 2022; Albergo et al., 2023a; Ma et al., 2024; Chen et al., 2024) similarly construct connections between the base density and the target, but allow for arbitrary base densities and provide a greater degree of flexibility in the construction of the connection.
Reducing simulation costs. There has been significant recent interest in reducing the cost associated with solving an ODE or an SDE for a generative model based on dynamical measure transport. One such approach, pioncered by rectification (Liu et al., 2022a;b), is try to straighten the paths of the probability flow, so as to enable more efficient. adaptive integration. In the limit of optimal transport, the paths become straight lines and the integration can be performed in single step. A second approach is to introduce couplings between the base and the target, such as by computing the optimal transport over minibatch (Pooladian et al., 2023; Tong et al., 2023), or by using data-dependent couplings (Albergo et al., 2023b), which can simplify both training and sampling. A third approach has been to design hand-crafted numerical solvers tailored for diffusion models (Karras et al., 2022; Zhang and Chen, 2023; Molicoeur-Martnneat et al., 2021; Liu et al., 2022c; Lu et al., 2022b), or to learn these solvers directly (Watson et al., 2021; 2022; Nichol and Dhariwal, 2021) to maximize efficiency. Instead, we propose to learn the flow map directly, which avoids estimating optimal transport maps and can overcome the inherent limåtations of numerical integration.
Consistency models. Most related to our approach are the classes of one-step models based distillation consistency; we give explicit mapping between these techniques and our own in Appendix C. Consistency models (Song et al., 2023) have been introduced as new class of generative models that can either be distilled from a pre-trained diffusion model or trained directly, and are related to several notions of consistency of the score model that have appeared in the literature (Lai et al., 2023a;b; Shen et al., 2022; Boffi and Vanden-Eijuden, 2023; Daras el al., 2023). These models learn a one-step map from noise to data, and can be seen as learning single-time flow map. While they can perform very well, consistency models do not benefit from multistep sampling, and exhibit training difficulties that mandate delicate hyperparameter tuning (Song and Dhariwal, 2023). By contrast, we learn two-time flow map, which enables us to smoothly benefit from multistep sampling.
Extensions of consistency models. Consistency trajectory models (Kim et al., 2024a) were later introduced to improve multistep sampling and to enable the student to surpass the performance of the teacher. Similar to our approach, these models learn a two-time flow map, but do so using a very different loss that incorporates challenging adversarial training. Generalized consistency trajectory models (Kim et al., 2024b) extend this approach to the stochastic interpolant setting, but use the consistency trajectory loss, and do not introduce the Lagrangian perspective considered here. Bidirectional consistency models (Li and He, 2024) learn two-time invertible flow map similar to our method, but do so in the score-based diffusion setting, and do not leverage our Lagrangian approach; all of these existing consistency models can be seen as special case of our formalism.
Progressive and operator distillation. Neural operator approaches (Zheng et al., 2023) learn a one-time flow map from noise to data, but do so by first generaling a dataset of trajectories from the probability flow. Progressive distillation (Salimans and Ho, 2022) and knowledge distillation (Luhman and Luhman, 2021) techniques aim to convert a diffusion model into an oquivalent model with fewer samples by matching several steps of the original diffusion model. These approaches are related to our flow map distillation scheme, though the object we distill is fundamentally different.
3

---

Published in Transactions on Machine Learning Research (05/2025)
3 Flow Map Matching
The central object in our method is the flow map, which maps points along trajectories of solutions to an ordinary differential equation (ODE). Our focus in this work is primarily on probability flow ODEs that arise in generative models, such as those constructed using flow matching via stochastic interpolants or score-based diffusion models (sce Appendix A for concise, self-contained review of these approaches). While many of our results apply more generally to other ODEs, we present our definitions and theorctical results in this specific context highlight their relevance to generative modeling. All proofs of the stalements made in this section are provided in Appendix B, with some additional theoretical results and connections to existing consistency and distillation techniques given in Appendix C.
3.1 Stochastic interpolants and probability flows
We first recall the main components of the stochastic interpolant framework (Albergo and Vanden- Vanden-Eijnden, 2022; Albergo et al., 2023a), which we use to construct the probability flow ODEs central to our study. We begin by giving the definition of a stochastic interpolant.
Definition 3.1 (Stochastic Interpolant). The stochastic interpolant 4. between probability densities po and is the stochastic process given by
= =0,0 + -,,.11 + 7
(3.1)
where a, 3,222 ((0, satisfy og = 31 1, an Bu Bu-=0 0, and To - 712 0. In (3.1), (2r, v1) is drown from coupling (1), 21) that satisfies the marginal constraints JR p( xn )dro = (1(11 and SRa p(zo, 21 ) (o())). Moreover, N(0, Id) with 1 (zo: æ1).
Theorem 3.6 of Albergo et al. (2023a) shows that the stochastic interpolant given in Definition 3.1 specifies an underlying probability flow ODE, as we now recall.
Proposition 3.2 (Probability Flow). For all t E [0, the probability density of I is the same as the probability density of the solution to
bz
It= = xo ~ po
(3.2)
where b: [0 Rd RRd is the time- dependenndd velocity field (or drift) given by
by (x -[ z²)
(3.3)
In (3.3), B[- [T, = x| denotes an expectation over the coupling x1) ~ p( x1) and z ~ N(O, Id) conditioned on the event I = 2.
Although (3.2) is deterministic for any single trajectory, taken together, its solutions are random because the initial conditions are sampled from Pa- We denote by Pt = Law(x.) the density of these solutions at time t. The drift can be learned efficiently in practice by solving a square loss regression problem (Albergo et al., 2023a)
argmin L (1PP idP]dt
(3.4)
where E denotes an expectation over the coupling (zo, ~ p( x1) and 2 ~ N(O, Id).
A canonical choice when /n = N(O, Id) considered in Albergo and Vanden-Eijnden (2022) corresponds to 0 = 1 1-t, 3t f, and 7t 0, which recovers flow matching (Lipman et al., 2022) and rectified flow (Liu et al., 2022a). The choice a = 0, 3
and 7 = vI - te corresponds to darianco-preserving diffusion model with the time-rescaling t - log where E [0, oo) is the usual diffusion time". A variance-exploding diffusion model may be obtained by taking a, = 0, 3r = 1, and = T t with i E [0, 1] and where = =T - is the usual diffusion time, though this violates the boundary conditions in Definition 3.1. For more details about the connection between stochastic interpolants and diffusion models, we again refer the reader to Appendix A.
Note that 70 = in this case, 80 that To
4

---

Published in Transactions on Machine Learning Research (05/2025)
3.2 Flow map: definition and characterizations
To generate sample from the larget density Pl, we can draw initial point. from Po and numerically integrate the probability flow ODE (3.2) over the interval E [0, 1]. While this approach produces high-quality samples, it typically roquires numerous integration steps, making inference computationally when b, is parameterized by a complex neural network. Here, we bypass this numerical integration by estimating the two-time flow map, which lets us take jumps of arbitrary size. To do so, we require the following regularity assumption, which ensures (3.2) has solutions that exist and are unique (Hartman, 2002).
Assumption 3.3. The drift salisfies the one-sided Lipschilz condition
3 C. E 10 : (be(x) - (v)) (a u) & Calz yl? for all (t,z,y) E [0, x Rd x Rd.
(3.5)
With Assumption 3.3 in hand, we may now define the central object of our study.
Definition 3.4 (Flow Map). The flow map Na : Rd - Rd for (3.2) is the unique map such that
(x,) It for all (s,t) [0, 1]3,
(3.6)
where (1t))E(0,1) is any solution to the ODE (3.2).
The flow map in Definition 3.4 can be seen as an integrator for (3.2) where the step size s may be chosen arbitrarily. In addition to the defining condition (3.6), we now highlight some of its useful properties.
Proposition 3.5. The flow map (x) is the unique solution to the Lagrangian equation
0,X,,((z)= (x))),
Xx,o x.
(3.7)
for all (s, E 10. 112 x Rd, In addition, it satisfies
Xx.(Xx,( = (x)
(3.8)
for all (, E [0, 113 x Rd. In particulor X (X(, x for all (s,t,z) E [0, 10. Rd, iec the flow mop is invertible.
Proposition 3.5 shows that, given an To ~ Po: we can use the flow map to generate samples from Pt for any E [0 in one step via di Xa, ~ Pt.
The composition relation (3.8) the consistency property, here stated over two times, that gives consistency models their name (Song et al., 2023). This relation shows that can also generate samples in multiple steps using It, -X 1.44 (Ita-1) Pt, for any set of discretization points (lo....,4x) with l E (0, 1] and K E N. We refer to (3.7) as the Lagrangian equation because it is defined in a frame of reference that moves with (x).
The flow map X,, also obeys an alternative Fulerian equation that is defined at any fixed point x E R& and which involves a derivative with respect to s. To derive this equation, note that since Xxs(1x)) = Ia by (3.6) we have (d/ds)X., 0, which by the chain rule can be written
d X,, (x.)) = O, (x.,,((2.) + (xx)) (x..( = -0
(3.9)
ds
where we usod (3.2) to set 4,- (x.). Evaluating this equation at x, = x gives the announced result. Proposition 3.6. The flow map X is the unique solution of the Eulerian equation,
(xX,,((x) (x)) (x),(x) =0, =0, Xxx(x)= = x,
(3.10)
for all (s, (s)), [0, 1J2 Rd.
5

---

Published in Transactions on Machine Learning Research (05/2025)
3.3 Distillation of known velocity field
The Lagrangian equation (3.7) in Proposition 3.5 leads to distillation loss that can be used to learn flow map for a probability flow with known right-hand side b, as we now show.
Corollary 3.7 (Lagrangian map distillation.. Let wsst E [[0, 113) be weight function satisfying waat > 0 and let Ig be the stochastic interpolant defined in (3.1). Then the flow map is the global minimizer over X of the loss
LLMD = [ Ws,, B10......
(3.11)
10.11
subject to the boundary condition that (x) for all x E I2 and [0, In (3.11), E denotes on expectation over the coupling (1) p( and N(0, Id).
The time integrals in (3.11) can also be written as an expectation over (s, ~ Wa,t (properly normalized):
LMD (8) Res.t)mw. I[ (B
(3.12)
We also note that the result of Corollary 3.7 remains true if we replace 1e by any process à, with density Po F Psy long it is strictly positive everywhere. In practice, it is convenient to use I, as it guarantees that we learn the flow map where typically need to evaluate it. Moreover, using I, avoids the need to generate a dataset from a pre-trained model, so that the objective (3.11) can be evaluated empirically in a simulation-free fashion.
When applied to a pre-trained flow model, Corollary 3.7 can be used to train a new, few-step generative model with performance that matches the performance of the teacher. When Xs, parameterized by neural network, the partial derivative with respect t can be computed efficiently at the same time as the computation of Xx using forward-mode automatic diffcrentiation. This procedure is summarized in Algorithm! 1.
For simplicity, Corollary 3.7 is stated for Was > 0 (e.g. Was, = 1), so that we can estimate the map X and its inverse Xt,s for all (s, t) E [0, 1]. Nevertheless, this weight can also be adjusted to learn the map for different pairs (s, of interest. For example, if we only want to estimate the forward map with < L, then we can set Wx,t = if E t and Ws,t otherwise.
By squaring the left hand-side of the Eulerian equation (3.10) in Proposition 3.6, may construct a second loss function for distillation.
Corollary 3.8 (Eulerian map distillation). Let wa,t and Ie be as in Corollary 3.7. Then the flow map is the global minimizer over X of the loss
LEMD (8) - =/0 (1,,E[0 + b.l.)) (Xa.(t))P]dsdtt
(3.13)
10.11
subject to the boundary condition (x) = x x for all x E Rd and for all s E R.
As in the discussion after Corollary 3.7, here we may also replace I, by any process &, with a strictly positive density. Writing the time integrals in (3.13) as an expectation over (s,t) wa, gives
( = (e.t) 8[0....... (1.) + (..)) (1.)
(3.14)
We summarize a training procedure based on Corollary 3.8 in Algorithm 2.
In Appendix C, we demonstrate how the preceding results connect with existing distillation based approaches. In particular, when (z) the velocity of the probability flow ODE associated with a diffusion model, Corollary 3.8 recovers the contimuous-time limit of consistency distillation (Song et al., 2023; Song and Dhariwal, 2023) and consistency trajectory models (Kim et al., 2024a), while Corollary 3.7 is new
PIn (3.10), the term (bx(x). vx),((a)), = (b,] (x), = (VX.,(()) - bx)1, corresponds to a Jacobian-vector product that can computed efficiently using forward-mode automatic differentiation.
6

---

Published in Transactions on Machine Learning Research (05/2025)
Algorithm 1: Lagrangian map distillation (LMD)
input: Interpolant coefficients Oit, P6 pre-trained velocity be, weight function Wa,t, batch size M. repeat
Draw batch (si, 14, wh. oi, 2:) from wa,t x p( xi) x N(2; Id).
Compute It, at, di + Ji +rito n, 2 and Lt, = + of + 7.
Compute LLMD = 1
M En ( ) - bi, (X.,, (I,, ))1
Take gradient step LLM to update X. until converged;
output: Flow map X.
Algorithm 2: Eulerian map distillation (EMD)
input: Interpolant coefficients ott, Be, 7t; pre-trained velocity by, weight function Wast batch size M. repeat
Draw batch (st, ti, ag. 2f, 24 = from wa, x p(zo, x1) x N(2; 0, Id).
Compute It, ug + 31,01 + 2 and lt =0 + Be,ai 2-
Compute LEMD 37 E;=1 10, 8 s. (e..:(1s))) (I,.) + V& (Ix)) (Is,)
Take gradient step on LEMD to update X. until converged;
output: Flow map X.
3.4 Wasserstein control
In this section, we show that the Lagrangian and Eulerian distillation losses (3.11) and (3.13) control the Wasserstein distance between the density Pe of the teacher flow model and the density Á. Xo, po of the pushforward of Po under the learned flow map (that is, O is the density of (x,( with zu Pu). When combined with the Wasserstein bound in Albergo and Vanden-Fijnden (2022), the following results also imply bound on the Wasserstein distance between the data density and the pushforward density for the learned flow map in the case where b is pre-trained stochastic interpolant diffusson model. We begin by stating our result for Lagrangian distillation.
Proposition 3.9 (Lagrangian error bound). Let Xa,t Hd + Rd denote the flow map for b, and let X,, IRd + Rd denote an approrimate flow map. Given Do Pu. let Đr be the density of Xo,1 and let p be the larget densily of (zo) Then,
(1 & t22 celds (LD)
(3.15)
where C. is the constant that oppears in Assumption 3.3.
The proof is given in Appendix B. We now state an analogous result for the Eulerian case.
Proposition 3.10 (Eulerian error bound). Let Xe,: :dd + Rd denote the flow map for b, and let Xes :
Rd + Rd denote an appronimate flow map. Given xo Pa, let Pr be the density of &u, and let pr be the largel densily of (2,) Then,
(21,)) < (X).
(3.16)
The proof is also given in Appendix B. The result in Proposition 3.10 appears stronger than the result in Proposition 3.9, because it is independent of any Lipschitz constant. Notice, however, that unlike (3.15) the bound (3.16) involves the spatial gradient of the map Xs,,, which may be more difficult to control. In our numerical experiments, we found the best performance when using the Lagrangian distillation loss, rather than the Fulerian distillation loss. We hypothesize and provide numerical evidence that this originates from avoiding the spatial gradient present in the Eulerian distillation loss; in several cases of interest, the learned map can be singular or nearly singular, that the spatial gradient is not well defined everywhere. This leads
7

---

Published in Transactions on Machine Learning Research (05/2025)
Algorithm 3: Flow map matching (FMM)
input: Interpolant coefficients Oit, Be, 76 weight function Wa,t; batch size M.
repeat
Draw batch (si, 14, uh. of, 21) from wa,t x p(an, xi) x N(2; Id).
Compute 4t, ai, 2ij + Fi, + n, 2 and Le, =o + oi + 7. Compute ČFMM- t Ei (1a, (Xts,s L. 1 18,, (Riuss (fr.)) 1r.1))
Take gradient step on FFMM to update &. until converyed;
output: Flow map X.
to training difficulties that manifest themselves as fuzzy boundaries on the checkerboard dataset and blurry images on image datascts.
3.5 Direct training with flow map matching (FMM)
We now give a loss function for direct training of the flow map that does not require a pre-trained b. Proposition 3.11 (Flow map matching). The flow map the global minimizer over & of the loss
LFMM( - -
10.112 Whst
- + E - -iPl) dsdt,
(3.17)
subject to the boundary condition (z)= = x for all a Rd and for all E R. In (3.17), Wa,t > 0 and is taken over the coupling (1) ~ x1) and N(O, Id).
In the loss (3.17), we are free to adjust the weight factor Wash. as illustrated in Figure 2. However, since we need to learn both the map Xs,t and its inverse Xt,e, it is necessary to enforce the symmetry property Wt,, = Wa.s. If we learn the map for all (s, t) E [o, 1|² using, for example, w,,t = 1, then we can generate samples from via (2o) with Tn ~ Po- If we learn the map in a strip, as shown in Figure 2, the learning becomes simpler but we need to use multiple steps to generate samples from P1; the number of steps then depends of the width of the strip. We note that the second term enforcing invertibility in (3.17) comess at no additional cost on the forward pass, because (2t. can be computed at the same time as (Xt,, (I+)) with standard Jacobian-vector product functionality in modern deep learning packages. A summary of the flow map matching procedure is given Algorithm
t
in one step
X8,28
Xo.5
Wsd >,>
1/K
X
Figure 2: Schematic illustrating the
weight Wa,t in the FMM loss, which can be tuned to arrive at different. learning schemes.
3.6 Eulerian estimation or Eulerian distillation?
In light of Proposition 3.11, the reader may wonder whether we could also perform direct estimation in the Fulerian setup, using for example the objective
0.1 W., IE (L) (.,) dadt.
(3.18)
The above loss may be obtained from (3.13) by replacing the appearance of b.(e)) with i,. Unfortumately, (3.18) is not equivalent to (3.13). To see why, we can expand the expectation in (3.18):
cl (I.) + i, Vt, ((.)
(3.19)
= =E(10 (1.)1 + 2(1. v&aa(la)) 0.8aa(1.) + 1i. (1.)
8

---

Published in Transactions on Machine Learning Research (05/2025)
Algorithm 4: Progressive flow map matching (PFMM)
input: Interpolant coefficients Ott, Be,7 76; weight wa,ti K-step flow map X; batch size M. repeat
= 8, (k- 1)(4+ 84) for A , K. Draw batch (st; 1e and compute t ===0
Compute LPFMM- Xf Σi (1xs,,t(Tn)) (Ts,) - (X1x_11tx o..0 8 (x..,))
Take gradient stop on ΣPFMM to update X.
until converged;
output: One-step flow map X.
The cross term is linear is so that we can use the tower property of the conditional expectation to see [[, V&,, 0,8,,, = 13 (E(i,A)) v.,,((4)) ((,)),
(3.20)
= (b,((( vX.,((I)) (I.)
However, the tower property cannot be applied to the last term in (3.19) since is quadratic in İg, i.e.
E Illi. [.,,, * (E(bo(L))
(3.21)
Since this term depends on 8, it cannot be neglected in the minimization, and the minimizer of (3.18) is not the same that of (3.13). Recognizing this difficulty, consistency models (Song al., 2023; Song and Dhariwal, 2023; Kim et al., 2024a) place a stopgrad(-) on the term I, V&,,((I.) when computing the gradient of the loss (3.18). The resulting iterative scheme used to update X then has the correct fixed point at X X, summarized in the following proposition.
Proposition 3.12 (Eulerian estimation). The flow map X.,t is critical point of the loss
( = S 10.132
Wa, E[0,,,,,,,11, stopgrad(i, (...(( dsdt,
(3.22)
subject to the boundory condition (2,(1T) = x for all x E TRd and for all 8 E TR. In (3.22), Ws,t > and E taken over the coupling (Io, x1) ~ p(ro, xT), ~ N(O, Id). The operator stopgrad(-) indicates that the aryument is ignored when compuling the gradient.
We note that it is challenging guarantee that the critical point described in Proposition 3.12 stable and attractive for iterative gradient-based scheme, the objective function is not guaranteed to decrease due to the appearance of the stopgrad operator. Nonetheless, some works have found success with approaches related to this objective on large-scale image generation problems (Song et al., 2023).
3.7 Progressive distillation
Empirically, found directly learning one-step map to be challenging in practice. Convergence significantly improved by taking Ws,r = Wh, = [[ s| & 1/K) for some K E N where denotes an indicator function. Given such a K-step model, it can be converted into a one-step model using a map distillation loss that is similar to progressive distillation (Salimans and Ho, 2022) and neural operator approaches (Zheng et al., 2023).
Lemma 3.13 (Progressive flow map matching). Let & be a two-time flow map. Given K N, let to s k=) s) s) for k 1,..., K. Then the unique minimizer over X of the objective
(XF) ( /..1p we w. ((..(())))))) - (Six-1 o... & (s) dsdt,
(3.23)
produces the same output in one step as the K- step iterated map X. Here w.tt > 0, and is taken over the coupling (z0, ) ~ 21) and N(O, Id).
9

---

Published in Transactions on Machine Learning Research (05/2025)
We note that X is fixed in (3.23) and serves as the teacher, so only need to compute the gradient with respect to the parameters of X. In practice, we may train X using (3.17) over a class of neural networks and then freeze its parameters. We may then use (3.23) to distill X into a more efficient model X, which can be initialized from the parameters of X for an efficient warm start. Alternatively, we can add (3.23) to (3.17) and train a single network X by setting X = X in (3.23) and placing a stopgrad(-) on the multistep term.
If K evaluations of X are expensive, we may iteratively minimize (3.23) with some number M < K evaluations of X and then replace X by X, similar to progressive distillation (Salimans and Ho, 2022). For example, we may take M 2 and then minimize (3.23) [log2 K] times to obtain a one-step map. Alternatively, we first generate dataset of (s,t, Ig, (Xtx-1.xx 0 (T.)) in parallel offline phase, which converts (3.23) into a simple least-squares problem. Finally, if we are only interested in using the map forward in time, we can set Wx,t = 1 if s < and Wgg = 0 otherwise. The resulting procedure is summarized in Algorithm 4.
4 Numerical Realizations
In this section, we study the efficacy of the four methods introduced in Section 3: the Lagrangian map distillation discussed in Corollary 3.7, the Eulerian map distillation discussed in Corollary 3.8, the direct training approach of Proposition 3.11, and the progressive flow map matching approach of Lemma 3.13. We consider their performance on two-dimensional checkerboard dataset, as well as in the high-dimensional setting of image generation, to highlight differences in their training efficiency and performance.
To ensure that the boundary conditions on the flow map Xs, defined in (3.7) are enforced, in all experiments, we parameterize the map using the ansatz
(x, = + (t (0d-8)02.1 (x),
(4.1)
where (x) [0, 17 x IRd + Id is a neural network with parameters 0.
4.1 2D Illustration
As simple illustration of our method, we consider learning the flow map connecting a two-dimensional Gaussian distribution to the checkerboard distribution presented in Figure 3. Note that this example is challenging because the target density is supported on a compact set, and because the target is discontinuous at the edge of this set. This mapping can be achieved, as discussed in Section 3, in various ways: (a) implicitly, by solving (3.2) with a learned velocity field using stochastic interpolants (or a diffusion model), (b) directly, using the flow map matching objective in (3.17), (c) progressively matching the flow map using (3.23), or (d/e) distilling the map using the Eulerian (3.13) or Lagrangian (3.11) losses. In each case, use fully connected neural network with 512 neurons per hidden layer and 6 layers to parameterize either a velocity field (xr) or a flow map (z), We optimize each loss using the Adam (Kingmaa and Ba, 2017) optimizer for 5 x 10t training iterations. The results are presented in Figure 3, where we observe that using the one-step (x) directly learned by minimizing (3.17) over the entire interval (s, 1) E [0. 1J2 performs worse than learning with It s| 0.25 and sampling with steps. With this in mind, we use the 4-step map as teacher to minimize the PFMM loss, which produces a high-performing one-step map. We also note that the EMD loss performs worse than the LMD loss when distillig the map from learned velocity field.
4.2 Image Generation
Motivated by the above results, we consider series of image generation experiments on the CIFAR-10 and ImageNet-32 x 32 datasets. For comparison, benchmark the method against alternative techniques that seek to lower the number of steps needed to produce samples with stochastic interpolant models, e.g. by straightening the ODE trajectories using minibatch OT (Pooladian et al., 2023; Tong et al., 2023). We train all of our models from scratch, so as to control the design space of the comparison. For clarity, we label when benchmark numbers are quoted from the literature.
For learning of the flow map, we use a U-Net architecture following (Dhariwal and Nichol, 2021). For LMD and EMD that require a pre-trained velocity field to distill, we also use a U-Net architecture for b. Because the
10

---

Published in Transactions on Machine Learning Research (05/2025)
SI, = 80
FMM, Dku: 0.017
FMM, N
PFMM, N
LMD, N
EMD,
Dk 0.016
D 0.008
DRu: 0.01
D 0.033
DEL 0.082
Figure 3: Two-dimensional results. Comparison the various map-matching procedures on the checkerboard dataset, with the results from the probability flow ODE of stochastic interpolant integrated using N 80 discretization steps as reference (top, first panel from left). The one-step map obtained by FMM when learning on (s, t) = [o, 1]2 (top, second panel) performs well as SI. Moreover, the accuracy improves if allow four steps instead of one (top, third panel). This four-step map can be accurately distilled into one-step map via PFMM (top, fourth panel). The one-slep map obtained by distilling the pre-trained via LMD (top, fifth panel) performs reasonably well too, and is better than the one-step map oblained by distilling the same via EMD (lop, sixth panel). A KL-divergence between each model distribution and the target provided quantify performance, indicating that FMM, its progressive distillation, and LMD are closest to the probability flow ODE baseline. The bottom row indicates, by color, how points from Gaussian base are assigned by each land outside the checkerboard. These results indicate that the primary source of error in each case is handling the discontiniity of the optimal map at the edges of the checker. See Appendix for more details.
respective maps. The yellow dots are points that mistakenly
flow map X..: is function of two times, modify the architecture, Both and t are embedded identically to t in the original architecture. The result is concatenated and treated like t in the original architecture for downstream layers. We benchmark the performance of the met using the Frechet Inception Distance (FID), which computes a measure of similarity between real images from the dataset and those generated by our models. In addition, we compute what we denote as the Teacher-FID (T-FID). This motric computes the same measure of similarity, but now between images generated by the teacher model and those generated by
A
N=
LMD, ImageNet-32x33
LMD, CIFARIO
Si
EMD, ImageNet-32x332
EMD. CIFARIO
100
LMD
1-step FTD
EMD
Loss
10a
PFMM
10'
10
10²
100000
200000
FMM
Training Step
Training Step
Figure 4: (A) Qualitative comparison between the standard stochastic interpolant approach (SI). Lagrangian map distillation (LMD), Eulerian map distillation (EMD), and progressive flow map matching (PFMM). SI produces good images for sufficiently large number of steps, but performs poorly for few steps. LMD performs well in the very-few step regime, and outperforms EMD significantly. PFMM performs well at any number of steps, though performs slightly worse than LMD in the very-few step regime. (B) Quantitative comparison between EMD and LMD both CIFAR-10 10 and ImageNet 32 x 32, Despite both having the same minimizer, LMD trains faster, and attains lower loss value and lower FID for fixed number of training steps.
11

---

Published in Transactions on Machine Learning Research (05/2025)
N=2
N
Method
Baseline
FID
T-FID
FID
T-FID
SI
112.42
34.84
5.53
EMD
48.32
34.19
44.35
30.74
5.53
LMD
7.13
1.27
6.04
1.05
5.53
PFMM
18.35
7.02
11.14
1.52
8.44
Table 1: Comparison of various distillation methods using FID and Teacher-FID metrics the CIFAR-10 dataset. Note that for PFMM, no velocity model (e.g. from a stochastic interpolant) is needed. It relies solely on the minimization of (3.17) and (3.23). Baseline indicates the FID of the teacher model (a velocity field for EMD and LMD integrated with an adaptive fifth-order Runge- Mutge scheme, and a flow map for PFMM) against the true data.
the distilled model, rather than leveraging the original dataset. This measure allows us to directly benchmark the convergence of the distillation method, as captures discrepancies between the distribution of samples generated by the teacher and the distribution of samples generated by the student. In addition, this allows benchmark accuracy independent of the overall performance of the teacher, as our teacher models were trained with limited compute.
Sampling efficiency In Table 1, we compute the FID and T-FID for the stochastic interpolant, Eulerian, Lagrangian, and progressive distillation models on 2 and 4-step generation for CIFAR-10. The stochastic interpolant was trained to a baseline FID (sampling with adaptive solver) of 5.53, and used as the teacher for EMD and LMD. The teacher for PFMM was FMM model trained with 5 < 0.25 to an FID of 8.44 using 8-step sampling. We observe that LMD and EMD methods can effectively distill their teachers and obtain low T-FID scores. In addition, the 2 and 4-step samples from these methods far outperform the stochastic interpolant. This sampling efficiency is also apparent in the left side of Figure 4, in which with just 1 to steps, the LMD and PFMM methods can produce effective samples, particularly when compared the flow matching approach.
Without any distillation, FMM can also produce effective few-step maps. Training an FMM model the TeageNet-32 x 32 dataset, we observe (Table 2) that FMM achieves much better few-step FID when compared to denoising diffusion models (DDPM), and better FID than mini-batch OT interpolant methods (Pooladian et al., 2023). In the higher-step regime, the interpolant methods perform marginally better.
DDPM
BatchOT
FMM (Ours)
20
63.08
7.71
9.68
232.97
15.64
12.61
275.28
22.08
14.48
362.37
38.86
16.90
Table 2: FID scaling with number of function evalua- tions to produce sample on ImageNet-32 x 32. We show comparison between DDPM (Ho al., 2020) and multi-sample Flow Matching using the BatchOT method (Pooladian et al., 2023) flow map matching. The first two columns quoted from Pooladian et al. (2023). Note that no distillation used here, but rather direct mininization of (3.17), using |z - s| < 0.25.
Eulerian vs Lagrangian distillation Remark- ably, we find a stark performance gap between the Eulerian and Lagrangian distillation schemes. This is evident in both parts of Figure 4, where we see that higher-step sampling with EMD only marginally improves image quality, and where the LMD loss for both CIFAR10 and ImageNet-322 x 32 converges an order of magnitude faster than the EMD loss. The same holds for FIDs over training, given in the right-most plot in the figure. Note that both LMD and EMD loss functions have a global minimum at 0, so that the loss plots suggest continued training will improve distillation quality, but at very different rates.
4.2.1 Efficient Style Transfer
To illustrate some the downstrcam tasks facilitated by our approach, we now describe mcans of performing style transfer with the tro-time invertible flow map that we call "consistency style transfer". A class conditional sample 21 with class label y can be partially inverted to earlier time s' via. X1,s (x1; v) for s' 1. From
12

---

Published in Transactions on Machine Learning Research (05/2025)
Car
Cat
Car
Bird
Frog
From
To
Boat
Dog
Truck
Boat
Horse
Figure 5: Consistency style transfer on CIFAR-10. Original class-condiiional images from the dalaset (Lop row) are pushed backward in time to = and then pushed forward using new class. The original styles of the images are maintained, while their subjects are replaced with those from the new class (bottom row).
here, replacing the class label with y' y, we can resample the flow map X (x1 y/) sample the conditional distribution associated to y while maintaining the style of the original class. We demonstrate this principle in Figure 5, which shows that maps discretized with classes. This serves both as verification of the cycle consistency as well as an illustration of the potential applications of the two-time flow map. In practice, the extent of the preserved style depends on how far back in time towards the Gaussian base at = 0 we push the original sample. Here, we use s' = 0.3.
step can be used convert between
5 Conclusion
In this work, we developed framework for learning the two-time flow map for generative modeling: either by distilling pre-trained velocity model with the Eulerian or Lagrangian losses, or by directly training in the stochastic interpolant framework. We empirically observe that while using more steps with the learned map improves sample quality, substantially lower number is needed when compared to other generative models built on dynamical transport. Future work will investigate how to improve the training and the neural network architecture so as to further reduce the number of steps without sacrificing accuracy, and to improve convergence for direct training of one-step maps.
References
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXio:2011. 13445,, 2020.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in neural information processing systems, volume 33, pages 6840-6851, 2020.
Jascha. Sohl-Dickstein, Fric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using noncquilibrium thermodynamies. arXin:1508.03385, 2015.
Yang Song and Stefano Ermon. Improved Techniques for Training Score-Based Generative Models. arXiv:2006.0900 09011, October 2020a.
Yang Song and Stefano Ermon. Generative Modeling by Estimating Gradients of the Data Distribution. arXiu:1907.06600 2020b.
Prafulla Dhariwal and Alex Nichol. Diffusion Models Beat GANs on Image Synthesis. arXiv:2105.05233, 2021.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEFE/CVF conference on computer vision and pattern recognition, pages 10684 10695, 2022.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion English, Kyle Laccy, Alex Goodwin, Yannik Marek, and Robin Rombach. Scaling rectified flow transformers for high-resolution image synthesis. arXi::2403.03200 2024.
13

---

Published in Transactions on Machine Learning Research (05/2025)
Vadim Popor, Ivan Vovk, Vladimir Gogoryan, Tasnima Sadekova, and Mikhail Kudinov. Grad-TTS: A Diffusion Probabilistic Model for Texe-lo-Speech. In Proceedings of the 98th International Conference Machine Learning, pages 8599 8608. PMLR, July 2021.
Myeonghun Jeong, Hyeongju Kim, Sung Jun Cheon, Byoung Jin Choi, and Nam Soo Kim. Diff-TTS: A Denoising Diffusion Model for Text-to-Speech. (Xx::2104.014000 2021.
Rongjie Huang, Zhou Zhao, Huadai Liu, Jinglin Liu, Chenye Cui, and Yi Ren. ProDiff: Progressive Fast. Diffusion Model for High-Qualit tepeeho-S In Proceedings of the 30th ACM International Conference on Multimedia, MM '22, pages 2595-2605, New York, NY, USA, October 2022. Association for Computing Machinery.
Yen-Ju Lu, Zhong- Qiu Wang, Shinji Watanabe, Alexander Richard, Cheng Yu, and Yu Tsao. Conditional Diffusion Probabilistic Model for Speech Fnhancement. In ICASSP 2022 2022 IEFE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 7402 7406, May 2022a.
Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexcy Gritsenko, Diederik P. Kingma, Ben Poole, Mohammad Norousi, David J. Fleet, and Tim Salimans. Imagen Video: High Definition Video Generation with Diffusion Models. arXiv:2210.0220 2022a.
Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J. Fleet. Video Diffusion Models. Advances in Neural Information Processing Systems, 35:8633 8646, December 2022b.
Andreas Blattmann, Robin Rombach, Huan Ling, Tim Dockhorn, Seung Wook Kim, Sanja Fidler, and Karsten Kreis. Align Your Latents: High- Ressutesoo Video Synthesis With Latent Diffusion Models. XrXiv223040888 2023.
Jay Zhangjic Wu, Yixiao Ge, Xintao Wang, Stan Weixian Lei, Yuchao Gu, Yufci Shi, Wynne Hsu, Ying Shan, Xiaohu Qie, and Mike Zheng Shou. Tune-A-Video: One-Shot Tuning of Image Diffusion Models for Text-to-Video Generation. 2rxi::221 2. 11565, 2023.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamn, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling, In The Eleventh International Conference Learning Representations, 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Represenlations, 2022a.
Michael Albergo and Eric Vanden-Fijuden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2022.
Michael S Albergo, Nicholas M Boffi, and Eric Vanden- Eijnnenn. Stochastic interpolants: A unifying framework for flows and diffusions. Xixev preprint arXiv:2303.08799 08797, 2023a.
Cheng Chi, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. Diffusion Policy: Visuomotor Policy Learning via Action Diffusion. arXin:2903.04 137, 2024.
Tan Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative Adversarial Nets. In Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(11): 139-144, October 2020.
Antonia Creswell, Tom White, Vincent Dumoulin, Kai Arulkumaran, Biswa Sengupta, and Anil A. Bharath. Generative Adversarial Networks: An Overview. JEEE Signal Processing Magazine, 35(1):53 -6, January 2018.
14

---

Published in Transactions on Machine Learning Research (05/2025)
Luke Metz, Ben Poole, David Pfau, and Jascha Sohl-Dickstein Unrolled Generative Adversarial Networks. arXiv:1611.02163, 2017.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein GAN. arXin:1701.07875, 2017.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine, Elucidating the Design Space of Diffusion-Bard Generative Models. arXin:2206.00 200066,0000064 2022.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency Models. 0rXiv:2303.0144 2023.
Yang Song and Prafulla Dhariwal. Improved Techniques for Training Consistency Models. arXiv:2310.1418) 2023.
Dongjun Kim, Chieh-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion. arXiv:2310. 02279, 2024a.
Tim Salimans and Jonathan Ho. Progressive Distillation for Fast Sampling of Diffusion Models. arXiv:2202.00512, 2022.
Hongkai Zheng, Weili Nic, Arash Vahdat, Kamyar Azizzaeneshhei,, and Anima Anandkumar. Fast Sampling of Diffusion Models via Operator Learning. arXiv:2211. 2023.
Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009.
Jean-David Benamon and Yann Brenier. A computational fluid mechanics solution to the monge-kantorovich mass transfer problem. Numerische Malhemalik; 84(3):375 393,, 2000.
Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-63):94, 2015.
Esteban G. Tabak and Eric Handen-Eijnden Density estimation by dual ascent of the log-likelihodd. Communications in Malhemalical Sciences, 8(1):217 (1):217-2333, 2010.
E. G. Tabak and Cristina V. Turner. A family of nonparametric density estimation algorithms. Communications Pure and Applied Mothematics, 66(2):145- -64,, 2013.
Danilo Rezende and Shakir Mohamed. Variational Inference with Normalizing Flows. In International Conference on Machine Learning, pages 1530 1538. PMLR, June 2015.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density Estimation Using Real NVP. In International Conference on Learning Representations, page 32, 2017.
Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, llya Sutskever, and David Duvenaud. FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models. arXiv:1810.0136) 2018.
Ricky T. Q. Chen, Yulia Rubanova, Josse Bettencourt, and David Duvenaud. Neural Ordinary Differential Equations. arXiv:1806. 07366. 2019.
Xingchao Liu, Lemeng Wu, Mao Yc, and Qiang Liu. Let us build bridges: Understanding and extending diffusion generative models. orXiv preprint XrXiu:2208.14699 2022b.
Nanye Ma, Mark Goldstein, Michael S. Albergo, Nicholas M. Bolfi, Eric -ajndn, and Saining Xie. SIT: Exploring Flow and Diffusion-basse Generative Models with Scalable Interpolant Transformers. arXiv:2401.087400 2024.
Yifan Chen, Mark Goldstein, Mengjian Hua, Michael S. Albergo, Nicholas M. Boffi, and Eric Canden-Fijnden Probabilistic Forecasting with Stochastic Interpolants and Föllmer Processes. 2rXii:2403.1379) 2024.
Aran-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Ente Enrich, Brandon Amos, Yaron Lipman, and Ricky Chen. Multisample flow matching: Straightening flows with minibatch couplings. arXiv preprint arXiv:2304.11 14777, 2023.
15

---

Published in Transactions on Machine Learning Research (05/2025)
Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brood Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2023.
Michael S. Albergo, Mark Goldstein, Nicholas M. Boffi, Rajesh Ranganath, and Eric Vanden-Fljnden. Stochastic interpolants with data-dependent couplings. arXiv:2310.037 03722, 2023b.
Qinsheng Zhang and Yongxin Chen. Fast Sampling of Diffusion Models with Exponential Integrator. arXiv:2204.13902, 2023.
Alexia Jolicoeur-Martineau, Ke Li, Rémi Pichó-Taillefer Tal Kachman, and Ioannis Mitliagkas. Gotta Go Fast When Generating Data with Score-Based Models. arXiv:2105. 14t:21005.440000 14080, 2021.
Luping Liu, Yi Ren, Zhijic Lin, and Zhou Zhao. Pscudo Numerical Methods for Diffusion Models on Manifolds. 200::22000977, 2022c.
Cheng Lu, Yuhao Zhou, Jianfei Chen, Chongxuan Li, and Jun Zhu. DPM-Solver: A Fast ODE Solver for Dilfusion Probabilistic Model Sampling in Around 10 Steps. ar NNi:2206.00992 2022b.
Daniel Watson, Jonathan Ho, Mohammad Norouzi, and William Chan. Learning to Efficiently Sample from Diffusion Probabilistic Models. arXiv:2106 200.00006,,302,, 2021.
Daniel Watson, William Chan, Jonathan Ho, and Mohammad Norouzi. Learning Fast Samplers for Diffissonn Models by Differentiating Through Sample Quality. arXin:2202.0555 05830, 2022.
Alex Nichol and Prafulla Dhariwal. Improved Denoising Diffusion Probabilistic Models. arXiv:2102.09679, 2021.
Chieh-Hsin Lai, Yuhta Takida, Toshimitsu Uesaka, Naoki Murata, Yuki Mitsufuji, and Stefano Ermon. On the Equivalence of Consistency-Type Models: Consistency Models, Consistent Diffusion Models, and Fokker-Planck Regularization. arXiv:2306.00 (000:2796000967, 2023a.
Chich-Hsin Lai, Yuhta Takida, Naoki Murata, Toshimitsu Uesaka, Yuki Mitsufuji, and Stefano Ermon. Improving score-based Diffusion Models by Enforcing the Underlying Score Bokker-Plancee Equation. arXin:2210.04296 2023b.
Zebang Shen, Zhenfu Wang, Satyen Kale, Alejandro Ribeiro, Aim Karbasi, and Hamed Hassani. Self- Consistency of the Fokker-Planck Equation. arXin:2206.00860, 2022.
Nicholas M. Boffi and Eric Eijuenn Probability flow solution of the Fokker Planck equation. Machine Learning: Science and Technology, 433)00351,, July 2023.
Giannis Daras, Yuval Dagan, Alexandros G. Dimakis, and Constantinos Daskalakis. Consistent Diffusion Models: Mitigating Sampling Drift by Learning to be Consistent. arXiv:2302. 00007 2023.
Beomsu Kim, Jaemin Kim, Jeongsol Kim, and Jong Chul Ye. Generalized Consistency Trajectory Models for Image Manipulation. 1155::240..19510 October 2024b.
Liangehen Li and Jiajun He. Bidirectional Consistency Models. arXiv:2243318033 September 2024.
Eric Luhman and Troy Luhman. Knowledge Distillation Iterative Generative Models for Improved Sampling Speed. arXie:2101.0238, January 2021.
Philip Hartman. Ordinary Differential Equations. Society for Industrial and Applied Mathematics, sccond edition, 2002. doi: 10.1137/1.97808988 97899999999999992222..
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXic:1412.6980, 2017.
M. Raissi, Perdikaris, and G. E. Karniadakis. Physics-informed neural networks: deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. JJurnal of Computational Physics, 378:686-700 2019.
16

---

Published in Transactions on Machine Learning Research (05/2025)
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising Diffusion Implicit Models. arXiv: 2022.
A Probability flow ODEs from stochastic interpolants and score-based diffusion models
For the reader's convenience, we now recall how to construct probability flow ODEs using either flow matching with stochastic interpolants or score-based diffusion models. We also recall the connection between these two formalisms.
A.1 Transport equation
Generative models based probability flow ODEs leverage the property that solutions to (3.2) push forward their initial conditions onto samples from the target:
Lemma A.1 (Transport equation). Let P Law(xz be the PDF of the solution to (3.2) assuming that Po- Then Pt satisfies
(xp V (b(()))(())) = 0, Pt-o(x) po(x))
(A.1)
where V denotes a gradient with respect x.
The interest of this result is that it can be used in reverse: we show that a PDF Pt satisfics the transport equation (A.1), then practice we can use the probability flow ODE (3.2) to sample pt al any time 0, so long as we can sample initial conditions from Pa-
Proof. The proof proceeds via the weak formulation of (A.1). Let & E CRR) denote an arbitrary test function. By definition,
/ (=) = E(((x=))
Vt [0.1]
(A.2)
where Fe is given by (3.1) and where the expectation on the right hand-side is taken over the law of the initial conditions It=o xo Po. Taking the time derivative of this equality, deduce that
J (())) Elit vo(z,)]
by the chain rule using the ODE (3.2)
= =E(ba(a)) Vo(x))
(A.3)
- be(x). (((x) by definition of p (x)
This the weak form of the transport equation (A.1). can be written as (A.1), since it admits strong solutions by our assumptions on (x()..
A.2 Stochastic interpolants and probability flows
The stochastic interpolant framework offers simple and versatile paradigm to construct generative models.
Definition 3.1 (Stochastic Interpolant). The stochastic interpolant I between probability densities po and pr is the slochaslic process given by
7, 04 Barr 722,
(3.1)
where a, 3.2 E ((O, 1)) satisfy as = 3, 1, aj 30 O, and 7o 71 = In (3.1), (æo. 21) is drawn from coupling (Iv, x1) p(z0.01) 21) that satisfies the marginal constraints Jgu p(au, x1)dro = p1 (1(11) and Sge p( on )dr (a())) Moreover, N(O, Id) with (1).
17

---

Published in Transactions on Machine Learning Research (05/2025)
Proposition 3.2 (Probability Flow). For all [0, 11, the probability density of is the same as the probability density of the solution to
it = b((r),
T==0 Fo ~
po
(3.2)
where b: [0 x gd + IRd is the time- denendepend velocity field (or drift) given by
(x(x) = 1
(3.3)
In (3.3), EI- [ = x] denotes an expectation over the coupling (zo, x1) ~ p( and & ~ N(O, Id) conditioned on the event I z.
Proof. Let Φ C; denote an arbitrary test function. By definition,
= (o)) =F(o() B(((T.)
vE [0
(A.4)
where Tt is given by (3.1) and where the expectation the right-hand side is taken over the coupling (Io, x1) ~ p(zo, 21) and ~ N(O, I). Taking the time derivative of this equality, we deduce that
(()))((())d E8L,- vo(l))
by the chain rule
E[E[iA Vo(I))
by the tower property of the conditional expectation by definition of (x()) in (3.3)
(A.5)
-II) vo(l.)
= -/e (x) (é(z))(())dr by definition of pz(x)
This is the weak form of the transport equation (A.1).
This result implies that the PDF P: of the stochastic interpolant Ie is also the PDF of the solution 2 of the probability flow ODE (3.2) with the velocity field (x) defined in (3.3). Algorithmically, this means that we can use the associated flow map as a generative model.
A.3 Score-based diffusion models
For simplicity, we focus on variance-preservingg diffusions; extensions to more general classes of diffusions are straightforward. These models are based on variants of the Ornstein-Uhlenbeck process, defined as the solution the stochastic differential equation (SDE):
dX, -xtt + V2dW,, Xi=o= Px:
(A.6)
where We is Wiener process. The solution to (A.6) for the initial condition X a is:
a + v2 f -4t dWy.
(A.7)
By the Itô Isometry, the second term on the right-hand side is a Gaussian process with mean zero and covariance given by
B ( aw )} 21d (' 224-21 di' = (1 22 e-20)
(A.8)
Therefore, at any time t, the law of (A.7) that of Gaussian random variable with mean ze + and covariance (I 22) That is, it can be represented as
X at 22
(A.9)
18

---

Published in Transactions on Machine Learning Research (05/2025)
where a Px, ~~ (0, Id), and - 2. (A.9) shows that the law of X, converges exponentially fast to that of the standard normal variable as & - oc for all initial conditions a.
The key insight of score-based diffuusin models is to time-reverse the SDE (A.6) to obtain a process that turns Gaussian samples into samples from the target. Alternatively, this time-reversal can be performed using the associated probability flow ODE. To do so, we can start from the evolution equation for the PDF of the solution the SDE (A.6). Denoting this PDF by Pe, it satisfies the Fokker-Planck equation (FPE)
Oipt V (xri) + Apt,
Po P==
(A.10)
To time-reverse the solution to (A.10) we define Pt = pr- for some large T > O, and derive an equation for Pe from (A.10):
eipe = -
V (xpT-1) Apr
(A.11)
-V. P((x V log pr p
- v-(2 V log pelp:)
where we used the identity Apr- = V (pr- V log 11) Equation (A.11) can be written as
0.p + V ([x Sy (x))ptt = 0, Pt=o = P N(0, Id)
(A.12)
where sx V (c(x) V log (x) is the score. This quantity can be estimated by regression using
the solutions to (A.6), which are explicitly available via (A.9). Equation (A.12) is in the form of (A.1) with velocity field (x) given by
(x()) = + (x).
(A.13)
Flow map matching can therefore be used to distill score-based diffusion models using this velocity field.
A.4 Connections
Here we show how ecore-based diffusion models can be related to stochastic interpolants. Recall that the law of the solution to the SDE (A.6) is given by (A.9). If we change time according to + - log L we map [0, oo) onto [1, 0) and we arrive
X Ing
d 4 =1=ot at VI t2
a~p ~~ Id), and a 1
(A.14)
This is a valid stochastic interpolant if we set a= ==, 3: 1, or = 0, and 7 = - in (3.1)3. Hence, we have shown that the probability flow of the coore-based diffusion model with velocity field (A.13) can also be studied by using the stochastic interpolant (A.14). Note that this reformulation has the advantage that it eliminates the small bias incurred by the finite value T in score-based diffusion models, since the stochastic interpolant (A.14) transports the Gaussian random variable time = 0 onto the data point a at time exactly.
B Proofs for Section 3
Proposition 3.5. The flow map xx.(x) is the unique solution to the Lagrangian equalion
(x),, (x)),,
Xx,(x)= x,
(3.7)
for all (,)) E [0, 1|2 x Rdd In addition, il satisfies
X (X,, Xz,
(3.8)
for all (,,,,,x) E [0, x Rd. In particulor X (X1,s (z)) x for all (s, 2)) e [0, 112 Rd L.e. the flow map is invertúble.
Note that 20 1 here, consistent with the fact that the base PDF Gaussian, that ag and can be lumped into single Gaussian variahle
19

---

Published in Transactions on Machine Learning Research (05/2025)
Proof. Taking the derivative with respect to t of (x,) 2t, deduce
0 0,X,, it (x())) (xx))))
(B.1)
where we used the ODE (3.2) to obtain the second equality. Conversely, since the solutions (3.2) and (3.7) from given initial condition are unique, if we solve (3.7) on [s, t] with the condition X, (2x) 2, must
have x (2,) Ia for all (s.t) [0, 10. Evaluating this expression at 2, gives (3.7). Also, for all (8,4,t) E [0, 7)s, we have
(x, (x.)) = Xe,((x)) = (x,)(1.))
(B.2)
Evaluating this expression at 2, = x gives (3.8).
Corollary 3.7 (Lagrangian map distillation). Let Wa,t E [[, 113) be weight function salisfying wa,t > 0 and let be the stochastic interpolant defined in (3.1). Then the flow map is the global minimizer over X of the loss
(L) = -ouup Wo,, B(i. (I,)) bi( (T.)
(3.11)
subject to the boundary condition that (x,(()) x for all x E IR and E 1). In (3.11), E denoles on expectation over the coupling (Io, x1) ~ x1) and ~ N(O, Id).
Proof. Equation (3.11) is a physics-informed neural network (PINN) (Raissi et al., 2019) loss that is minimized only when the integrand is zero, i.c., when (3.7) holds.
Proposition 3.6. The flow map X is the unique solution of the Fulerian equation,
(x) + (x), =0, Xt,(xx) = =I
(3.10)
for all (s,t,x) 1|. x Rdd
Proof. Taking the derivative with respect to of (Xt,,)()))))) and using the chain rule, we deduce that
d
0 Xx, X (X,, 0,X., (X, (x)) + 0,X,, (z) (X,. =0,X.. (x)) + (x). (X..
ds
(B.3)
where we used (3.7) to get the last equality. Exaluating this expression at Xt,(xx) = y, then changing inlo 2, gives (3.10).
Corollary 3.8 (Eulerian map distillation). Let wa and Is be as in Corollary 3.7. Then the flow map is the global minimizer over X of the loss
LEMD (8) - = -/o W1o,E[008aa.11, (T..)) (X..(t))P
(3.13)
subject to the boundary condition (x) = x for all x E Rdd and for all E IR
Proof. Equation (3.13) is PINN loss that is minimized only when the integrand is zero, i.e, when (3.10) holds.
Proposition 3.9 (Lagrangian error bound). Let Xas Rd - Rd denote the flow map for b, and let x,,: T ddd denote an approrimate flow map. Given zo Po, let Pr be the densily of X (2o) and let pr be the target density of Xo,1 Then,
(3 < c 2+2 (clle (L)
(3.15)
where Ct is the constant that oppears in Assumption 3.3.
20

---

Published in Transactions on Machine Learning Research (05/2025)
Proof. First observe that, by the one-sided Lipschitz condition (3.5),
ar(Xa,(D) (v)) (x) Xx,((v)) (b((Xx,((z)) (((X).(()))), 2C, - xx,(
(B.4)
Equation (B.4) highlights that (3.5) gives a bound on the spread of trajectories. We note that we allow for C₂ < 0, which corresponds to globally contracting maps. Given (B.4), we now define
aat
Xx,(
(B.5)
where we recall that (x) satisfies (x),,((z)) = (x(,,(((z)) and Xx,o(T) = x. Taking the derivative with respect to of (B.5), we deduce
O, 10c.=2B[[(x (Xs, (l,((1))) (.,((()))(bb(((.... -x) 0(),,
(Is,
(X - (,) -lx,,,
- (,,(( (
(b((...()))) by (X.,()))))
(B.6)
< E - BB - (1.)
- -(,,(1
(X - (l.((s)) (b. (b.(Xx,,(L)) -
(12,(())))
= Fg,t + SLMD + 21F (X (1.)- bx (T.)))
Above, we defined the two-time Lagrangian distillation error,
BLD
( - O x,, (T,)
(B.7)
By definition, the LMD loss can be expressod as (L)) SO.TTp 1o SLMD dsdt. Using (3.5) in (B.6), we
obtain the relation
O,Eo, & (1 (1+ 2C,) + 5LMD
(B.8)
which implics that
0.(e -t-2 2. Cadu EsS) Se -t-2 C.du gLMD
(B.9)
Since E.,e = O this implies that
Eg & '' (t-w)+2 S Cidt BLMD du ++2 CCildt [ SLMD du.
(B.10)
Above, we used that (t, u) E [0, ff2 that (t- u) & t. This bound shows that E..1 & <e' which can be written explicitly as (using instead of u dummy integration variable)
1+2 |Ce|de
So sLMD du,
(zo)()))))) -x.t)))))) < 1+2 ladd S TElla - a 2.
(B.11)
Now, observe that by definition,
(e & - (2o))
So,1
(B.12)
because the lent-hand side is the infimum over all couplings and the right-hand side corresponds to a specific coupling that pairs points from the same initial condition. This completes the proof.
Proposition 3.10 (Eulcrian error bound). Let Xx, : Rd + Rdd denote the flow map for b, and let X : Rdd + Rdd denote approcimate flow map. Given xo pa, let ir be the density of Xa,1 (xo) and let pr be the
target density of (2,) Then,
(2(0), e LEMD(
(3.16)
21

---

Published in Transactions on Machine Learning Research (05/2025)
Proof. We first define the error metric
E, = E - -x,1()))
(B.13)
It then follows by direct differentiation that
O,E,, = 2 (4,) (L,(())) (L.),, vX..(4.) (.,,(())))...(())) - (0x ([
E 2 (Ta) (T,) (1. VX..((.) -0.8 (T.) (lo))) (Ta.s((s))
-Eg - E (1.) + bo())) (1,) (0.8s, (1.) + (.,()) VX (.).(L.))
= - - 8,8
JEMD 8EMD
Above, we used the tower property the conditional expectation, the Eulerian equation 0, b, (I,, = 0, and defined the two-time Eulerian distillation error,
8EMD
E (0. (1.).(44))
(B.14)
This implies that
0z( Ex,)) e 8EMD
(B.15)
Using that E,,t O for any t E [] and integrating with respect to s from to t,
e' Ei, + Es, & [ e enDEMD du.
(B.16)
Tt then follows that
F.s E S e" en-sSEED
SEMD du,
(B.17)
and hence, using that [0 and that SEMD > 0,
(o. - (z) < ee S (128a111))+bb4) (lo, 27 ds.
(B.18)
The proof is completed upon noting that
(31, < Xo. - 8o, (2
(B.19)
because the left-hand side is the infimum over all couplings and the right-hand side corresponds to a particular coupling.
Proposition 3.11 (Flow map matching). The flow map is the global minimizer over X of the loss
LFMM( - - Whst
- 1 ER,, - I1l²] I+ dsdt.
(3.17)
10.11
subject to the boundary condition (x) x for all a gd and for all E R. In (3.17), Wa,t > 0 and is taken over the coupling (I), 21) and N(O, Id).
Proof. We start by noticing that
E ( (T)) (T), 1.P,
21, 9, 8.8., (8.,( + +1117,
= ( -2E(i.)) 0,8 (x,, + 1Td]].
(B.20)
=B[iats.,1 (Xe,,
= -E( (Ix)) 2b. 0, & (T)) + +1.2]
22

---

Published in Transactions on Machine Learning Research (05/2025)
where used the tower property of the conditional expectation to get the third equality and the definition of (x( in (3.3) to get the last. This means that the loss (3.17) can be written as
(X)
JUAP I Want [10, (Xt,,)) (x())) (8,((X),())) - (
(B.21)
+ f Mo.r wan 11 (I,) dsdt.
where Pr La)( The second integral does not depend on X, so it does not affect the minimization of (x) Assuming that Wa,t > O, the first integral is minimized if and only if we have
A (s, 2.) E [0, 112 Rd 0,&,, (X,,x)))))(())) = (x)) and (x))()
(B.22)
From the second these equations it follows that: (i) (x, = x, and (ii) if we evaluate the first equation at (x), this equation reduces to
[ (s,t,)) E [0, Rd
0.8aclu) = be(Raa(v))
(B.23)
which recovers (3.7).
Proposition 3.12 (Eulerian estimation). The flow map X.,t is critical point of the loss
10.11 LEE(X) - - Wa ((0.X.(.1...111) stopgrad(i. (8aa((a))) dodl,
(3.22)
subject to the boundory condition (x)(x)= = x for all x E Rd and for all s E R. In (3.22), Watt > and E taken over the coupling (Io, x1) ~ p(ro, xI), N(O, Id). The operator stopgrad(.) indicates that the argument is ignored when computing the gradient.
Proof. The functional gradient of (3.22) with respect to 8as is given by
0, (wa,( (0.8.. (z) E v&..(t)) I. -1= ====\\]((())))
(B.24)
D (We.e (..., (x) + (z() (=),,(
The flow map zeroes this quantity since it solves the Eulerian equation (3.10) that appears in it, and hence is critical point of the objective.
Lemma 3.13 (Progressive flow map matching). Let & be two-time flow map. Given K E N, let t = A (t-s)
-5) s) for k =, I,...,.. K. Then the unique minimizer over X of the objective
CPFMM (x) /
J0.11 Wx, B
(Xtx-1AK ....
dsdt,
(3.23)
produces the same output in one step the K- K-step iterated map X. Here watt > 0, and E is taken over the coupling (mo, m1) ~ p(x0, 21) and ~ N(O, Id).
Proof. Equation (3.23) is a PINN loss whose unique minimizer satisfies
A (s,t,z) 00, x R
Xx((x) (8ex-1 o...o ...
(B.25)
which establishes the claim.
C Relation to existing consistency and distillation techniques
In this section, we recast consistency models and several distillation techniques in the language of our two-time flow map Xs to clarify their relation with work.
23

---

Published in Transactions on Machine Learning Research (05/2025)
C.1 Flow maps and denoisers
Since Law(Xx,:(4)) Law(L.), it is tempting replace (1,)(Is) by 1g in the loss (3.17) and use instead
= - Wa,,
Liteniio [8] = 10.112
I, dsdl,
(C.1)
minimized over all such that (2) =. However, the minimizer of this objective is not the flow map Xs,,, but rather the denoiser
(d)noise(x)) - -E/4// = z].
(C.2)
s,
This can be seen by noticing that the minimizer of (C.1) is the same as the minimizer of
denoise [8] = /
1..11 ws, (I.) -[[,,,, dedt,
(C.3)
J Jp.pp
L W6 (1,,8 -
--I-I, =
which follows from an argument similar to the one used in the proof of Proposition 3.11. The minimizer of (C.3) satifies
(x) = = 0, = x],
(C.4)
which implies (C.2) by the boundary condition xx, = The denoiser (C.2) may be useful, but it is not consistent generative model. For instance, if 20 and x1 ~ are independent in the definition of 4, since Iu = 2g and 1 = 21 by construction, for = and t = we have
denoise (x) = E[v1] X
(C.5)
i.e. the one-step denoiser only recovers the mean of the target density Pt-
C.2 Relation to consistency models
Noising process. Following the recommendations in Karras et al. (2022) (which are followed by both Song et al. (2023) and Song and Dhariwal (2023)), we consider the variance-exploding process"
i = + tz, t E [0, fmax],
(C.6)
where a ~ Pl (data from the target density) and N(O, I). In practice, practitioners often set fmx 80. In this section, because we follow the score-based diffusion convention, we time so that 0 recovers P1 and that Gaussian is recovered
co. The corresponding probability flow ODE given by
& - (ig))
Ōτ=0 = P1
(C.7)
where Pr(x)= Law(āz). In practice (C.7) is solved backwards in time from some terminal condition ". To make contact with our formulation where time goes forward, we define 2, = 1mmx-b leading to
f = (lmax V log Pt. Ptmee-
Fr-o N( l
(C.8)
To make touch with our flow map notation, we then define
0xX,,((x) = (fmax t) V log Ptm... (X
(x) = x.
(C.9)
1Oftentimes 0 set. Amin
numerical stability, choosing egg. fmin 2 -0
24

---

Published in Transactions on Machine Learning Research (05/2025)
Consistency function. By definition (Song et al., 2023), the consistency function fo :R: - IR is such that
f(
(C.10)
where on denotes the solution of (C.7) and a ~ P1- To make connection with our flow map formulation, let us consider (C.10) from the perspective of Tt.
fa fe( = Flmax:
(C.11)
which is to say that
f((z) Xtmax tilmas (x).
(C.12)
Note that only one time is varied here, i.e. fe(x), cannot be iterated upon: by its definition (C.10), it always maps the observation o, onto a sample
Discrete-time loss function for distillation. In practice, consistency models are typically trained in discrete-time, by discretizing (min mmax] into a set of N points min = l < l2 <... < IN maxx- According to Karras et al. (2022), these points are chosen as
t, =
+
i-1 N --1
(1/n 1/n mmax
mmin
(C.13)
with 1- Assuming that we have at our disposal a pre-trained estimate sx()) of the score V log (zr), the distillation loss for the consistency function f(( is then given by
( Σ E [I (Fterr) fr (ow)))
(C.14)
11 lit1
b. .1.4 (t: t,+1) 144184441 (Zt44);
where E is taken over the data a~ P1 and 2 N(O, I). The term ft, is an approximation of ž., computed by taking single step of (C.7) with the approximate score model sx()). In practice, the square loss in (C.14) can be replaced by an arbitrary metric d: Rd Rd 2z,, such as a learned metric like LPIPS or the Huber loss.
Continuous-time limit. In contimuous-time, time, it is easy to see via Taylor expansion that the consistency loss reduces to
S
CE lim N (D N
- (x) vf((z))² ((((=)))((()dddd))) (x)ddatt
(C.15)
fmin
where Wt = ( 1/1)6111/7 is a weight factor arising from the nonuniform time-grid. This a particular case of our Eulerian distillation loss (3.13) applied to the varianco-cxploding setting (C.6) with the identification (C.12).
Estimation distillation of the consistency model. If we approximate the exact
V Vlog(((>)
E
5t
(C.16)
12
by single-point estimator
V Vlogp((>) a-Fr Ft
(C.17)
may the expression
G1 + 42
(C.18)
25

---

Published in Transactions on Machine Learning Research (05/2025)
in (C.14) to obtain the estimation loss,
N-1 (E Σ E ( Sez (fx,
(C.19)
T4. = a aa+42 1.5.
This expression does not require a previously-trained score model. Notice, however, that (C.19) must be used with stopgrad fr, (*,,) so that the gradient is taken over only the first Sta (#t, This is because (C.14) and (C.19) are different objectives with different minimizers, even at leading order after expansion in 1/N, for the same reason that (3.13) differs from (3.18). To see this, observe that to leading order,
(Gx.) fūnn .) (OeSais (Tt
+ + Vien (Zt.) (4: 4411 +(44+ (411)))
(C.20)
which gives the continuous-time limit
fmax
/
(f) - Jim C
wxlOrfe(x)+ + V (x)l ( (r(rddddl.
(C.21)
N
Amin
Observing that 0,ă, shows that (C.21) recovers the Eulerian estimator described in Section 3.6, which does not lead to gradient descent iteration.
C.3 Relation to neural operators
In our notation, neural operator approaches for fast sampling of diffusson models (Zheng et al., 2023) also estimate the flow map Xo.t via the loss
( = [/
S 1Xos(()) - Xox(z)) (o (d)dd,
(C.22)
where Xo,t is parametorized by Fourier Neural Operator and where Xu,t the flow map obtained by simulaling the probabilily flow ODE associaled with pre-trained (or given) be(x). To avoid simulation at learning time, they pre-generate dataset of trajectories, giving access to (z, for many initial conditions
Pn- Much of the work focuses on the architecture of the FNO itself, which is combined with U-Net.
C.4 Relation to progressive distillation
Progressive distillation (Salimans and Ho, 2022) takes a DDIM sampler (Song et al., 2022) and trains a new model to approximate two steps of the old sampler with one step of the new model. This process is iterated repeatedly to successively halve the number of steps required. In our notation, this corresponds minimizing
(X) = [ 1-221
[ [XLc422((x) - (X1+A14428 (z)+ (x()ddadt
(C.23)
where is a pre-trained map from the previous iteration. This then iterated upon, and At increased, until what left few-step model.
Additional Experimental Details
D.1 2D checkerboard
Here, we provide further discussion and analysis of our results for generative modeling on the 2D checkerboard distribution (Figure 3). Our KI--divergence estimates clearly highlight that there a hicrarchy of performance. Of particular interest is the large discrepancy in performance between the Eulerian and Lagrangian distillation techniques.
26

---

Published in Transactions on Machine Learning Research (05/2025)
(L(PillĀ1)
()( P1)
L error
SI
0.020
0.026
0.0
0.000
LMD
0.043
0.059
0.032
0.085
EMD
0.079
0.029
0.010
0.011
FMM, N =
0.104
0.021
0,026
FMM, = 4 0.045
0.014
0.024
PFMM, N= ==1 0.043
0.014
0.023
Table 3: Comparison of (L(pllp) and W (p. p), where á is the pushforward density from the maps No, (200) for the methods listed above. Additionally included is comparison of L2 expected error of the distillation methods against their teacher 1 given as [8[ (a) (x)))) Intriguingly, LMD performs better being distributionally correct, measured by the KL-divergence, but worse in preserving the coupling of the teacher model. The roles flipped for EMD. This may highlight KL as a more informative metric in our case, as our aims are to sample correctly in distribution. See Figure 6 for visualization.
As noted in Figure 3 and Table 3, LMD substantially outperforms its Eulerian counterpart in terms of minimizing the KT-divergence between the target checkerboard density Pi and model density ý = 15 Interestingly, while LMD more correct in distribution, EMD better preserves the original coupling (ro, X (zr)) (ro)) of the teacher model X6 as measured by the W distance and the expected L2 reconstruction error, defined B[] (x) (z))(= Where this coupling is significantly not preserved is visualized in Figure 6. For each model, we color code points for which [ (z) (z) > 1.0, highlighting where the student map differed from the teacher. We notice that the LMD map pushes initial conditions to an opposing checker edge (purple) than where those initial conditions are pushed by the interpolant (blue). This much less common for the FMD map, but its performance is overall worse in matching the target distribution.
SI, 80
FMM, N
FMM, N ==
PFMM, N=
LMD, N=1
EMD, N
Figure 6: Visualization of the difference in assignment the maps (xo) for the various models as compared to the teacher/ pacher/groond ground truth model 881 (zu) for the same initial conditions from the base distribution 2p. Points that lie in the region (&S (zo) (zo)( 2 > 1.0 are colored compared to the blue points, which represent where the stochastic interpolant teacher mapped the same red initial conditions. This gives us an intuition for how well each method precisely maintains the coupling (æo, X8 (2o)) from the teacher. Note that
treating Xa the ground truth map here, as it is close to the exact map. The models based on FMM either don't have a teacher or have FMM, NN
teacher, but all should have the coupling at the minimizer.
D.2 Image experiments
Here we include more experimental details for reproducing the results provided Section 4. We use the U-Net from the diffusion OpenAl paper (Dhariwal and Nichol, 2021) with code given at https://githull cot.huuaacoo/ooaaaaaaaaaaoaaaaaa wwwm///paaa/guidddddifiiiii guudcdddiiffusion. We use the same architocture for both CIFAR10 and ImageNet- 32 x 32 experiments. The architecture is also the same for training velocity field and for training flow map, barring the augmentation of the time-step embedding in the U-Net to handle two times (s, t) instead of one. Details of the training conditions are presented in Table 4.
27

---

Published in Transactions on Machine Learning Research (05/2025)
CIFAR-10
ImageNet 32x32
Dimension
32x32
32x32
# Training point
X 104
1,281,167
Batch Size
256
256
Training Steps (Lagrangian distillation)
1.5x 10
2.5x105
Training Steps (Eulerian distillation)
1.2x10
2.5x 10
Training Steps (Flow map matching)
N/A
1x10
Training Steps (Progressive flow map matching)
1.3
N/A
U-Net channel dims
256
256
Learning Rate (LR)
0.0001
0.0001
LR decay (every 1k epochs)
0.992
0.992
U-Net dim mult
[1,2,2,2
[1,2,2.2]
Learned time embedding
Yes
Yes
# GPUs
4
Table 4: Hyperparameters and architecture for image datasets.
28


Reviews:
[{"review_id": "u25IBbWI6d", "reviewer": "Reviewer_NFyf", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n- [S1] To the best of my knowledge, this is the first work to leverage the Lagrangian equation to learn few-step generators along neural-ODEs.\n\n**Weaknesses**\n- [W1] Theoretical concerns. There is no proof that the flow map is the only map which satisfies the Lagrangian equation. Specifically, the authors only show that flow map satisfies the Lagrangian equation. I was unable to find a proof that a map which satisfies the Lagrangian equation is the flow map. Hence, it is difficult to see whether training a neural net to satisfy the Lagrangian equation will truly yield one-step or few-step generators along neural ODEs.\n- [W2] Weak experiments. I do not find the experimental results very convincing. Indeed, in Figure 3, FMM with N=1 is highly inferior to SI, which enhances my concern in [W1] - that training a neural net to satisfy the Lagrangian equation may not yield one-step or few-step generators along neural ODEs.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes Flow Map Matching (FMM) for one or few-step generation along neural ODEs. The authors show that the flow map, a function which translates points between two arbitrary time-points of neural ODEs, satisfy a certain Lagrangian equation (e.g., Eq. (3.5)). Inspired by this observation, the authors propose optimizing a neural net to satisfy the Lagrangian equation to learn the flow map. The authors numerically show that FMM is capable of learning two-dimensional distributions or 32x32 resolution images.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Addressing both weaknesses are critial to securing my recommendation for acceptance. Please provide clarifications or additional proofs to address [W1], and provide further experimental results to address [W2].", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "HrMY6fQuqb", "reviewer": "Reviewer_ZNnM", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\nOverall, the paper is nicely written and reasonably easy to follow, even for someone who is not an expert in dynamics-based generative modelling like me. The figures in the text are also very nice.\n\nThe results showing the existence and uniqueness of the solutions to their proposed objectives follow in a reasonably straightforward fashion from previous results, but they are important to establish nonetheless. I had a cursory check of the proofs, and apart from two inequalities (see below), I believe they are correct. \n\n## Weaknesses\nI don't think the paper suffers from any major weakness that should prevent it from appearing in TMLR.\n\nHowever, there are a few points that could be improved to strengthen the paper:\n - Related works: The authors' proposed approach seems very similar to the work of [1] and [2], but these are not cited. Could the authors explain how their method is different from these? At a glance, it seems that the authors have more theoretical results/their framework is slightly more general, whereas these papers do more experiments.\n - I appreciated the error bounds in Section 3.3, but how tight these are is unclear. Are there comparable bounds or bounds that are special cases of the ones in the paper present in the literature? Or are these the first such error bounds to appear? In any case, the authors should note this.\n - The four objectives proposed by the authors are always compared as equals in the experiments. However, it is unclear whether this means we should always prefer the LMD formulation (the best-performing one) or if other potential factors should guide this decision. It would be great if the authors could include a paragraph or two discussing when to apply which objective or present the necessary assumptions and design decisions in a table.\n - Why do the inequalities in A.11 and between A.18 and A.19 hold?\n\n## References:\n - [1] Li, L., & He, J. (2024). Bidirectional Consistency Models. arXiv preprint arXiv:2403.18035.\n - [2] Kim, B., Kim, J., Kim, J., & Ye, J. C. (2024). Generalized Consistency Trajectory Models for Image Manipulation. arXiv preprint arXiv:2403.12510.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper tackles the issue of evaluating generative models defined via differential equations, such as diffusion and flow matching models, sampling from which usually relies on expensive numerical integration.\nTo solve this problem, the authors propose to learn what they call the *two-time flow map* of the model, i.e. the integrator of the dynamics that define the generative model. In other words, instead of handling the model dynamics and integrating them directly, they propose learning a function that takes an integration start time $s$, end time $t$ and the state at the start time $x_s$ and outputs the state $x_t$ at time $t$.\nThey propose four possible objectives for learning the flow map based on different formulations of the problem depending on the information available to the practitioner. They benchmark their solution on learning a 2-dimensional toy distribution and an image generation problem on CIFAR-10 and ImageNet 32x32.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Necessary:\n - Please address my concerns in the weaknesses section.\n - Bottom of page 4: \"by providing access to alternative samples from\" - I believe it should be \"by providing an alternative to access samples from\".\n - Please define what T-FID is.\n - Please increase the legend size in Fig 4/B; it should be easily readable at 100% zoom.", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Q5bDwjVnrg", "reviewer": "Reviewer_bE2k", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "See requested changes section for specific questions.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Authors propose using two types of losses (Lagragian map distillation-based loss and Eulerian map distillation-based loss) to obtain a flow map, which is an image generation function. \n\nThe formulation of these two distillation losses are motivated from an autonomous system ODE and conditions of density functions of incompressible flow. \n\nThe main gap authors want to bridge is the necessity of multiple steps of generation from extant methods. \n\nAuthors presented experimental results to generate synthetic images similar to those of CIFAR10 and ImageNet, and showed Frechet Inception Distance metrics to show the generated images are of better quality compared to those by DDPM and BatchOT.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Questions:\n\n- How is the _velocity field_ $b_t$ pre-trained? In the context of image generation, for instance based on CIFAR-10 and ImageNet datasets, how is this $b_t$ pre-trained and what physical meaning does it have?\n\n- In section 3.4, how is $\\dot{I}_t$ computed? Even taking a step back, please explain clearly the relationship between $I_t$, $\\rho_t$ and $b_t$. How are they pre-computed? How are they prepared in the numerical experiments for PFMM and FMM methods, respectively?\n\n- What is \"_identification $t$_\" above proposition 3.11?\n\n- In order to make method comparison fair, also to highlight the main motivation of this work, does it make sense to benchmark with other methods that also claim the benefit of few-step generations? See [1-8] and references in the Related Work section in [1]. I am not asking authors to make comparison to all those methods, while some other methods which aim to address the computationally expensive multi-step simulation problems should be proper baselines.\n\n- The benefits or drawbacks of the proposed methods are not clearly explained in the experimental section. What is the \"_Baseline_\" in Table 1. Table 1 shows that in terms of the FID metric, proposed methods underperforms against the \"_Baseline_\". \n\nAuthors should give more careful thoughts distinguishing the advantages and disadvantages of each proposed method, and give a characterisation on when to use which. Roughly talking about their computational time difference is insufficient. It leaves the impression that methods are simply stacked up without being well understood.\n\n\nMinor comments:\n\n- Define notations used in the manuscript clearly, such as $X_{0,1}\\sharp\\rho_0$, $\\perp$, $\\mathrm{Law}(I_t)$.\n\n- Define clearly what $\\hat{X}$ and $\\check{X}$ are.\n\n\n*************************************\n[1] Clément Chadebec, Onur Tasar, Eyal Benaroche, Benjamin Aubin, Flash Diffusion: Accelerating Any Conditional Diffusion Model for Few Steps Image Generation, 2024, arXiv:2406.02347v2 \n\n[2] Yin et al., One-step Diffusion with Distribution Matching Distillation, arXiv preprint arXiv:2311.18828, 2023.\n\n[3] Shanchuan Lin, Anran Wang, and Xiao Yang. Sdxl-lightning: Progressive adversarial diffusion distillation. arXiv preprint arXiv:2402.13929, 2024.\n\n[4] Xingchao Liu, Chengyue Gong, et al. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2022\n\n[5] Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, et al. Instaflow: One step is enough for highquality diffusion-based text-to-image generation. In The Twelfth International Conference on Learning Representations, 2023.\n\n[6] Dongjun Kim, Chieh-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning probability flow ode trajectory of diffusion. In The Twelfth International Conference on Learning Representations, 2023.\n\n[7] Simian Luo, Yiqin Tan, Suraj Patil, Daniel Gu, Patrick von Platen, Apolinário Passos, Longbo Huang, Jian Li, and Hang Zhao. Lcm-lora: A universal stable-diffusion acceleration module. arXiv preprint, arXiv:2311.05556, 2023.\n\n[8] Axel Sauer, Frederic Boesel, Tim Dockhorn, Andreas Blattmann, Patrick Esser, and Robin Rombach. Fast high-resolution image synthesis with latent adversarial diffusion distillation. arXiv preprint, arXiv:2403.12015, 2024.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: As the reviewers acknowledge, the proposed method is novel, and the theoretical results are solid. The reviewers had several concerns regarding the theoretical contributions and presentation of the results, but the authors adequately addressed the concerns in the revised version. Although the concerns regarding the scale of the experiments remain valid, the strengths of the paper outweigh this weakness. Therefore, I recommend the paper for acceptance at TMLR.
Decision: Accept (poster)