How to build a consistency model: Learning flow maps via self-distillation
Authors: Nicholas Matthew Boffi
Michael Samuel Albergo
Eric Vanden-Eijnden
Venue: NeurIPS 2025 poster ()
Keywords: consistency models
diffusion models
stochastic interpolants
flow matching
Primary area: deep_learning
TL;DR: We introduce a framework for training accelerated, few-step generative models that includes consistency models, shortcut models, and mean flow as special cases.

Abstract: Flow-based generative models achieve state-of-the-art sample quality, but require the expensive solution of a differential equation at inference time. Flow map models, commonly known as consistency models, encompass many recent efforts to improve inference-time efficiency by learning the solution operator of this differential equation. Yet despite their promise, these models lack a unified description that clearly explains how to learn them efficiently in practice. Here, building on the methodology proposed in Boffi et. al. (2024), we present a systematic algorithmic framework for directly learning the flow map associated with a flow or diffusion model. By exploiting a relationship between the velocity field underlying a continuous-time flow and the instantaneous rate of change of the flow map, we show how to convert any distillation scheme into a direct training algorithm via self-distillation, eliminating the need for pre-trained teachers. We introduce three algorithmic families based on different mathematical characterizations of the flow map: Eulerian, Lagrangian, and Progressive methods, which we show encompass and extend all known distillation and direct training schemes for consistency models. We find that the novel class of Lagrangian methods, which avoid both spatial derivatives and bootstrapping from small steps by design, achieve significantly more stable training and higher performance than more standard Eulerian and Progressive schemes. Our methodology unifies existing training schemes under a single common framework and reveals new design principles for accelerated generative modeling. Associated code is available at https://github.com/nmboffi/flow-maps.

How to build a consistency model: Learning flow maps via self-distillation
Nicholas M. Boffi Carnegie Mellon University
Michael S. Albergo Harvard University
Eric Vanden-Fijnden Courant Institute of Mathematical Sciences
Abstract
Flow-based generative models achieve state-of-the-wee sample quality, but require the expensive solution of differential equation at inference time. Flow map mod- els, commonly known as consistency models, encompass many recent efforts to improve inference-time efficiency by learning the solution operator of this differ- ential equation. Yet despite their promise, these models lack a unified description that clearly explains how to learn them efficiently in practice. Here, building on the methodology proposed in Boffi et al. (2024), we present a systematic algo- rithmic framework for directly learning the flow map associated with flow diffusion model. By exploiting relationship between the velocity field underlying a continuous time flow and the instantaneous rate of change of the flow map, we show how to convert any distillation scheme into a direct training algorithm via ceff-distillation, eliminating the need for pre-trained teachers. We introduce three algorithmic families based on different mathematical characterizations of the flow map: Fulerian, Lagrangian, and Progressive methods, which we show encompass and extend all known distillation and direct training schemes for consistency mod- els. We find that the novel class of Lagrangian methods, which avoid both spatial derivatives and bootstrapping from small steps by design, achieve significantly more stable training and higher performance than more standard Eulerian and Pro- gressive schemes. Our methodology unifies existing training schemes under a single common framework and reveals new design principles for accelerated generative modeling. Associated code is available at https://gitab.coo/mmb/fff//tiiit www.f..//ooommmp......p maps.
1 Introduction
Generative models based on dynamical systems, such as flows and diffusions, have achieved remark- able successes vision (Song al., 2020; Rombach et al., 2022; Ma et al., 2024: Polyak et al., 2025), language (Lou et al., 2024), protein structure prediction (Abramson et al., 2024), weather forecast- ing (Price et al., 2024), and materials design (Zeni al., 2025). While highly expressive, dynamical models leverage the solution of a differential equation for sample generation, which typically requires repeated evaluation the learned model. This computational bottleneck has limited the application of flows and diffusions in domains where rapid inference is crucial, such real-time control (Black et al., 2024; Chi et al., 2024) and image editing, and as a result has led to intense interest in accelerated inference. One particularly promising approach, which underlies consistency models (Song et al., 2023; Kim et al., 2024), is to estimate the flow map associated with the deterministic probability flow equation instead of the velocity field governing its instantaneous dynamics. The flow map is defined by the integrated flow and can be used to generate samples in as few as one model evaluation, leading to inference that can be 10- 100x faster than traditional dynamical models. This dramatic speedup
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

A
B
lim (x,
x. X,
v,00
1-8
xx,
Xaa(() x,000) Xartr)
X,10) -x
iim, v2,00) b(0x)
Figure 1: Overview. (A) Schematic of the two-time flow map Xs,e and the tangent condition (Lemma 2)1), which provides relation between the map and the drift of the probability flow. The flow map is composable, invertible, and has the property that f - 8, its time derivative recovers the drift b. from (2). (B) Illustration of our proposed parameterization. The function v,,, estimates the slope the line drawn between two points on a trajectory of the probability flow, and can be directly trained efficiently via the tangent condition.
potential motivates the central question: is there a principled methodology for training flow maps, and how can we do so efficiently in practice? In this work,
We introduce a direct training framework for flow maps, eliminating the need for pre-trained teacher models while maintaining the training stability of distillation.
Recently, there have been broad efforts to learn the flow map either directly or through distillation of a pre-trained model (Boffi al., 2024; Frans et al., 2024; Zhou et al., 2025; Salimans et al., 2024). Distillation-basee approaches perform well empirically, but require two-phase learning setup in which the performance of the student is limited by the performance of the teacher. In these methods, the practitioner first learns score (Song and Ermon, 2020; Ho et al., 2020) or flow (Lipman et al., 2022; Albergo and Vanden-Fijnden, 2022; Albergo et al., 2023; Liu et al., 2022) model, and then converts it into flow map via secondary training algorithm that considers the pre-trained model as a "teacher" for the "student" flow map. To avoid this complication, we aim to design learning schemes in which the flow map can be trained similarly to standard flow matching. In this endeavor, the fundamental challenge is a lack of a unified mathematical characterization that reveals how to learn the flow map efficiently, which has led to complex pipelines that require extensive engincering to overcome unstable optimization dynamics outside of the distillation setting (Lu and Song, 2025).
To address this challenge, we introduce a mathematical framework that exposes a landscape of novel training schemes. Our key insight is simple relation, the tangent condition (Figure 1), that explicitly relates the velocity of the probability flow equation to the derivative of the flow map. Using this insight, we develop self-distillation framework where the flow map is learned by simultaneously training and distilling its implicit velocity. The result is simple pipeline that leverages off-the-shelff training procedures for flows to learn a model with accelerated inference. Our framework reveals the fundamental design principles for learning flow maps, enabling practitioners to build few-step generative models as systematically as standard flows. Our main contributions are:
1. Algorithmic framework. We provide three equivalent mathematical characterizations of the flow map, showing how consistency models and other recent few-step methods including consistency trajectory models (Kim al., 2024), shortcut models (Frans et al., 2024), mean flow (Geng et al., 2025). and align your flow (Sabour et al., 2025) emerge as special cases of our methodology.
2. Self-distillation algorithms. Leveraging our description of the flow map, we introduce three new algorithmic families Eulerian (ESD), Lagrangian (LSD), and Progressive Self-Distillation (PSD) and discuss their connections to existing direct training schemes. We prove that each has the correct unique minimizer, and provide guarantees that the loss values bound the 2- Wasserstein error of the learned one-step model for ESD and LSD.
3. Empirical analysis. We study the performance of each method as a function of the number of spatial and time derivatives that appear in the objective function. We find that LSD, which avoids both spatial derivatives and self-consistent bootstrapping from smaller steps, attains the best performance across standard benchmarks including the synthetic checkerboard dataset, CIFAR-10, CelebA-64, and AFHQ-64.

---

2 Theoretical framework
In this work, we study the flow map of the probability flow equation, which is a function that jumps between points along a trajectory (Figure 1). Given access to the flow map, samples can be generated in single step by jumping directly to the endpoint, or can be generated with an adaptive amount of computation at inference time by taking multiple steps. Below, we give detailed mathematical description of the flow map, which we leverage to design a suite of novel training schemes. We begin with review of stochastic interpolants, which we use to build efficient flow-based generative models.
2.1 Stochastic interpolants and probability flows
Let {= with each æf Rd x pi denote a dataset drawn from a target density P1. Given D. our goal is to draw a fresh sample ž و_ from a distribution أ1 ~ learned to approximate the target. Recent methods for accomplishing this task leverage flows, which dynamically evolve samples from simple base distribution Pu such as Gaussian until they resemble samples from P1.
To build a flow-based generative model, we leverage the stochastic interpolant framework (Albergo et al., 2023), which we now briefly recall. We define a stochastic interpolant as stochastic process : [0, 1]
Interpolants.
RRd x de Hd that combines samples from the target and the base,
(x x1) 0, + B
where a, 3: [0, [ |0, 1] are continuously differentiable functions satisfying the boundary condi- tions 0kg = 1, a1 = 0, 30 0, and 31 1. In (1), the pair (Io. x1) ~ p(xo, 21) is drawn from cou- pling satisfying the marginal constraints Jo= p( 21) drz, = P1 (x1) ) and JRa P(Zg, I1 )dz = Po (2o). By construction, the probability density fr = Law( defines a path in the space measures between the base and the target. This path specifies probability flow that pushes samples from po onto PI.
z = (zz),
To ~ Po.
(2)
which has the same distribution as the interpolant, Hg Pt for all E 1]. The drift in (2) is given by the conditional expectation of the time derivative the interpolant, be(x) E[i,/4, = z). which averages the "velocity" of all interpolant paths that cross the point at time t. A standard choice of coefficients is a, = 1 and 3. t (Albergo and Eijnenn 2022: Albergo al., 2023), which recovers flow matching (Lipman et al., 2022) and rectified flow (Liu et al., 2022). Many other options have been considered in the literature, and in addition to flow matching, preservin and variance-exploding diffusions can be obtained as particular cases.
Learning. By standard results probability theory and statistics, the conditional expectation b, can be learned efficiently in practice by solving a square loss regression problem.
b argmin (b),
(b)
Ezn.e. [ic((To) iil3]dt.
(3)
b
Above, Ezg,21 denotes an expectation over the random draws of (ro, in the interpolant (1).
Sampling. Given an estimate i obtained by minimizing (3) over a class of neural networks, we can generate an approximate sample d by numerically integrating the learned probability flow i, = (āx) until time t from an initial condition ag ~ Po. This approach yields high-quality samples from complex data distributions in practice, but is computationally expensive due to the need to repeatedly evaluate the learned model during integration; here, we aim to avoid this solve.
2.2 Characterizing the flow map
The flow map X : [0, 1]2 x Rd + Rd is the unique map satisfying the jump condition
X (x,)) - wt for all (s, E [0, 1|2
(4)
where (zi (e)))c(o.1] is any solution of (2). The condition (4) means that the flow map takes "steps" of arbitrary size t E-s along trajectories of the probability flow. In particular, single application Xo, (zo)
3

---

with Po yields a sample from Pt, avoiding numerical integration entirely. Moreover, we may also increase the number of steps by composing X over a grid to < t the presence of model errors, which enables us to trade inference-time compute for sample quality.
t in
In what follows, we give three characterizations of the flow map that each lead an objective for its estimation. As we now show, these characterizations are based on a simple but key result that shows we can deduce the corresponding velocity field be from given flow map X.,t.
Lemma 2.1 (Tangent condition). Let X,, denote the flow map. Then,
lim (x)) bx(x) 8-
Vt E [O,
Va E IR",
(5)
i.c. the tangent vectors to the curve (z))(x))) give the velocity field be(x)) for every a.
As illustrated in Figure IA, Lemma 2.1 highlights that there a velocity model "implicit" flow map. To leverage this algorithmically. we propose to adopt an Euler step-like parameterization that takes into account the boundary condition 5,, x,
X (x) = x + (t (2)..()).
(6)
In (6), : [0, 1]2 x Rd + Rd is the function we will estimate parametrically. Despite its similarity to first-order Taylor expansion, the representation (6) corresponds to shift and rescaling of X and hence without loss of expressivity. In addition enforcing that Xx.t recovers the identity on the diagonal 8 = I, (6) implies that limg-se O,X (v,)(z), which gives an elegant connection between V,,t and the drift field by.
Ur, (x) be(x),
Vt E [O, IJ,
Vo E
Geometrically, es describes the "slope" of the line drawn between and a single ODE trajectory (Figure IB). The condition (7) states that the slope between two infinitesimally-spaces points is precisely the velocity b. A key insight is that this relation indicates Ur can be estimated using the objective (3). To learn the map X... then remains to estimate Us.t away from the diagonal = t. To this end, we leverage the following result, which relates Ux U1 for +
Proposition 2.2 (Flow map). Assume that X., is given by (6) with Us satisfying (7), and assume that u. is continuous both time arguments. Then, Xs, the flow map defined in (4) if and only if any of the following conditions also holds:
(i) (Lagrangian condition): X,,t solves the Lagrangian equation
(xX,,(z)) (x).
(8)
for all (s,t) E [0, 1]2 and for all x E Rdd.
(ii) (Eulerian condition): Sxt. solves the Eulerian equation
(xX,,(()) VX, (x)),())))))
(9)
for all (s,t) E [0, 1]2 and for all x E E.d.
(iii) (Semigroup condition): For all (s,t,u) E [0, and for all x E Rd
(x,,(Xx, = (x))
(10)
The Lagrangian and Eulerian conditions in Proposition 2.2 categorize the flow map X as the solution of an infinite system of ODEs or as the solution of a PDE, each of which describes transport along trajectories of the flow (2). The semigroup condition states that any two jumps can be replaced by a single jump. Sections B to D provide a review of the flow map matching framework (Boffi et al., 2024), and describe how these three characterizations are the basis for consistency (Song et al., 2023; Kim et al., 2024; Geng et al., 2024) and progressive distillation (Salimans and Ho, 2022a) schemes that have appeared in the literature. In the following, we show how each and in fact how any distillation method that produces a flow map from a velocity field b can be immediately converted into a direct training objective for single network model X via the concept of self-distillation.

---

2.3 A framework for self-distillation
Our framework augments training E1,t on the diagonal 8 via the objective (3) and the identity (7)
with penalization term for one or more of the conditions in Proposition 2.2 along the off- oifgonaa * t. This leads a of objectives that can each be used learn the flow map.
Proposition 2.3 -Self-distillationn.. The flow map xx defined in (4) is given for all 0 & < by (x) (t ())),(()) where Us, (x,)(x) the unique minimizer over i of
LSD(ū) = Co(e) + Lo(e),
(11)
where Co( given by
Cz( E (1e) 111)dt
(12)
and where LD( is any of the following three objectives.
(i) The Lagrangian self sef-distillation (LSD) objective, which leverages (8),
LLDD (0) = -S,- [ Ee.... ||0.8., éi (x., dsdt;
(13)
(it) The Eulerian self-distillation (ESD) objective, which leverages (9),
LESD -f00 Eao..n [1,, + V.g,, 9 dsdd;
(14)
(iii) The progressive self-distillation (PSD) objective, which leverages (10),
LPSD (i) = - J/'''''
10 Brn.
(1) - (Is) [2] dudsdt.
(15)
Above, (x((1)= = x (t - ,)0,,((x) and Ex,.,, denotes an expectation over the random draws of (zo, 21) in the interpolant defined in (I).
The proof follows directly from Proposition 2.2 and is given in Section E; the resulting algorithmic approach summarized graphically in Figure 2. In Proposi- tion 2.3, we focus on & i, which is all that is re- quired to generate data. Training over the entire unit square (s, t) E [0, 1J2 enables jumping backwards from data noise along trajectories of (2) in addi- tion to standard generation from noise to data. The derivatives with respect to space and time required to implement the LSD and ESD losses can be com- puted efficiently via standard jvp implementations. As we will see in our experiments, an advantage of the schemes (13) and (15) they naturally avoid derivatives with respect to space, which leads to sig- nificantly improved training stability.
LD(0)
t
(6
S
(B) + CD
Figure 2: Self-distillation. Our plug-and-plas approach pairs any distillation objective Lo the off off-diagonal s # of the square (0, 1]a with flow matching objective Ls on the diagonal 8 = to obtain direct training algorithm for the flow map.
We now provide theoretical guarantees that the ob- jective value bounds the accuracy of the model for LSD and ESD. We were unable obtain similar guarantee for PSD due to issues of compounding errors and distribution shift associated with bootstrapping small steps to large steps. which we believe be a fundamental difficulty the algorithm's 's construction. This difficulty is consistent with the observed reduced performance of PSD in comparison to L.SD in our experiments (Section 5).
Proposition 2.4 (Wasserstein bounds). Let Xx,( = + (t- ()),,(),(()) denote candidate flow map, let وı = Eoo denote the corresponding one-step generated distribution, and let É denote the spatial Lipschitz constant of (-,(()) uniformly in t. First assume Lu(ō) + (i) & Then,
W (p1, P1) & 4212
(16)
Now assume that Lo( + (0) < Then,
W (1), 2e (1 + e 2i))s.
(17)
5

---

Algorithm 1: Learning flow maps via self-distillation
input: Dataset D; interpolant coefficients ot, Be; batch size M; diagonal fraction yf distillation method Lo E {LLSD: LESD, LPSD)-
repeat
Sample Ma LnM] pairs (1 æf) p( æ1) and times t; ~ ([0, I]);
Compute interpolants 4t at, zy + St, xi and velocities Lt. dq, Jb + Be
Compute diagonal loss: Lb 1 Ma E
e 1 (Fe,) ) - 121 + wi.i
Ma
Sample M. M - Ma pairs (zg. x() p(2o, and times (85, 44) ~ Undi
Compute interpolants Is; as, + Bs,
Compute distillation loss: LD = M. /==1
Ma ebyyi 20 (i) Wsysts?
Compute LSD Cb + LD:
Update and w using L5D: until converged:
output: Trained flow map Xx,(xx) = x + (1-,,,,,r(o)
The above result highlights that the model's accuracy improves systematically the loss is min- imized for both ESD and LSD. The proof follows by combining guarantees for distillation based algorithms (Boffi et al., 2024) with guarantees for flow-based algorithms (Albergo et al., 2023), which can be stitched together by our assumption on the value of the loss.
3 Algorithmic aspects
We now provide practical numerical recommendations for an implementation of self-distillation. Our aim is not to provide single best method, but devise general-purpose framework that can be used to build high-perfor high-performing flow maps across data modalities. We provide general algorithmic prescription in Algorithm with specific instantiations for LSD, ESD, and PSD in Section F.6.
Choice of teacher. The ceff-distillation objectives in Proposition 2.3 are obtained by squaring the residuals of the properties in Proposition 2.2. While the minimizers correct, the associated training dynamics may not be optimal because the losses are nonconvex in ô. The flow of information should be from the diagonal ti,,, where there is an external learning signal via 1,, to the off- off-diagonal 0g which bootstraps the signal in ,,,.- To enforce that ú..t learns match Ür,,, rather than vice-versa, use the stopgrad operator to match the distillation setting, where the off-diagonal would adapt entirely external teacher. Detailed descriptions of the recommended placement are given in Section F.4
Relation to existing methods. The generic framework described by Proposition 2.3 and Algorithm recovers most existing schemes for training consistency models and their extensions. In particular, by proper choice of the distillation objective and the teacher, we can obtain standard training for consistency models (Song et al., 2023), consistency trajectory models (Kim et al., 2024), and shortcut models (Frans et al., 2024). These connections are given in detail Sections and D.
Loss weighting, The loss (11) can be written explicitly as an integral over the upper triangle s < t,
(S)(2) = ['' I'" (c(())((s - t) + cg (6)) dadt, (Tt)) where ()) = E,,,,,//e,.
(18)
21 denotes (12) restricted to f, and where Lo denotes the distillation term restricted (s, t). We find that loss values at different pairs t).t) can have gradient norms that differ significantly, introducing undesirable variance. To rectify this, we incorporate a learned weight Ws.., generalizing the EDM2 weight (Karras al., 2024) to the two-time setting,
(0)( [() (e---- (C)((())))) - t) ())( + w..e) dsdt.
(19)
In (19), Wast can be interpreted as an estimate of the log-variance of the loss values; at the global minimizer, ensures that all values of (s, I) contribute on a similar scale. We find that using ws,t significantly stabilizes the training dynamics and enables the use of larger learning rates.

---

Temporal sampling. In addition to the weight, we introduce sampling distribution Pa.t.
(i) Ep.. [e [e- Wo.. (C)(0))(s - t) (6)( + we.e]-
(20)
While the weight w,, normalizes the variance across times, Pe,c chooses how select times randomly for each batch. Let Ud denote the uniform distribution on the diagonal s = t, and let Uod denote the uniform distribution
the upper triangle s < t. In our experiments, we leverage the mixture distribution ps,t = r + (1 7) which places a fraction v, of the batch uniformly at random on the diagonal and fraction (1 77) uniformly random in the upper triangle. Because our distillation losses reduce to the flow matching loss in the limit as 8 - t (Section F.2), we only use La on the diagonal the majority of the computational effort towards learning the flow and proportionally less towards distilling it. The distillation objectives in Proposition 2.3 are more expensive than the interpolant objective La because they require multiple evaluations of the network and Jacobian-vector products. As result, 7) can also be used to tune the cost of each training step (Section F.3).
f. We found 77 0.75 to work well in early experiments, which puts
PSD sampling. For PSD, we introduce a proposal distribution Pu over the intermediate step, <P50 (0) = Epp. Exo.. [X (Ts) (X.
(21)
We parameterize u = 78 + (1 - 7 as a convex combination for 7 E [0, 1] and define the proposal distribution by sampling over %. In our experiments, compare uniform sampling (PSD-U) with y ~ U(0, 1]) to midpoint sampling where y ~ 61,2 so that = 1/2 deterministically (PSD-M).
PSD scaling.
We show in Section F that (21) may be rewritten entirely in terms of é as
Cpso (0) (t) Ep, Ezo. (T.) (1 - . (To) - yuu, (
(22)
The form (22) eliminates factors u - 8, 1 - 8, and 1 - u. that appear due to the parameterization (6). We found these terms introduced higher gradient variance because they cause the loss to scale like (t - s)², which changes the effective learning rate depending on the timestep (t (s=s); s); we drop this
factor of (t s)² in practice. (22) preconditions the loss and removes this additional source of variability, leading to improved training stability.
Conditioning and guidance. Flow maps be made conditional by incorporating conditioning argument c as x (z; c) = x + (t (t-& ,),,,((2; c). We can use this observation to define classifier-frs free guided (CFG) flow maps, as we now show (Ho and Salimans, 2022). To do so, let Ø correspond to unconditional generation, and let let az c) = by Ø) a (1; - bx 2)) be the CFG velocity field at guidance strength a. This velocity has flow map X (x; (x; c) satisfying Ve, (x; o) = (2; a,, c), which may be learned via self-distillation by incorporating additional random sampling over the guidance scale (Section F.5). In this work, we focus solely on unconditional, unguided generation and leave the usage of guidance for future study.
+ (t s)v,,, ar, e)
Multiple models and general representations. In Proposition 2.3, we leverage single model Xs, defined in terms of uz... While this leads to greater efficiency through (7), it requires one model to learn both the velocity and its flow map X, which may require more network capacity than traditional flow-based generative models. Instead, it is also possible to use two models one parameterizing b and one parameterizing X which can be trained simultaneously. Higher-order parameterizations can be designed that leverage both, such as Xx.(x)- = x + (t - ()) + { s) (z), which can use frozen pre-trained model b, or can train à from scratch tandem with v.. More generally, any parameterization Xz,t satisfying X (x) x may be used in practice, where in this setting we use lim,-, 0 0,,,,. (zr) in place of ū,,, in Algorithm This may be computed via automatic differentiation jvp in f. at s f. In this work, we focus the representation (6), which requires only a single model and gives a computationally efficient way to evaluate La; we leave these more general and oigeer-order parameterizations to future work.
4 Related work
Flow matching and diffusion models. Our approach builds directly on methods from flow match- ing and stochastic interpolants (Lipman et al., 2022; Albergo and Vanden-Eijnden, 2022; Albergo
7

---

Step Count
Dataset
Method
2
8
16
LSD
0.086
0.077
0.071
0.070
0.071
Checker (KL !)
ESD
0.098
0.092
0.083
0.082
0.075
PSD-M
0.146
0.089
0.081
0.072
0.069
PSD-U
0.111
0.107
0.075
0.073
0.068
LSD
8.100
4.370
3.340
3.330
3.570
CIFAR-10 (FID ()
PSD-M
12.810
8,430
5.960
5.070
4.640
PSD-U
13.610
7.950
6.030
5.320
5.160
LSD
12.220
5.740
3.180
2.180
1.960
CelebA-64 (FID L)
PSD-M
19.640
11.750
7.890
6.060
5.090
PSD-U
18.810
11.020
7.470
6.000
5.630
LSD
11.190
7.780
7.000
5.890
5.610
AFHQ-64 (FID !)
PSD-M
18.860
14.750
14.400
13.260
11.070
PSD-U
14.500
10.730
10.990
12.020
11.470
Table I: Benchmark results. Performance sampling step counts for the low-dimensional checker dataset (KL divergence) and natural image datasets (FID). Best method per dataset and step count shown bold.
et al., 2023; Liu et al., 2022) as well as the probability flow equation associated with diffusion models (Song al., 2020; Ho et al., 2020; Maoutsa et al., 2020; Boffi and Vanden-Eijnden, 2023). These methods define an ordinary differential equation whose solution evaluates the flow map at single time. Due to the computational expense associated with solving these equations, line of recent work asks how to resolve the flow more efficiently with higher-order numerical solvers (Dockhorn et al., 2022; Lu et al., 2022; Karras et al., 2022; Li et al., 2024) and parallel sampling schemes (Chen al., 2024; Bortoli et al., 2025). Our approach instead estimates the flow map to enable accelerated sampling by avoiding the differential equation solve altogether.
Consistency models. Appearing under several names, the flow map has become central object of study recent efforts to obtain accelerated inference. Consistency models (Song et al., 2023; Song and Dhariwal, 2023) estimate the single-time flow map to jump from any time to data, given by X in our notation. Consistency trajectory models (Kim al., 2024; Li and He, 2025; Luo et al., 2023) estimate the two-time flow map, enabling multistep sampling. Both approaches implicitly leverage the Eulerian characterization (9), which we find leads to gradient instability, explaining recent engineering efforts for stable training (Lu and Song, 2024). Progressive distillation (Salimans and Ho, 2022a) uses the semigroup condition (10) to train a model that can recursively replicate steps of pre-trained teacher. Progressive flow map matching (Boffi al., 2024) enforces this iteratively over flow map after pre-training, while shortcut models apply a discretized semigroup condition (Frans al., 2024). In concurrent work, Geng et al. (2025); Sabour al. (2025) introduce distillation and direct training schemes that reduce to a particular case of our Eulerian formulation. Details on these methods and their connection to our framework may be found Sections D.
5 Numerical experiments
We test LSD, ESD. PSD-U, and PSD-M on the low-dimensional checkerboard dataset, as well as in the high-dimensional setting of unconditional image generation on CIFAR- 10, CelebA-64, and AFHQ-64. In each case, we study performance at fixed training time obtain a fair comparison. We emphasize that our aim is not to obtain state of the art performance, but to understand the trade-offs of each approach and compare them on an equal footing; with further engineering, quantitative metrics could be lowered significantly for all methods. For image datasets, we find ESD to be unstable due to spatial gradient, leading to poor performance without gradient stabilization schemes. We find that LSD obtains uniformly the best performance on all problems tried. This is consistent with our theoretical results Proposition 2.4, where we were able obtain stronger theoretical guarantees for LSD than for PSD. Full network and training details are provided in Section G and Table 2.

---

LSD
ESD
PSD-M
PSD-U
Figure 3: Checker dataset. Qualitative results for the two-dinensional checker dataset. LSD performs the best across all step counts except 16 (Table 1). methods improve as number of steps increase. ESD and both PSD variants fail to capture the sharp boundaries at small N, introducing artifacts and driving KL higher.
Checkerboard. While synthetic, the checkerboard dataset exhibits multimodality, sharp boundaries, and low-dimensionality that make it useful testbed for exact visualization of how few-step samplers capture complex features in the target. Qualitative results shown in Figure 3, while quantitative results obtained by estimating the KI. divergence (for details, again see Section G) between generated samples and the target are shown in Table 1. L.D performs best across all sampling steps tried except for N = 16, where all methods perform well. The performance of LSD also saturates around N = 4 sampling steps. By contrast, ESD, PSD-U, and PSD-M all see increased performance up to steps with reduced performance for fewer steps. The qualitative results in Figures 3 and 6 highlight that the higher KL values result from failure to capture the sharp features present in the dataset, with ESD blurring the boundaries and the PSD methods introducing artifacts that connect the modes.
16
CIFAR-10. Figure 4, we study the parameter gradient norm as function of the training iteration on CIFAR-10. LSD and PSD, which avoid computing spatial derivatives of the network during training, maintain significantly more stable gradients than ESD even when using sg (.). We found the high gradient norm of ESD to induce training instability, ultimately leading to divergence. This consistent with earlier work on consistency models, where careful annealing schedules, clipping, and network design has been necessary to stabilize continuous-time training (Lu and Song, 2025).
We track the quantitative performance of each method as measured by FID in Table I; we do not report FID values for ESD due to training instability. We find that LSD ob- tains the best performance across all step counts followed by PSD. PSD-U and PSD-M trade places depending on step count. A qualitative visualization of sample quality shown in Figure 5 (Top) as function of the number of sampling steps. We see that each method obtains im- proved quality as the number of steps increases, and that all methods produce similar images for fixed seed.
103
Gredient Norm
10a
10
100
100000
200000
CelebA-64. As shown in Table 1, LSD also obtains the best performance across all step counts on CelebA-64 (Liu et al., 2015), with FID scores ranging from 12.22 at N = to 1.96 at N 16. The gap between L.SD and the PSD variants is more pronounced on CelebA-64 than on CIFAR-10, particularly for low step counts. PSD-U mostly outperforms PSD-M, with PSD-M only obtaining higher-
Training Step
Figure 4: CIFAR-10: Parameter gradi- ent norms. Spatial and temporal representa- tions
flow map impact parameter gra- dient norms of self-distillation methods that require network time and space derivatives.

---

LSD
PSD-M
PSD-U
N=1
2=
N=4
N
N=1
2=
N=4
N
N=16
N=1
2=
N=4
N=8
N=16
Figure 5: Progressive refinement. Sample quality as function sampling steps using same eight fixed noise samples across all methods for fair comparison. (Top) CIFAR-10, 10. (Middle) CelebA-64, (Bottom) AFHQ-64. LSD consistently produces coherent samples across all datasets and step counts.
performing 16-step map. A qualitative visualization is shown in Figure 5 (Middle). All methods show systematic improvement as the step count increases, with faces becoming sharper more detailed.
AFHQ-64. Finally, we evaluate on AFHQ-64, a more challenging dataset with greater visual diversity than CelebA-64 that includes variation across animal categories (Choi et al., 2020). As shown Table 1, LSD again achieves the best FID scores across all step counts, ranging from 11.19 at N to 5.61 N 16. PSD shows notably higher FID scores this dataset, particularly PSD-M, which struggles at low step counts. PSD-U again mostly outperforms PSD-M, with PSD-M obtaining slightly higher- pethhrrnrforninggg 16-step map but worse performance otherwise. Qualitative results are shown Figure (Bottom), where we again see that higher step counts lead to generated images with increasing levels of detail.
Conclusion
In this work, we expose and investigate the design space of a class of flow-based generative models with accelerated inference known as flow maps. These models generalize and extend consistency models to include multiple training paradigms and principled multistep inference. Rather than learning the velocity field typical of flows and diffusions, flow maps learn the solution operator the probability flow equation, obviating the need to solve differential equation for inference. We show that learning can be performed directly by pairing flow training with any of three characterizations of the flow map, approach we refer to self alf-distillation. Self-distillation can be incorporated with minimal additional overhead, making flow maps an appealing new paradigm. While we systematically categorize the design space of flow map models, the main limitation of our contribution is that we were unable to systematically test each component empirically due to the large associated computational expense. Critical aspects deserving further experimentation include ablations over the flow map parameterization and architecture; stabilization, annealing, and stopgradient schemes for training; and hybrid approaches that combine multiple of our self-distillation objectives.
10

---

Acknowledgments
MSA is supported by Junior Fellowship at the Harvard Society of Fellows well as the National Science Foundation under Cooperative Agreement PHY-2019786 (The NSF AI Institute for Artificial Intelligence and Fundamental Interactions, http://aaifi.rr NMB would like to thank Max Sim- chowitz, Andrej Risteski, Stephen Huan, Jerry Huang, Chaoyi Pan, Giri Anantharaman, and Gabe Guo for helpful conversations.
References
Nicholas M. Boffi, Michael S. Albergo, and Eric Vanden-Eijnden, Flow Map Matching: unifying framework for consistency models. arXiv:2406.07550 June 2024.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. MiXiv:2011.13456 2020.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pages 10684-10695, 2022.
Nanye Ma, Mark Goldstein, Michael Albergo, Nicholas M. Boffi, Eric Vanden- Eajden-Eijnden, and Saining Xie. SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers. Ttt.::20010087444 2024.
Adam Polyak, Amit Zohar, Andrew Brown, Andros Tjandra, Animesh Sinha, Ann Lee, Apoorv Vyas, Bowen Shi, Chih- Ma. Ching- Yao Chuang. David Yan, Dhruv Choudhary, Dingkang Wang, Geet Sethi, Guan Pang, Haoyu Ma, Ishan Misra, Ji Hou, Jialiang Wang, Kiran Jagadeesh, Kunpeng Li, Luxin Zhang, Mannat Singh, Mary Williamson, Matt Le, Matthew Yu, Mitesh Kumar Singh, Peizhao Zhang, Peter Vajda, Quentin Duval, Rohit Girdhar, Roshan Sumbaly, Sai Saketh Ramb- hatla, Sam Tsai, Samaneh Azadi, Samyak Datta, Sanyuan Chen, Sean Bell, Sharadh Ramaswamy, Shelly Sheynin, Siddharth Bhattacharya, Simran Motwani, Tao Xu, Tianhe Li, Tingbo Hou, Wei- Ning Hsu, Xi Yin, Xiaoliang Dai, Yaniv Taigman, Yaqiao Luo, Yen-Cheng Liu, Yi-Chiao Wu, Yue Zhao, Yuval Kirstain, Zecheng He, Zijian He, Albert Pumarola, Ali Thabet, Artsiom Sanakoyeu, Arun Mallya, Baishan Guo, Boris Araya, Breena Kerr, Carleigh Wood, Ce Liu, Cen Peng, Dimitry Vengertsev, Edgar Schonfeld. Elliot Blanchard, Felix luefei-Xu, Fraylie Nord, Jeff Liang, John Hoffman, Jonas Kohler, Kaolin Fire, Karthik Sivakumar, Lawrence Chen, Licheng Yu, Luya Gao, Markos Georgopoulos, Rashel Moritz, Sara K. Sampson, Shikai Li, Simone Parmeggiani, Steve Fine, Tara Fowler, Vladan Petrovic, and Yuming Du. Movie gen: cast of media foundation models. wrXiv:2410.13720 2025.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution. vrXiv:2310.16823 June 2024.
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J. Ballard, Joshua Bambrick, Sebastian W. Bodenstein, David A. Eyans, Chia-Chun Hung. Michael O' O'Neill, David Reiman, Kathryn Tunyasuvunakool, Zachary Wu, Akvilé Zemgulyté, Eirini Arvaniti, Charles Beattie, Ottavia Bertolli, Alex Bridgland, Alexey Cherepanov, Miles Congreve, Alexander I. Cowen-Rivers, Andrew Cowie, Michael Fig- urnov, Fabian B. Fuchs, Hannah Gladman, Rishub Jain, Yousuf A. Khan, Caroline M. R. Low, Kuba Perlin, Anna Potapenko, Pascal Savy, Sukhdeep Singh, Adrian Stecula, Ashok Thillaisun- daram, Catherine Tong, Sergei Yakneen, Ellen D. Zhong, Michal Zielinski, Augustin Zidek, Victor Bapst, Pushmeet Kohli, Max Jaderberg, Demis Hassabis, and John M. Jumper. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 630((016):493-55)) 5000 2024.
Ian Price, Alvaro Sanchez-
Ferran Alet, Tom Andersson, Andrew El-Kadi, Dominic Masters, Timo Ewalds, Jacklynn Stott, Shakir Mohamed, Peter Battaglia, Remi Lam, and Matthew Willson. GenCast: Diffusion-based ensemble forecasting for medium-range weather, May 2024.
Claudio Zeni, Robert Pinsler, Daniel Zügner, Andrew Fowler, Matthew Horton, Xiang Fu, Zilong Wang, Aliaksandra Shysheya, Jonathan Crabbé, Shoko Ueda, Roberto Sordillo, Lixin Sun, Jake

---

Smith, Bichlien Nguyen, Hannes Schulz, Sarah Lewis, Chin- Wei Huang, Ziheng Lu, Yichi Zhou, Han Yang, Hongxia Hao, Jielan Li, Chunlei Yang, Wenjie Li, Ryota Tomioka, and Tian Xie. A generative model for inorganic materials design. Nature, 639(8055))62446333 632, 2025.
Kevin Black, Noah Brown, Danny Driess, Adnan Esmail, Michael Equi, Chelsea Finn, Niccolo Fusai, Lachy Groom, Karol Hausman, Brian Ichter, Szymon Jakubczak, Tim Jones, Liyiming Ke, Sergey Levine, Adrian Li-Bell, Mohith Mothukuri, Suraj Nair, Karl Pertsch, Lucy Xiaoyang Shi, James Tanner, Quan Vuong, Anna Walling, Haohuan Wang, and Ury Zhilinsky. un: A Dision-Languuge--e 2024.
Flow Model for General Robot Control. arXiv:2410.24164, October
Cheng Chi, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. Diffusion Policy: Visuomotor Policy Learning via Action Diffu- sion. arXiv:2303.04131 2024.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ila Sutskever. arXiv:2303.01469, 2023.
Consistency Models.
Dongjun Kim, Chich-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion. arXiv:2310.02279, 2024.
Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel. One step diffusion via shortcut models. arXiv:2410.12557. 2024.
Linqi Zhou, Stefano Ermon, and Jiaming Song. Inductive moment matching. arXiv:2503.07565, 2025.
Tim Salimans, Thomas Mensink, Jonathan Heek, and Emiel Hoogeboom. Multistep Distillation of Diffusion Models via Moment Matching. arXiv:2406.04100 June 2024.
Yang Song and Stefano Ermon, Generative Modeling by Estimating Gradients of the Data Distribution. arXiv: 0.060000 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.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference Learning Representations, 2022.
Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference Learning Representations, 2022.
Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.0879) 2023.
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 Learning Representations, 2022.
Cheng Lu and Yang Song. Simplifying, stabilizing and scaling continuous-time consistency models. MwXiv:2410.11081, 2025.
Zhengyang Geng, Mingyang Deng, Xingjian Bai, J. Zico Kolter, and Kaiming He. Mean Flows for One-step Generative Modeling. arXiv:2505.13447, May 2025.
Amirmojtaba Sabour, Sanja Fidler, and Karsten Kreis. Align Your Flow: Scaling Continuous- Flow Map Distillation. arXiv:2506. June 2025.
Zhengyang Geng, Ashwini Pokle, William Luo, Justin Lin, and J. Zico Kolter. Consistency Models Made Easy. arXiv:2406. 2e0e24466..444. 14548, 2024.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arX1v:2202.00512, 2022a.
12

---

Tero Karras, Mikk Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyzing and Improving the Training Dynamics of Diffusion Models. arXiv:2312.02696, March 2024.
Jonathan Ho and Tim Salimans. Classifier-Free Diffusion Guidance. arXiv:2207.12559 July 2022.
Dimitra Maoutsa, Sebastian Reich, and Manfred Opper. Interacting particle solutions of Fokker- Planck equations through -radestt-log-density estimation. Entropy, 22(8):802, July 2020.
Nicholas M. Boffi and Eric Vanden-Fijndenn Probability flow solution of the Fokker-Planck equation. Machine Learning: Science and Technology, 4(3):035012, July 2023.
Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Genie: Higher-order denoising diffusion solvers. arXiv:2210.05475. 2022.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: fast ode solver for diffusion probabilistic model sampling in around 10 steps. arXiv:2206.0000 2022.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. arXiv:2206.00366 2022.
Gen Li, Yu Huang, Timofey Efimov, Yuting Wei, Yuejie Chi, and Yuxin Chen. Accelerating convergence of score-based diffusion models, provably. arXiv:2403.03852, 2024.
Haoxuan Chen, Yinuo Ren, Lexing Ying, and Grant M. Rotskoff. Accelerating diffuusion models with parallel sampling: Inference at sub-linear time complexity. arXiv:2405.159%0 2024.
Valentin De Bortoli, Alexandre Galashov, Arthur Gretton, and Amaud Doucet. Accelerated diffusion
models via speculative sampling. arXiv:2501.0533 2025.
Yang Song and Prafulla Dhariwal. Improved Techniques for Training Consistency Models. arXiv:2310. 14199, 2023.
Liangchen Li and Jiajun He. Bidirectional consistency models. arXiv:2403./8033 2025.
Simian Luo, Yiqin Tan, Longbo Huang. Jian Li, and Hang Zhao. Latent consistency models:
Synthesizing nigh-resolution images with few-step inference. arXiv:2310.04378, 2023.
Cheng Lu and Yang Song. Simplifying, Stabilizing and Scaling Continuous- Time Consistency
Models. 1xXx::2410.110/1 October 2024.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes the wild. In Proceedings of International Conference Computer Vision (ICCV) December 2015.
Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung-Won Ha. Stargan v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2020.
Ben Poole, Ajay Jain, Jonathan Barron, and Ben Mildenhall. DreamFusion: Text-to-3D using 2D Diffusion. arXiv:2209. 14988, September 2022.
Tim Salimans and Jonathan Ho. Progressive Distillation for Fast Sampling of Diffusion Models. 12X1v:2202.00552, 2022b.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Yes, we introduce a novel class of self-distillation algorithms for learning flow maps, and study their performance numerically.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Yes, we discuss how the main limitation of our work is lack of state of the results due limited computational capabilities, and that our conclusions can be architecture and dataset dependent. We plan to improve upon the quantitative values in the revision.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
14

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We include several theoretical results in the paper, each of which clearly states the assumptions and includes correct proof.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We include full experimental details in the main text and appendix, and we plan release our code, checkpoints, and configuration files.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
15

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: Yes, we plan to release a well-documented open source release with the camera-ready version. Included in this will be open-source jax implementation of the EDM2 neural network.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nipp cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //nips "tipp.cccpublliii. wtc/pullic/aao
wweeeeeeeoeeeaiiaeeeeeeeee ississss ssi aesionnPoicyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: All experimental details are provided in the main text with further details in the appendix, and will be included in the released configuration files.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: We did not report error bars for our quantitative metrics, but we used a very large batch for calculating the estimated KL divergence, did not observe any variability empirically. We reviewed the literature and found that it common to not report error bars for FID values, but we happy include them in the revision.
Guidelines:
The answer NA means that the paper does not include experiments.
16

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results are out range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provided the number and type of GPUs used to run the experiments in the main text.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /nturrrurs...// cc cctppblic//aeeecccccceeceeeaetet Caublic/Ethii htte//ttisssiiiiiieeeeeeeee csQuidelines!
Answer: [Yes]
Justification: Our work is primarily about the theoretical and empirical study of training moderate-scale generative models, and does not come with significant ethical implications outside of the usual caveats surrounding generative AI broadly speaking.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
17

---

As stated in the previous answer, outside of the usual caveats surrounding generative AI, we do believe there be significant ethical or broader impacts related to our work.
Justification:
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g.. pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
We do not believe that this paper poses such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
We do not use any existing assets outside of the EDM2 network, which is cited clearly.
Justification:
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
18

---

The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Answer: [Yes]
Justification: We have released an open-source implementation of our method that repro- duces all experimental results, which comes with documentation.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/mmodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our work does not involve crowdsourcing nor research with humans subjects.
Guidelines:
19

---

The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: Our work does not involved LLMs as a core component.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: / /tt///nerrpppp//e/ /ttarrips.cc/ooofeeeeeeeeeeeeeee cccc//oofeeeencce/222///LLL for what should or should not be described.
20

---

A Background on stochastic interpolants
For the reader's convenience, now recall how to construct probability flow equations using stochastic interpolants. We remark that score-based diffusion models can also be cast in the form of a stochastic interpolant (1), though they do not usually satisfy the exact boundary conditions f 0 and = 1, and the direction of time is opposite by convention. For example, the variance exploding and EDM processes (Karras al., 2022) naturally fit within this form as X Xu + 0g2, while the variance preserving process can be cast in this form after solving the Brnstein-Uhlenber process + /
dX;- = -xdtt + V2dW, in distribution X, d In both cases, we may define I = X to flip the direction of time over the horizon T.
with N(O, I).
The key property of the interpolant construction is that solutions to the probability flow (2) push forward their initial conditions onto samples from the target by matching the time-dependent density of (I), we now show.
Lemma A.1 (Transport equation). Let Pt. = Law(Le) be the density of the stochastic interpolant (1). Then Pr satisfies the transport equation
(.pe(a)) V (Nx()))(())) O,
Pt=o(x)) po(x))
(23)
where V denotes a gradient with respect to z, and where bo(x) Bii, 1 x].
Proof. The proof proceeds via the weak form of (23). Let Φ E C2 (RR) denote an arbitrary continu- ously differentiable and compactly supported test function. By definition,
J (()) E[o(T)]
Vt E [0, 1]
(24)
where 15 is given by (1). Taking the time derivative of this equality, we deduce that
[ (z))) E( VΦ(I)
E[b,(II) Vo(I))
(25)
f (z) Vo(z)p((z)dr.
_
The first line follows by the chain rule, the second by the tower property of the conditional expectation and the definition of the drift b,, and the third by definition of pt. The last line is the weak form of the transport equation (23).
nearly-identical derivation (simply dropping the tower property step) shows that the probability flow equation (2) satisfies Law( = Pt Law(()). Together, these results imply that we can sample from any density Pt solving a transport equation of the form (23) by solving the corresponding ordinary differential equation (2). In practice, we may implement this algorithmically by approximating b, with neural network via minimization of (3) to obtain model b,, and then solving the associated differential equation à (āte) from approximate samples from Pl.
with an initial condition do ~ Po to obtain
B Background on flow map matching.
As discussed Boffi et al. (2024), given a pre-trained velocity field 6, we may leverage the three properties in Proposition 2.2 design efficient distillation schemes by minimizing the corresponding square residual. In the following, we use the notation C(X; b) or X) to denote loss function for the flow map X given the teacher (which remains frozen during training).
Stopgradients. Because is pre-trained teacher, its parameters are frozen during training. The dittilistillation schemes we introduce in this work replace the teacher network b, by self-consistent implicit teacher i.... eliminating the need for the pre-trained model entirely. Inspired by the distilla- tion setting, Section F.4, we will use a stopgradient operator sg (-) in the context of self-distillation schemes to create similar effect to a frozen teacher and to control the flow of information within the model. Nevertheless, for training stability, it has been observed that it can useful to use additional sg (+) operators even for distillation, which we discuss after introducing each loss.
21

---

B.1 Lagrangian distillation.
The first approach the Lagrangian map distillation (LMD) algorithm, which is based on (8) and is the basis for the LSD algorithm,
(LM (''
Ep. [0..t, (1.) by (Xx,,((4))) dsdt.
(26)
The Lagrangiam scheme (26) was introduced Boffi et al. (2024), and to our knowledge has not appeared in other works. While b is frozen, the loss (26) is nonconvex in X' due to the nonlinearity of Moreover, computing the gradient of (26) with respect to X (or its parameters) requires computing the spatial Jacobian of b, which has been observed to be problematic for large generative models such as image synthesis systems (Poole et al., 2022). For these reasons, it is common to use the modified loss function
LMD ( (X;b) - _ E,,
- be -x [1x,,1(.....111 (l.) (Xsx:(.))) p3] dsdt.
(27)
The effectiveness of (27) over (26) depends the data modality and the neural network architecture, as the spatial Jacobian is only problematic in some contexts depending on the pre-trained teacher. We refer to the gradient of a loss function such (27) which includes the sg (-) operator semigradient.
B.2 Eulerian distillation.
A second scheme is the Eulerian map distillation (EMD) method based on (9),
LEMD (X;:) J/ Ep, TX.
(L.),,
dsdt.
(28)
Unlike the Lagrangian approach, (28) convex in X. Nevertheless, taking the gradient with respect to the parameters X requires backpropagating through its spatial Jacobian, which can be similarly problematic as the setting described for (26). One fix is to use semigradient based on
LEMD LEMD(X:))) (X;b) - f./. J. Ep, [Ia.X..(..) + sg ( (V.,,((.)(..((.)) |²] dsdl,
(29)
which avoids backpropagating through the spatial Jacobian entirely. While this helps training stability, it has been observed by Boffi al. (2024) that the Lagrangian schemes (26) and (27) are more stable than the Eulerian schemes (28) and (29), which is consistent with our experiments in Section 5.
B.3 Progressive flow map matching.
We now describe the progressive flow map matching (PFMM) algorithm, which is inspired by progressive distillation (Salimans and Ho, 2022b) for diffusion models, but adapted to the stochastic interpolant and two-time flow map setting. Let X,, denote pre-trained teacher flow map, assumed to be valid over the range 0 & & & To obtain such map at initialization, we may take + (t (x) with a pre-trained flow map b, corresponding to single Euler At. Our aim is "extend" X over a larger range, say 0 & & < 2r, by training a second flow map Xs.: to match two steps of X .. To do so, we consider the objective
At
and set (x.(x)) step of size (t s)
LPFMM (x;MM(X;;] (x; = 1.ah 2T 1'" J. Bpp.
[1xax.1.) [
-N., (Xa,a(Ts)) [2] dudsdl,
(30)
which is based on the semigroup property (10). In words, (30) teaches X replicate two jumps of X in one larger jump. We may also apply (30) self-consistently, where X itself serves the teacher,
LPFMM () (8)= / ". [ (1x..(...
sg (x..(((())))))
|e] dudsdi,
(31)
after the first round where is used, and extend over the course of optimization according to pre-defined annealing scheme. Our general self-distillation framework described in Section 2.3 may be obtained by using one of the above distillation schemes in tandem with direct training of é, and where we use è the teacher velocity field for the student flow map model X.
22

---

C Connection to consistency models.
The approaches (28) and (29) directly related consistency distillation in the continuous-time limit (Song and Dhariwal, 2023; Lu and Song, 2024). Consistency models estimate the single- time flow map from noise data, which in our notation given by Xs,1. Consistency trajectory models (Kim et al., 2024) use the same approach to learn the two-time map Xs,ti for agreement with the main text, we focus this setting here.
C.1 Consistency distillation and Align Your Flow.
We first take a continuous time limit of the discrete-time consistency distillation objective. Discrete- time consistency distillation considers the loss
(cD(X:b) L E [Re((Ts) - sg (X++se((f+tae)) Pl
(32)
1. 015 B,011,
3 As = 1, + Asd, (T.)-
In words, LCD aims to make X.,, "consistent" on trajectories of the teacher's probability flow ig by using shorter step of size (t As) teacher for slightly larger step of size (t (5-)) s). Taking
the gradient with respect X,,,, find
8CCD (8;b) (o,( (e+sor(Rotaa))
(33)
8X s.t
To obtain the gradient with respect to the parameters 0 of X, we have by the chain rule that VoLco V.X..s &Cm, so we focus on the functional derivative for notational simplicity. Taylor "X..t" expanding, we find that
Xs+de.((sstaa)
(4, Asb,
-X) + (TX + +(As),
(34)
= + A) d(As)
=xaa(Tss As 0, + V (T.),,(TT) o(As)
In the last line, used that AsV Sn, (s, AsV o(As). With this, we find
lim 1 5LCD (X;b) + +V...(.).6..
(35)
As 8X,
which is simply the negative Eulerian residual.
We now ask if the semigradient (35) be obtained from the Eulerian distillation objective (28) with a certain choice of sg (-). To do so, we consider the specific parameterization (6) given by X xx,(x) + (t - s)ú., (z). In this case, the Eulerian equation becomes a, 8, (x) + VX (x)
(36)
(x) + (t s)0,ú,,(z)) + -t (t (z)
As result, the Eulerian map distillation loss (28) becomes
LEMD b)
-/. (t) (t- (-
(37)
dsdt.
We consider variant that avoids backpropagating through any spatial temporal gradient
LEMD (0: i)
(38)
0 Jo Ep. [|- (o.((Lo) + sg (tt (t- (L)) (1.)+ (t- (t-s)VEaa) s))6,.((4)).(44) p2] dsdt.
23

---

This yields the semigradient,
8 LEMD
(0: b)
80g
=0,,((e) (t- (t-s)),,0.o11) (.()) vt
(39)
-0,8s,, (I,
In the last line, applied (36), which agrees with (35). Hence, the objective (38) is equivalent the consistency distillation objective in the continuous-timg limit after a suitable rescaling of gradients. We note that (39) is identical the "Align Your Flow" update considered by Sabour et al. (2025).
C.2 Consistency training and mean flow.
Consistency training aims to train a model directly, avoiding access to a pre-trained teacher (Song and Dhariwal, 2023; Lu and Song, 2024). The associated loss follows from an identical derivation, except it uses two points on the same interpolant trajectory (rather than 2 + which requires access to 6,),
I, = 0, + 3,X1,
(40)
1+ as+A. BstASTI I. Asl, + o(As).
In (40), 2o and iı shared between I, and 1s+As, which yields the second equality for As+As- Following the same steps as for consistency distillation, the final result to replace the semi- gradient (39) by a Monte-Carlo approximation that leverages i, in place of the true vector field (x())) E i, I,
VCT= - (0.--008,,,114))) ((...))))))
(41)
For the parameterization (6), (41) becomes
Ver Vcr = 0., (t
(T.))
(42)
The semigradients (41) and (42) are higher-variance than (35) as Monte-Carlo approximations, but on average give access the ideal flow b rather than the pre-trained, approximate flow h. The gradient (42) is identical to the "mean flow" update recently considered by Geng et al. (2025).
D Connection to shortcut models.
Shortcut models (Frans et al., 2024) correspond to a subset of our proposed PSD scheme (15), which itself is based PFMM. To touch base with the formulation of PFMM in (31), well as the discussion of PSD in the main text, we place shortcut models in our notation here.
Shorteut models consider fixed grid of times 0 = lo < 4 <... < IN 1 spaced dyadically, so that 4 - 4++1 2(4++1 - 44). Observing similar relation to the tangent identity (7), they train te, like a flow matching model and leverage (31) as a bootstrapping mechanism.
Cs(8) = / E |oc. - i1p] + [Rer (Le)) - sg (Xumiiie (SenterrHe)))
(43)
{x,(x)= (t- s)).))).())).
Clearly, the second term in (43) reduces to (31) with s, 4, restricted to a fixed grid. Similarly, (43) corresponds (15) with discretization in time and the specific proposal distribution pi. = dz.- The second term in (43) can also be written entirely in terms of i using the preconditioning discussed later in Section F.I, which leads the exact form the objective discussed in Frans al. (2024).
E Proofs
In this work, assume that all studied differential equations satisfy the following assumption.

---

Assumption E.1. The drift satisfies the one-sided Lipschitz condition
3C>0: (b.(x) bu(w)) (x - v) < - v12 for all (t,z,y) E [0, RR x R (44)
Under Assumption E.I. the classical Cauchy-Lipschitz theory guarantees that solutions exist and unique for all 2g E R and for all t [0, 1].
We first provide self-contained proof of the following proposition, which first appeared in Boffi et al. (2024). We will then apply this result to prove the primary claims of the main text.
Proposition F.2. Let x, denote the flow map (4) for the probability flow equation it = by (2re) Then N,,, satisfies the Lagrangian equation,
0,X (x) (X.((()))
A (2,8,1) Rd x [0, 1|2
(45)
the Eulerian equation,
(z),,( + (x),(())b,())= 0.
A (2, 10,
(46)
and the semigroup property
Xx,(x) = X Xx,,
A (2, e R 10,1]
(47)
Proof. Repeating (4) for ease of reading, the flow map satisfies the jump condition
(x,) = Mt,
(s,() E [0, 1]3,
(48)
where v, denotes trajectory of the probability flow (2). The proof of each condition relies on careful manipulation of this equation.
We first prove the semigroup condition. Observe that
X (X,,(Xe.u = (u,)(zv) = =x (x,(
(49)
Because I, was arbitrary, the result follows.
We now prove the Lagrangian condition. Taking derivative of (48), with respect to t and applying the probability flow (2), we find
0, 0.Xa,(1.)) = -t,
(xx)),
(50)
= be (X,, (xx)).
Because F, was arbitrary, we obtain the Lagrangian condition (45)
Last, we prove the Eulerian condition. Taking total derivative of (48) with respect to we find that
d
X.,( 0,X,, + V.,,((F))aa
ds
(51)
0, 3. (xx) + VX (b,
Again, because v, arbitrary, the result follows.
We now provide a simple proof of the tangent condition we leverage in the main text.
Lemma 2.1 (Tangent condition). Let Xs,t denote the flow map. Then,
lim (x)xa(r)= = g-2t
Vt E [0,
VI E IRd,
(5)
i.e. the tangent vectors to the curve (X,, give the velocity field be(x) for every x.
Proof. By Proposition E.2, we have that the flow map satisfies the Lagrangian equation (45). Taking the limit as + f, and assuming continuity of the flow map, we find
lim 0gX = lim by (x),(x)) by (x,((x) = by (z). --tt
(52)
--
Above, used that (z) for all Rd and for all t [0, 1].
25

---

We now prove Proposition 2.2, which extends Proposition E.2 to the representation (6).
Proposition 2.2 (Flow map). Assume that X is given by (6) with vaa satisfying (7), and assume that ua is continuous in both time arguments. Then, Xs, is the flow map defined in (4) if and only if any of the following conditions also holds:
(i) (Lagrangian condition): S,, solves the Lagrangian equation
(x),,((:)= (x))
(8)
for all (s,t) [0, 1J2 and for all E2.
(ii) (Eulerian condition): Xe,t solves the Eulerian equation
0,X,,(()) + V X (x)00,,,,, = O,
(9)
for all (s,t) E [0, 1J and for all E R..
(iii) (Semigroup condition): For all (s, u) [0, and for all x E Rd.
(u,(XX,,(z)) Xx,(1).
(10)
Proof. We start with the Lagrangian condition (8). By assumption of (7), vt,(2x) b, (x), so that (8)
is equivalent to (45). It follows that the flow map must satisfy (8) by Proposition E.2, which proves the forward implication, To prove the reverse implication, observe that by Assumption E.1, solutions to (8) are unique, so that any solution must be the flow map.
The proof of the Eulerian condition is similar. For the forward implication, we observe that (9) is equivalent to (46), so that the flow map solves (9). Now, let X solve (9) (along with (6) and (7)). We would like to prove that X is the flow map. Let observe that by assumption,
xx,) =0
(53)
ds
where x, is any solution of the probability flow. Integrating both sides with respect to from s to t, find that
(x,( (x,)(x))= =0 - 3e = (x.)
(54)
This is precisely the definition of the flow map.
Last, prove the final property. By Proposition E.2, we have that the flow map satisfies (10), which proves the forward implication. To prove the reverse implication, let X be any map satisfying (6), (7) and (10). Define the notation (v)),( = lim, (y) = ( by by(). Then, consider
Taylor expansion of the infinitesimal semigroup condition for 8x, E Rd x [0, 1] arbitrary,
X (x)+ (x.tt((s,,xx)))
(x,, + h&, +X,,,..X..() o(h),
X2,(x) hö, (X,,, o(h),
(55)
(x)(1) (x,,,((x,...). (x)) o(h),
(x) + hb, (x)) + o(h).
Note that the above Taylor expansion implicitly uses that Us,e is continuous in (s. t) to write 12,4+
2t, + Oh). This rules out the discontinuous solution
(x)
= t
(z,
10
(56)
+ t
which corresponds to (x) = x for all (x, 8, 1) E Rd x [0, 1|2 and satisfies the semigroup condition trivially.
Re-arranging the last line of (55), we find that
(as+h(x)= Xx,( br(X),((z)) + o(1),
(57)
h
so that
Xx)+h(x)- Xx) = (x) = (x)),,(z)).
lim
(58)
h0
h
Equation (58) precisely the Lagrangian equation, whose unique solution is the ideal flow map. This completes the proof.
D
26

---

Given the above developments, we now recall our main proposition. Proposition 2.3 (Self-distillation). The flow map a defined (4) is given for all < by (x)( (t s))aa(()) where (x,)( the unique minimizer over i of
CSD( (o( LD(D).
(11)
where Li( is given by
(o( = Eza.en (T)). - Iel2]dt
(12)
and where Lo(ū) is any of the following three objectives.
(i) The Lagrangian self-distillation (LSD) objective, which leverages (8),
LLSD 10)=1( L. Ezo.er [[0r8s,((s) - Cr, (4,)( dsdt;
(13)
(ii) The Eulerian self-distillation (ESD) objective, which leverages (9),
(E) 1:00)=(())
[ E..... 110,8, + (o (
dsdt;
(14)
(iii) The progressive self-distillation (PSD) objective, which leverages (10),
LPSD (8) = S (.
Ero. (Tx) ( (1a))) dudsdt.
(15)
Above,
s) s) (x) and Exg,, denotes an expectation over the random draws of (2o, 21) in the interpolant defined in (1).
Proof. We first prove the statement for the LSD algorithm. Observe that for any by and any xx,((z) = x + (x))
(b) > Cb(b),
(59)
LLSD (0) > 0.
where bx(x)= Eii = x] is the ideal flow. This follows because Lg is convex in ó with unique global minimizer given by b, while LLSD is a square residual term on the Lagrangian relation (8). From this, we conclude
(t) LD( LISD(() Lb(b). By Lemma 2.1 and Proposition 2.2, the ideal flow map Xs,, satisfies
(60)
(x,((x) = bx())
(x, E Rd x 10, 1], A (x, Rd x 10.
(61)
(xx),((x) (Xx,(())())
From (61), we see that
= = + LLSD = Lo(w) = Lo(b),
(62)
so that X., achieves the lower bound (60) and is therefore optimal. Moreover, any global minimizer must satisfy (61), and by Proposition 2.2 therefore must be the flow map.
We now prove the statement for the ESD algorithm, which is similar. We first observe that for any 0,
Cu(ū) Lb(b),
(63)
LESD (0) > 0.
From above, we conclude
(G) Lb(() LESD(O) Lb()).
(64)
Moreover, by Lemma 2.1 and Proposition 2.2,
(x) (z)
((z E nd [0,
(65)
0,X,,((x) = (X),
(( 10.
27

---

From (65), it follows that the ideal flow map satisfies
(SD(U) Lu(w) LESD(U) Lo(u) Lo(b).
(66)
Equation (66) shows that X achieves the lower bound (64) and hence is optimal. Moreover, any global minimizer must satisfy (65) and therefore by Proposition 2.2 is the ideal flow map.
Finally, we prove the result for the PSD approach. The proposition is stated for uniform proposal distribution over u, but holds for any distribution with full support over [s, t). First, we observe that
LsD i) (o) LPsD(t) Lb(.).
(67)
By the semigroup property (10), we have that the true flow map satisfies
(v) Cs(v) + (PSD(v) (s(v) Cb)))
(68)
so that X optimal. Now, let X be any map satisfying
(xD(X*) La(b),
(69)
i.., any global minimizer of the PSD objective. It then necessarily follows that La(e") Lv(v),
(70)
CpsD(e*) 0.
By Proposition 2.2, under the assumption that v* is continuous, (70) implies that X the ideal flow map X.
We now recall our theoretical error bounds for the LSD and ESD algorithms.
Proposition 2.4 (Wasserstein bounds). Let (x,((x)= = x + (t =)),,(())) denote a candidate flow map, let ف1 = Xoo denote the corresponding one-step generated distribution, and let L denote the spatial Lipschitz constant of (c,(() uniformly in First assume Co(ō) + LisD(0) & Then,
W2 (1),
P1) & 4e142L..
(16)
Now assume that (o( (6) & c. Then,
W2 (11, < 2e (1 + c2h)e.
(17)
For ease of reading, we split the proof of Proposition 2.4 into two results, one for each algorithm. We begin with LSD. Proposition E.3 (Lagrangian self-distillation). Consider the Lagrangian self-distillation method,
(sD(X) = (o()) + (L))(
(71)
Let X denote a candidate flow map satisfying L & and let Pa denote the corresponding pushforward Pı = -No.EPo- Let i. denote the spatial Lipschitz constant of (,)(() uniformly in time, i.e.
[0,,(x)) 0e,c & L|x- -v vl
V (x. Rdd x Rd x 10.1].
(72)
Then,
W (pz, P1) & 4e 122L.
(73)
Proof.
Observe that (X) & implies that both (D( & and (isD(i) & c. We first note that
Lb(0)= = ( Ep, [I0eA(T)) -r13]] dt,
(74)
/ Ep, (10),((1) - (..)) dl
Ep, [1
(b((T))) dt.
In (74), we used that bx(x)) Eli,/e x] along with the tower property of the conditional expectation. It then follows that the L2 error from the target flow is bounded by
S." Ep. [T)),e(T)-- b((T.)1)] dt < E / Ep. [LAP [b,( dt.
(75)
28

---

We now observe that, again by the tower property of the conditional expectation,
S. Ep. 1b.
dt
Ep. [ia [bc( I dt,
(76)
E lis - bo(())) I dt,
Ep.
0.
Equation (76) shows that the term subtracted in (75) conditional variance, and therefore is nonnegative. Combining the two, we find that
L. Epp (10),(11)- be E
(77)
We now consider the learned probability flow
= ,τ (#.),
to ~ Pa-
(78)
By Proposition Albergo and Vanden Mijudnn (2022), (77) implies that
W2 (01. Pi) < el
(79)
where pf = Law(()). Now, by Proposition 3.7 of Boffi et al. (2024), ( & & implies
W3 (pi. P.) < e2
(80)
By the triangle inequality and Young's ' inequality, we then have
W2 (i) < 2 (p1,0)) w (e) (1)))
(81)
$ 4e 122L c.
This completes the proof.
We now prove similar guarantee for the ESD method. Proposition E,4 (Eulerian olff-istillaation). Consider the Eulerian self-distillation method,
LsD Lz(t) + (t).
(82)
Let X denote a candidate flow map with the same properties as in Proposition E.3. Then,
W2 P1) & + ce(1+
(83)
Proof. As in Proposition E.3, our assumption LSD (0) < implies that both Lo( < and LESD (f) & E. Defining the flow åg as in (78), we have a bound identical to (79) on W (pt, p.).
Now, leveraging Proposition 3.8 in Boffi et al. (2024), we have that
W2 (Pf. P1) & ec.
(84)
Again applying the triangle inequality and Young's inequality yields the relation
Wh (p1,P)) < 2 (W2 (pf. w (pi, ())
(85)
< 2e(1 + c2h)s.
This completes the proof.
Further details on self-distillation
In this section, we collect some additional results and detail on some of the topics discussed in the main text.
29

---

F11
Semigroup parameterization for PSD.
By definition, we have that
(2,(2)) = x (t - ())o,()()))
(86)
We then also have that
(x) = x + (u s)Us,a(z), (xx((x,(()) (x,)(r) + (t u)Va, (X
(87)
= (u s) s)us,a (z) + (t - (X
By the semigroup property (10), it follows that
(x) = Xu, (x,,
(88)
from which we see that
+ (t - s))s, (x) = x (u s)v,u(() + (t u (X
(89)
Re-arranging and eliminating, we find that
Us,a(at) 1 u
(x))
ua,t var( (X (Xx.(x)))
(90)
which provides a direct signal for Us.t. Choosing = 78 + (1 7) for y E 10, [] leads to the simple relations
u s = 1
t u
t-s %
(91)
which can be used to precondition the relation (90) as (x)) (1 (e,,ux) + you (x).
(92)
In the numerical experiments, we use (92) to define a training signal for i for the PSD algorithm.
F.2 Limiting relations and annealing schemes.
Limiting relations. As shown in the proof of Proposition 2.2, application of the semigroup property with (s, u, t) = (s, t, + h) for a fixed (s, t) recovers the Lagrangian equation at order h. As shown in the proof of the tangent condition Lemma 2.1, the Lagrangian condition recovers the velocity field in the limit as s + Similarly, if we consider the Eulerian equation in the limit as +
lim 0, X,,,,, r (x) + (x),,(x)b,,) [-
= a (I) bx(x) 0,
(93)
so that 0 X., (x) -.,, - b, (z). In this way, all three characterizations reduce to the flow matching objective for U4,t as the diagonal is approached.
Annealing and pre-training. As a result of (93), we can view training the flow be only on the diagonal /. as pre-training scheme for the map X. This also means that we initialize or, from a pre-trained model in a principled way via appropriate duplication of the time embeddings.
The relations (7) and (93) imply that the off-diagonal self-distillation terms represent natural extension of the diagonal flow matching term. This suggests a simple two-phase curriculum which the flow matching term is trained alone for Nim steps as a pre-training phase, followed by smooth conversion from diagonal training into self-distillation by expanding the sampled range of |t - from 0 to 1 over the course of Namneal steps. This can be accomplished, for example, by drawing (s, t) uniformly on the off-diagonal and then clamping t = min(t, s + 5(k)) where & denotes the iteration and 8(k) is the maximum value of L s|. for example ō(k) k/ 4/Nanneal- For simplicity, we trained directly without any annealing in our experiments, but expect this to simplify and speed up training for large datasets where overfitting is not concern.
F.3 Further details on loss sampling and computation
In this section, we provide further detail on how the choice of 1 E [0. 1], which distributes the batch between the diagonal flow matching term and the off off-diagonal self-distillation term, affects training time. The factor 7) can be chosen based on the available computational budget to systematically trade off the relative amount of direct training and distillation per gradient step. We focus here on the computational cost of forward pass of the objective function; the complexity of a backward pass will depend on the specific choice of sg (-) operator used.
30

---

Flow matching. Evaluating the interpolant loss L. on single sample requires single neural network evaluation (4), leading B network evaluations on batch.
LSD. The LSD objective requires a single partial derivative evaluation (I.)..(I.) and two network evaluations - one for ot and one for Xx,r (L,) per sample. The time derivative is constant factor C' 1.5 more than forward pass, and with standard computational tools such as jvp, can be computed at the same time as X,, (I,). The LSD objective thus requires (1 + C) network evaluations. Adding the diagonal and off off-odiagnna parts, we find complexity of (1==))(+C)++)) for the full self-distillation objective.
PSD. The PSD objective requires three neural network evaluations, so that its expense is 3B. Combining this with the diagonal component, we have (3(1 n) + 17) network evaluations.
ESD. The ESD objective requires a partial derivative evaluation (Lo), a neural network evaluation (Ls), and a Jacobian- product (Xe,,4))) Observing that (8,, V) can be used as one augmented (d + 1)-dimensional gradient, and then observing that (I
V.X,,((ss)0,, (1.) Vo.xXe.((L) (T,) this can be computed as single Jacobian-vector
product. This gives complexity of (1 C, identical to LSD. Adding the diagonal component, we find ((1+ C1 7) B)
F.4 Stopgradient recommendations
The choice of sg (.) operator in the loss is delicate and empirical, as it is very difficult to ascertain the convergence properties of an algorithm operating on an objective leveraging sg (-) a-priori. Nevertheless, in practice, we find it critical for nigh-dimensional tasks such as images to use sg (.) to control the flow of information from the teacher network on the diagonal s = to the off- off-diagonal. For large-scale neural networks, we find empirically that backpropagating through Jacobian-vecte products- in particular spatial Jacobian-vector products leads to significant instability, which can be avoided with sg (.). For low-dimensional tasks with simple neural networks, we found instability to be less of concern.
Following these observations, we found it useful to take insight from the distillation setting described in Section B, leading the configurations
(0) /"" / E [0.8..(1.)) - sg
LESD (f) /' Jo E 11..8a
(V. (1,) []]
(94)
(PsD(0) / S. Er, (Ta) sg (11 - ( To ( (Z.)) PP]
It is also possible to avoid backpropagating through the partial derivative with respect to and with respect t in ESD and LSD by expanding the definition of (x) x + (t- (x) as described in Section C. This reduces the memory overhead even further by avoiding backpropagating through a backward pass of the network.
EMA teacher. addition to the use sg (-), an important consideration is the choice of parameters for the teacher, which provides an alternative perspective and method to implement the sg (-). Making explicit the student parameters 0 and teacher parameters o, we can write for LLSD (with analogous expressions for the other choices),
LLSD (6) (((())//) //' L f
(95)
The recommendation (94) corresponds to taking the gradient of (95) with respect to and then evaluating the result at Φ 0. A second option would be to evaluate & at an exponential moving average 0,
Φk _0 + (1 8)0k.
8 0,
(96)
31

---

where & denotes the optimization step and where 6 denotes forgetting factor such as ó = 0.9999. While in practice we found improved samples by evaluating the learned flow map over EMA parameters (see Section G for exact values), we found that the use of EMA for the teacher parameters offered no gain and sometimes led to instability early experiments. For this reason, we use instantaneous parameters Φ = 0, corresponding to (94) or ė = 0 in (96).
F55 Classifier-free guidance
In this section, we describe how train a flow map with classifier-free guidance. For the derivation, we focus on the LSD algorithm to avoid replicating the loss functions in each case, but the other choices are identical. To this end, let (2;;; e) denote conditional velocity field. We first observe that we may train a conditional flow map via the objective function
(6) - ['00)-/''[ [iiiiaatree)-Et 191 S. [/ 110...... ( [P] (97)
In (97), I denotes conditional interpolant (i.e., with 1;= = a(1)zo + 3()aff with of ~ pi(- drawn conditionally on c), and E now includes expectation over the value of c. To train model that is both conditional and unconditional, we may include = Ø in the expectation.
As in the main text, let us now define the CFG velocity field at guidance strength a E R as
da (x; b, Ø) + (b; ba(1;; Ø)), Ut,,(2; Ø) + a (vt,( c) "t, (x; ();
(98)
Given the second line (98), we define the current estimate of the guided velocity,
de (z; or, c) 0t,e Ø) + a (2;; c) (x; Ø)).
(99)
We now observe that because (99) is constructed entirely in terms of the known 0r, we only need to modify self-distillation term rather than the flow matching term to train CFG flow map. To this end, we self-distill the guided velocity &r over range of This leads to the objective function
CFG = 1/'
(100)
/
E ()) - sg (g.r (Xst P].
where â denotes maximum guidance scale of interest. Following the same derivation, we may obtain the CFG ESD objective
CFG ESD
(
[(na(1;))) i8p]
Jn
(101)
E [1008e + (8
[Pede. (4);a.))
as well as the CFG PSD objective.
f'" E (r, -
CPSD (0)
(102)
Jo [|0aa(15; -sg ((- (117))0sm(((aa (X. p2]
+
F.6 Detailed algorithms for each self-distillation method
Here, we provide detailed algorithmic implementations for each self-distillation method using the recommendations provided in (94). Each algorithm computes the flow matching loss Ls(o) over batch of size 7 and the distillation loss (D( over a batch of size (1 M, comprising a total batch size of M.
32

---

Algorithm 2: Lagrangian Self-Distillation (LSD)
input: Distribution TT); model Us,t coefficients at fe batch size M; diagonal fraction y weight Ws.t-
repeat
Sample Ma |nM] pairs 21) p( 21);
Sample M, times 4 ~ ([0, 1));
Compute interpolants 4,, a1 v + Br, xj and velocities, Lt, òe, + Bi
Compute diagonal loss Lb = Ma
1 Li=l
Ma e 1 (4,,) - i. 12 Wi,b.
Sample M,,- M Ma pairs (wr) x1) ~ (1).
Sample M, pairs (z).4)) ~ Undi
Compute interpolants: Is, ass vb + 3e,
Compute simultaneously via jvp: Xx..t; (1, 0e (4ay)
Evaluate teacher transported point: by, 0r X (Ley)):
Compute residual: ry = 0, X s5,,5(1a)) (Taz) sg (12)
Compute LSD loss LLSD M. Σド=
Mo (e w lryl w Wssti
Compute self-distillation loss LsD Lo + LLSD:
Update 0 and w using LsD:
until converged;
output: Trained flow map Xx,((z) + (t (2),,((z)
Algorithm 3: Eulerian Self-Distillation (ESD)
input: Distribution p(zo, .i ); model U..s3 coefficients or, 3n batch size M; diagonal fraction weight Ws.t-
repeat
Sample M2 = Ln pairs (rg. a1) ~ p(aro x1): Sample M, times 4 U(O, 1]);
Compute interpolants 4, = a1, zb + Be, x and velocities i,, = di, zb + Be, e- (e-mi.h (Te,) i,, we,,t):
Compute diagonal loss: Lo = Me
1 En4
Sample M, M - Ma pairs (n) xt)) p(x00
Sample Mr, pairs (s),4z) ~ Uodi
Compute interpolants: I,, as xy + Bz zt;
Evaluate teacher velocities: b., = 0, s, (Ie,):
Compute simultaneously via single augmented jup: 0gS.y,1s (Ta;); VX (Ta))
Compute Bulerian residual: r; = 0.8 (la,) + sg VX (K, 1o.))
Compute ESD loss: LESD = 1 Ar, EM, (e le- |ryl wo
Compute self-distillation loss LsD Lo + LESDi
Update and w using LLs:
until converged;
output: Trained flow map Xx((x) x + (t - (x)
33

---

Algorithm 4: Progressive Self-Distillation (PSD)
input: Distribution p(zo, T1); model og.t coefficients ar fe batch size M; diagonal fraction y weight Wx,t3 sampling method Py-
repeat
Sample Ma |nM] pairs 21) p(zo, 21);
Sample M, times 4 ~ ([0, 1));
Compute interpolants 4,, at, z + Se, xj and velocities Le. d + Be,ofi
Compute diagonal loss: Ls Ma Ei=1 (e-14,4:00,, (Te.) te, 12 wiit)):
Ma
Sample M, = M - Md pairs (az) x1) ~ p(au,
Sample M, pairs (sy.tz) Uodi
Sample intermediate fractions: Vi ~ p. (e.g., ( 1]) or 81/2);
Compute intermediate times: uj = 85 + (1 - 7/) Yy;
Compute interpolants: L,j Cxa- 20 + 3e, x1;
Evaluate model at student points: Bsy, 13 (Js;):
Evaluate model at first segment: 0o
Compute intermediate flow maps: Xo,,u, (L Lo, + (uj - sy)o,y (Ls,):
Evaluate model at second segment: Day. (X 0y,us (I,z) ))
Compute preconditioned teacher signals:
greacher (1 7t (2)00,,,,, (I,,) + 73 Uy
ag (Ing)
sy.fg
Compute residuals: ry = 0z 1, ) - sg (tencher
Compute PSD loss: LPSD = M. Ej= (e Mey
1ry1 + Ws w,y,;);
Compute self-distillation loss: LSD = Lo + LPSD:
Update i and using LLD;
until converged;
output: Trained flow map X,, xx,(())= x (2)
G Further details on numerical experiments
Here, we provide a complete description of the numerical experiments performed in the main text. A concise summary of each experiment is given in Table 2.
G.1 Checkerboard Details
Experimental setup. We compare the LSD, ESD, and PSD algorithms on the two-dimensional checkerboard dataset. For PSD, we evaluate both uniform sampling (r ~ U(O0, 1)), denoted PSD-U) and midpoint sampling (y = 1/2, denoted PSD-M). We generate a dataset with 107 samples and train for 150, 000 steps with a batch size of 100, 000 and learning rate of 10-3 with square root decay after 35, 000 steps. We use diagonal fraction of 0.75, allocating 75% of each batch to the flow matching loss La and 25% to the self-distillation loss. The network architecture consists of a 4-layer MLP with 512 neurons per hidden layer and GELU activation functions. We use the linear interpolant with a = 1 t. and Bg = t with Gaussian base distribution Zo ~ N(0, /) with adaptive scaling to normalize by the variance of the target distribution. Times are sampled uniformly over the upper triangle without annealing, and we apply gradient clipping at 10.0. All methods use the stopgradient configurations described in (94). We visualize model samples produced from an exponential moving average of the learned parameters with decay factor 0.999. Each experiment was run on a single 40GB A100 GPU. A full qualitative visualization of the tabular results discussed in the main text is shown in Figure 6.
KL Computation. To compute the KL divergence, we leverage that (a) the checkerboard density known analytically as uniform density over the selected squares, (b) the low-dimensionality of the dataset means that histogramming the model samples gives a good approximation of the model density, and (c) the low-dimensionality implies that quadrature can used to compute high-accuracy, deterministic approximation of KL. To this end, first compute 64, 000 samples
34

---

Checker
CIFAR-10
CelebA-64
AFHQ-64
Dataset Properties
Dimensionality
3 32 32 50k
3 64 64 203k
64 64 16k
Samples
107
Network
Architecture
4- 4-yyyer MLP
EDM2 128
EDM2 128 21 per resolution 16 x 16, 8 x 0.0
EDM2 128 21, 3, per resolution 16 x 16, 8 0.0
Hidden/base channels
512
Channel multiplicrs
[2, 2, 2] per resolution 16 x 16 0.13
Residual blocks
Attention resolutions
Dropout
Hyperparameters
Batch size
100,000
512 400,000
256 800,000
256 800,000
Training steps
150,000
Total samples
25 x 10
204.8 x 10 RAdam
204.8 x 10 RAdam
204.8 x 10 RAdam 10-2
Optimizer
RAdam
Learning rate
10-3
10-2
10-2
LR schedule
Sqrt decay at 35k
Sqrt decay at 35k 1.0
Sqrt decay at 35k 1.0
Sqrt decay at 35k 1.0
Gradient clipping
10.0
Diagonal fraction y
0.75
0.75
0.75
0.75
EMA decay
0.999
0.9999
0.9999
0.9999
Evaluation
Metric
KL divergence 64,000
FID 50,000
FID 50,000
FID 10,000
Sample count
Methods
Algorithms
LSD, ESD, PSD-U, PSD-M
LSD, PSD-U, PSD-M
LSD, PSD-U, PSD-M
LSD, PSD-U, PSD-M
Table 2: Experimental setup. Summary experimental configurations across all datasets. experiments use uniform sampling of (s, f) pairs over the upper triangle and leverage the sg (-) choices described in (94),
from each model for each number of steps N. We then histogram these samples using an M x M grid with M = 50 over the range [1, To approximate the KL, we use the quadrature formula
KL (1) 1=)== flog
(1) (1) Pi (x)da
M
M
log
pi pi(3)))
pi (Irrs)
(103)
p (x:z)
where (103) Fiz denotes the center of bin (i,j) used to compute the histogram. We note that because of the uniformity of #ı and ghist, this quadrature rule exact, i.e.
pI (i( pAiot(
log
p (( = KL(PI ll 1
(104)
G.2 CIFAR-10 Details.
We evaluate the LSD, ESD, and PSD algorithms the CIFAR-10 dataset. Again for PSD,
compare uniform sampling (7 ~ U((0, 1)), denoted PSD-U) and midpoint sampling (7 = 1/2, denoted PSD-M). All methods use uniform sampling over the upper triangle (s,t) Uod without annealing. We train for 400, 000 steps with a batch size of 512 and an initial learning rate of 10-2 with square root decay after 35, 000 steps. We use diagonal fraction of 17 0.75, allocating 75% of each batch to the flow matching loss and 25% to the self-dsttilattion loss. The network architecture is based on EDM2 in Configuration (Karras et al., 2024) and NCSN+ (Song al., 2020), using 128 base channels, channel multipliers [2, 2, 2). and residual blocks per resolution. We use positional embeddings time, as we found that Fourier embeddings led greater training instability (Lu and Song, 2025). We embed and (L x) rather than and I, which we found to perform better in early experiments, add these embeddings together, and otherwise use standard FiL. conditioning in the EDM2 network. We apply attention the 16 x 16 resolution and use dropout of 0.13 following EDM
35

---

N=2
1.0
PS0-1 PS1D-1
Figure 6: Checker: full qualitative results Full visualization of sample quality as a function of number of steps on the two-dimensional checker datasct.
recommendations for CIFAR-10 (Karras et al., 2022). We employ a learned weight function ws,t with 128 channels to normalize gradient variance. The interpolant uses = 1 - and B3 f with Gaussian base distribution, setting the variance the Gaussian adaptively to match the variance of the training data. We apply gradient clipping at 1.0 and use the stopgradient configurations described (94) for LSD and PSD; ESD was unstable in every sg (-) configuration tried. We evaluate sample quality using FID computed on-the-fly fly every 10, 000 steps with 10, 000 generated samples, using NFE E 21,2,4, 8, 16] for the flow map. Models were trained from random initialization without pre-training, and we track EMA parameters with decay factors 0.999 and 0.9999. We re-compute FID over 50, 000 generated samples for post-processing and take the best checkpoints for each number of sampling steps optimizer with default settings. Minimal hyperparameter tuning was applied to the algorithms due to well-established training practices for CIFAR-II available the literature.
the entire training range with EMA factor 0.9999. We use the RAdam
G.3 CelebA-64 Details
We compare LSD and both PSD variants (uniform and midpoint) on the CelebA-64 dataset. As for CIFAR-10, we found ESD be uniformly unstable and do report results. We train for 800, 000 steps (corresponding to 204.8M samples) with batch size of 256 and an initial learning rate of 10-2 with square root decay after 35, 000 steps. We use diagonal fraction of 0.75, allocating 75% of each batch to the flow matching loss and 25% to the self-distillation loss. The network architecture is based on EDM2 in Configuration with 128 base channels, channel multipliers [1, 2,3, 4), and residual blocks per resolution, corresponding to the "ImageNet-S" variant reduced from 192 channels to 128. We apply attention at resolutions x 16 and CIFAR-10, we use positional embeddings for time and embed and (l - s) with standard FiLM conditioning. We use the linear interpolant with a, distribution and adaptive scaling normalize to the variance of the target density. Times points (s,t)
and not use dropout. As with
and Br = with Gaussian base
36

---

are sampled uniformly over the upper triangle and no annealing or pretraining is used. We apply gradient clipping at 1.0 and leverage the stopgradient configuration (94) for all methods. We use the RAdam optimizer with default settings. We evaluate online sample quality using FID-10K computed every 10, 000 steps, and then compute FID-50k post-hoc find the best model, following the same steps as for CIFAR-10. Models were trained from random initialization without pre-training, and we track EMA parameters with decay factor 0.9999. FID-50K scores were computed with this EMA factor.
G.4 AFHQ-64 Details
We compare LSD and both PSD variants (uniform and midpoint) on the AFHQ-64 dataset. As with 6-4, we found ESD to be unstable and do not report results. We train for 800, 000 steps (corresponding 204.8M samples) with batch size of 256 and initial learning rate of 10-2 with square root decay after 35, 000 steps. We use a diagonal fraction of 17 = 0.75, allocating 75% of each batch to the flow matching loss and 25% to the self-distillation loss. The network architecture based on EDM2 in Configuration G with 128 base channels, channel multipliers [1, 2, 4], and 3 residual blocks per resolution, matching the architecture used for CelebA-64. We apply attention at resolutions 16 16 and 8 x 8, and do not use dropout. As with CIFAR- 10, we positional embeddings for time and embed s and (t ---] s) with standard FiLM conditioning. We use the linear interpolant with a, = 1 - and 3, f with Gaussian base distribution and adaptive scaling to normalize to the variance the target density. Time points (s, t) are sampled uniformly over the upper triangle and annealing or pretraining is used. We apply gradient clipping at 1.0 and leverage the stopgradient configuration (94) for all methods. We use the RAdam optimizer with default settings. We evaluate online sample quality using FID-10K computed every 10, 000 steps, and then compute FID-50k post-hoe to find the best model, following the same steps as for CIFAR-10 10 and CelebA-64. Models were trained from random initialization without pre-training, and we track EMA parameters with decay factor 0.9 0.9999. FID-50K scores were computed with this EMA factor.
37


Reviews:
[{"review_id": "vYLalWMGoa", "reviewer": "Reviewer_oUCA", "summary": "This paper introduces a unifying framework that encompasses a wide range of existing distillation schemes.  By analyzing how different conditions imposed by their characterisation of the flow map, alongside the insight that the velocity field corresponds to the instantaneous rate of change of the flow map, the authors reveal connections between established distillation algorithms. Building on this perspective, the paper proposes a novel self-distillation algorithm to train flow maps. The authors prove that their training scheme and the associated loss functions are well-posed—meaning each admits a unique solution. Additionally, they provide theoretical guarantees in the form of bounds on the 2-Wasserstein errors from the loss values for 2 of the loss function proposed in their paper. Empirical evaluations demonstrate that the proposed training approach and loss functions achieve promising results for unconditional generation, validated on both synthetic datasets (such as checkerboard patterns) and the CIFAR-10 benchmark.", "questions": "1. Conduct and report systematic hyperparameter sweeps for key variables such as learning rate, model depth, regularization parameters, and optimizer hyperparamters and choice. Analyze how sensitive the training dynamics and performance are to these settings. Highlight scenarios where the method exhibits instability or performance drops.\n2. Experiment with alternative flow map architectures—varying network depths, widths, or activation functions—to identify which design choices most affect robustness. Discuss not only successes but also observed failure modes linked to specific architectures.\n3. Evaluate the method using more challenging datasets, such as high-resolution or multi-modal benchmarks (e.g., CelebA) or even datasets for discrete random variablaes. This will help uncover common failure modes that may not be apparent with simpler data.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n1. The paper presents a comprehensive framework that unifies a wide range of distillation methods, offering a cohesive perspective that deepens the understanding of existing approaches and an approach to how new distillation algorithms can be systematically derived.\n2. The authors provide theoretical guarantees for their framework, which also offer insights into its relationship with related algorithms.\n3. The framework opens up interesting possibilities, such as incorporating multiple models and exploring additional extensions.\n\nWeakness\n1.  While the results on the 2D checkerboard dataset are promising, the empirical evaluation is restricted to relatively simple datasets. The case for the method’s effectiveness would be much stronger if it were demonstrated on more complex, high-resolution, and multi-modal datasets, such as CelebA\n2. The descriptions of the algorithmic aspects—particularly loss reweighting, as well as temporal and PSD sampling—are at times difficult to follow; the paper would benefit from clearer exposition or additional explanation in these sections.\n3. The discussion around hyperparameter sweeps and their impact when training flow maps for CIFAR-10 in the appendix could be expanded to provide more useful insights for the reader.", "was_revised": true, "final_justification": "The authors have done more experiments and added more theoretical results to strengthen their claims as well as clairfy them. This justifies the changes in score.", "rebuttal": "We thank the reviewer for their careful reading and evaluation of our work. Below, we address the main questions and weaknesses raised. If our arguments are found convincing, we would greatly appreciate the reviewer raising their score. We would equally appreciate further suggestions for us to improve our work, and will happily provide further clarification on any aspects of our paper or our answers.\n\n## Weaknesses\n*While the results on the 2D checkerboard dataset are promising, the empirical evaluation is restricted to relatively simple datasets. The case for the method’s effectiveness would be much stronger if it were demonstrated on more complex, high-resolution, and multi-modal datasets, such as CelebA.*\n\nWe agree with the reviewer that the initial evaluation was limited, and we thank them for drawing our attention to this point. Since receiving the reviews, we have launched a round of experiments on CelebA-64. In these experiments, we first anneal by training only on the diagonal $s=t$ for 24 hours on 4x H100s with a batch size of 256. Sampling directly with the learned $\\hat{v}\\_{t, t}$ achieves an FID of 4.3 using 64 steps of the Heun sampler. We leverage an architecture based on the small EDM2 configuration for ImageNet64 [1], using $128$ channels for the estimated log-variance weight factor $w_{s, t}$, $128$ channels at the first layer of the UNet hierarchy, a four-layer hierarchy with channel multipliers in the pattern $(1, 2, 3, 4)$, three residual blocks per layer, positional encoding, attention at the $8\\times 8$ and $16\\times 16$ resolutions, and zero dropout. We then expand the range of $|s-t|$ linearly away from the diagonal for $10,000$ gradient steps and train for an additional 24 hours using the full self-distillation methodology. We study the performance as a function of the algorithm and the number of sampling steps. \n\nOur empirical FID scores are shown in the table below. What we find confirms our results on CIFAR-10, again highlighting dataset-dependence and differences between stylized low-dimensional examples and images. Our results demonstrate that LSD and PSD outperform ESD, which is again unstable. In this case, LSD performs the best at all numbers of sampling steps.\n\n| Method | 1 step | 2 step | 4 step | 8 step | 16 step |\n|--------|--------|--------|--------|--------|---------|\n|   LSD    | **15.77**       | **8.53**       |  **4.70**      |  **3.20**      | **2.66**        |\n|   PSD-M     |  17.24      |   9.34     |   5.77     |  3.97      |  3.25       | \n|   PSD-U     |  17.94      |  9.19     |   5.55     |   4.11    |   3.55      | \n\nGiven time constraints and an inability to upload new visual results with this year's rebuttal format, we chose to focus solely on Celeb-A. For the camera ready version, we also plan to launch a round of experiments on the LSUN-Bedrooms dataset to further validate our findings, and are currently running some experiments with the SiT architecture on ImageNet-256 [2].\n\n*The descriptions of the algorithmic aspects—particularly loss reweighting, as well as temporal and PSD sampling—are at times difficult to follow; the paper would benefit from clearer exposition or additional explanation in these sections.*\n\nWe thank the reviewer for bringing our attention to this point. We have significantly re-worked this section to improve the clarity of presentation. In particular, we have:\n\n- Clarified how to make use of multiple models. To this end, we have added a section to the appendix explicitly providing the loss $\\mathcal{L}(\\hat{b}, \\hat{X}) = \\mathcal{L}_b(\\hat{b}) + \\mathcal{L}_D(\\hat{X}; \\hat{b})$ that simultaneously trains a flow model $\\hat{b}$ and distills it into a teacher.\n- Clarified our representation. In particular, we have highlighted that our functional form $X_{s, t}(x) = x + (t-s)v_{s, t}(x)$ is not an approximation: it can be used without loss of generality, and was chosen for computational efficiency, because it leads to the expression $\\lim_{s\\to t}\\partial_t X_{s, t}(x) = v_{t, t}(x)$. We have also explained more clearly how taking additional terms allows us to split the role of the velocity field $b_t$ and the \"residual\" term $\\psi_{s, t}$, for example in the representation $X_{s, t}(x) = x + (t-s)b_t(x) + \\frac12(t-s)^2\\psi_{s, t}(x)$.\n- We have updated the description of the learned loss weight $w_{s, t}$ to clarify its role. Originally introduced in the EDM2 paper [1], this term learns to weight different times in the loss equally, leveraging insights from multitask learning. We have also added a section to the appendix, generalizing the derivation from [1] to the two-time setting, highlighting that the minimum over $w_{s, t}$ normalizes the variance of the loss between the different time points $(s, t)$.\n- We have clarified the role of the sampling distribution $p_{s, t}$ and how it differs from the weight $w_{s, t}$. The sampling distribution $p_{s, t}$ defines how we \"select\" time points for the Monte-Carlo approximation of the loss, while the weight $w_{s, t}$ defines how much each such pair contributes to the loss. We have also significantly simplified our discussion of how we select the sampling distribution $p_{s, t}$, which we agree was cumbersome in the original submission. We have clarified the differences between the uniform distribution on the diagonal and uniform on the off-diagonal, as well as the role of the $\\eta$ parameter in tuning the computational effort, and moved many of the technical details to the appendix. For further computational detils on the choice of $\\eta$, please see our response to Reviewer cfDL.\n- We have extended and strengthened the discussion of PSD sampling to make this important design decision more clear. In PSD, the flow map model $\\hat{X}$ learns to make a jump from time $s$ to time $t$ equal to the two-step jump from $s$ to $u$ and then from $u$ to $t$. The sampling distribution $p_u$ then defines how this intermediate time $u$ is selected. A natural way to parameterize this time $u$ is to write it as a convex combination of $s$ and $t$, $u = \\gamma s + (1 - \\gamma)t$ for $\\gamma \\in [0, 1]$. This parameterization is useful because $\\gamma$ is always on the same scale independent of the selected values of $s$ and $t$. In the original version, we show in the appendix that re-writing the objective function in terms of $\\gamma$ preconditions the loss by placing terms on the same scale. In the revision, we have updated this discussion, highlighting more explicitly how this preconditioning removes terms such as $t-u$ that can be on a different scale depending on the values of $(s, u, t)$.\n    \n*The discussion around hyperparameter sweeps and their impact when training flow maps for CIFAR-10 in the appendix could be expanded to provide more useful insights for the reader.*\n\nWe agree with the reviewer that the original discussion on experimental details for CIFAR-10 was somewhat sparse, which we have significantly expanded as a result. We have clarified the specific network architecture by providing a reference to Configuration G of [1], along with exact definitions of each architectural hyperparameter. In addition, we have added further details on the annealing schemes -- in particular, the precise configuration of loss functions we use over the training trajectory, how the anealing is performed, how many GPUs were required, and the timing information. We have added an analogous section for our new results on the Celeb-A dataset.\n\n### References\n\n[1] Karras, Tero, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. “Analyzing and Improving the Training Dynamics of Diffusion Models.” arXiv:2312.02696.\n\n[2] Ma, Nanye, Mark Goldstein, Michael S. Albergo, Nicholas M. Boffi, Eric Vanden-Eijnden, and Saining Xie. SiT: Exploring Flow and Diffusion-Based Generative Models with Scalable Interpolant Transformers. arXiv:2401.08740."}, {"review_id": "dDx942WnWZ", "reviewer": "Reviewer_cfDL", "summary": "This paper proposes a two-time flow-matching approach to generative modelling, closely related to diffusion and flow matching. Instead of learning only a vector field (that acts in the continuity equation), the paper proposes learning a two-time flow map (that acts in the pushforward equation), which, in theory, allows one to map directly from noise to data in one function call. The paper models this flow map with a Euler-step discretisation of a vector field model.\n\nBased on the theoretical properties of the flow map, the paper proposes a loss function (Eq. 11) that decomposes into a standard flow matching loss (Eq. 12) and a combination of three other derived losses (Eq. 13-15) – LSD (matches final time derivative to vector field), ESD (matches the initial time derivative to the vector field), and PSD (matches the semigroup property of the flow map). The paper then devises four different models – LSD, ESD, PSD-M, and PSD-U – whose loss function corresponds to the three derived losses (with PSD-{U, M} being sub-variants). \n\nThe paper then tests them on the Checkers dataset (in terms of KL-divergence) and CIFAR-10 dataset (in terms of FID). Results show that in the Checkers dataset (lower-dimensional, higher batch size) that the PSD approach wins for higher step sizes, while for the CIFAR-10 dataset (higher-dimensional, lower batch size) that the PSD approach wins for lower step sizes.", "questions": "1. Can you please justify the first-order (Euler) flow-map parameterization (Eq. 6)?\n    - Why is it sufficient, and were higher-order forms tried?\n    - Can you provide an ablation comparing Euler vs. a 2-step Heun/midpoint alternative, or provide a theoretical argument that higher-order terms offer no practical gain?\n    - This assumption seems to underpin every experiment – if it is a hidden bottleneck, results might not generalise.\n2. Can you report the compute-performance trade-off of $\\eta$-sampling (mentioned in line 243)?\n    - Can you provide an experiment that sweeps $\\eta$ and plots validation loss/FID vs. training cost?\n3. Can you evaluate the proposed \"Multiple-model\" and \"General representation\" variants (mentioned in lines 209 and 218)?\n    - Providing at least a small ablation (e.g., on MNIST) would show whether the single-network design is a conscious trade-off.\n    - If impractical, can you provide guidance on what the effects of representations would be and when to prefer them?\n    - These variants are positioned as advantages, but were not validated.\n4. Can you provide more details about the KL divergences in Table 1?\n    - Can you specify how the KL is estimated, which subset of data is held out, and report standard error over at least 3 runs?\n5. Can you address the fairness of the comparison between Checker (Table 1) and CIFAR-10 (Table 2)?\n    - Why is it that the trends in Tables 1 and 2 seem to be opposite? Vector field supervision appears to be more effective for small steps in Checkers, whereas flow map only supervision is more effective for small steps in CIFAR-10, and vice versa.\n    - How do these trends relate to the data dimension? Could you do an experiment that scales dimension and/or batch size?\n    - How fair is comparing these trends, given that the batch size in Checkers is 64000 and in CIFAR-10 is 512?\n    - In Table 2, why does PSD-M get better than PSD-U for more steps? Also, why does this trend not seem to show in Table 1?\n6. Can you elaborate on why we do not also train on the lower diagonal of the unit square (ref Fig. 2) for reverse steps (data to noise direction)?\n    - What would be the effect of doing so?\n6. Miscellaneous:\n    - Can you elaborate on what \"Lagrangian\" and \"Eulerian\" perspectives mean in this context?\n    - Can you please elaborate on lines 209-211? Why do we have to \"learn strictly more than the velocity field\"?\n    - Regarding line 248: Can you explicitly show the correspondence to shortcut models?\n    - Regarding line 262: How do you do the smooth conversion from diagonal training to self-distillation?\n    - Regarding lines 271-272: What do you mean by \"use the instantaneous parameters as the teacher\"? And can you elaborate on how EMA is implemented here?\n    -  Why is the flow map a function of two times?\n    - Regarding line 598: How are the hyperparameter sweeps done?\n    - Can you clarify: Is the loss in Eq. 12 enforced for all models (LSD, ESD, PSD-U, PSD-M)? Please clarify exactly which losses are used for these four models.", "limitations": "I felt that the technical limitations could use more elaboration – see above comments.", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n1. Further elucidates the theoretical relation between the flow map and the vector field in the context of flow matching and diffusion models.\n2. Provides some insightful practical details, e.g., about gradient norms under each of the proposed models.\n3. Achieves positive results on CIFAR-10.\n\nWeaknesses\n1. The trend observed in Checkers (Table 1.) – PSD better than LSD for more steps – seems to be opposite of the trend observed for CIFAR-10 – PSD better than LSD for fewer steps. More justification or elaboration would help to make conclusions more concrete.\n2. No means and standard deviations are given for Tables 1 and 2 over multiple runs. This would make the results more robust and perhaps better characterise the trends mentioned in the previous point.\n3. Assumes a specific parameterisation of the flow map (Euler step). It would help to justify this more and elaborate on related limitations.\n4. Assumes training only on the upper triangle – time pairs in the noise-to-data direction. No discussion, as far as I can tell, on training in the full unit square of time. Also, no discussion on the choice of a two-time flow map rather than one-time.\n5. On first read, the notation is sometimes a bit hard to keep track of, and the correspondence between the loss functions and the tested models could be clearer.\n6. Discussions of some ideas in the paper – e.g., Multiple Models (line 209) and  General Representations (line 218) – give the impression that they have been included in experiments (since they are included in the section \"Algorithmic aspects\". I would suggest making it more explicit which discussions pertain to ideas for future work and which pertain to experiments implemented in the paper.", "was_revised": true, "final_justification": "The authors provided reasonable explanations to my questions that bolster the contribution.", "rebuttal": "We thank the reviewer for their thorough and thoughtful comments on our work. We provide aggregate responses to their mentioned weaknesses and the questions related to them below. Due to the 10k character limit, we had to remove some, but have answers to all and are happy to provide them in the discussion period. If satisfactory, we would greatly appreciate an improved score.\n\n**CIFAR-10 vs. Checkerboard.** Our goal in this work is to introduce a systematic framework for training consistency models. We hypothesize that the optimal training scheme depends on the network architecture, the dataset, and the modality, which is why we lay out a family of methods rather than one. Our results support this hypothesis, with low-dimensional datasets exhibiting distinct trends from higher-dimesional vision datasets. To make this more clear, we have added additional text to the discussion and introduction that explicitly describes these distinct trends. As described in our reply to Reviewer oUCA, we have also performed additional experiments on the CelebA-64 dataset, and these new results support our claims even further: LSD performs best on Celeb-A, even though it is higher dimensional than CIFAR, suggesting the performance is very task dependent.\n\n**Error bars.** In our experience, variance in quantitative results due to network initialization is small. This is supported by common practices for results in the field: we performed a literature review and found that none of the major papers on consistency models included error bars [1-5]. For this reason, during this rebuttal period, we have focused on the generation of additional results rather than on quantifying error for our existing results. Nevertheless, we are happy to include them for the camera-ready version.\n\n**Euler parameterization.** While the parameterization that we exploit is reminiscent of an Euler step, we emphasize that the ansatz $X_{s, t}(x) = x + (t-s)v_{s, t}(x)$ can be used without loss of generality and introduces no approximation error. To see this, we may observe that the function $v_{s, t}(x) = \\frac{1}{t-s}(X_{s, t}(x) - x)$ is simply a time-dependent shift and rescaling of the flow map. We have amended the text to stress there is no limitation in using it, it just allows us to avoid using autograd to take a time derivative.\n\nWe also stress our algorithmic formulation is general and allows for an arbitrary parameterization of the map, though this choice is natural for the reason just described.\n\n**Upper triangle.** We trained over the upper-triangle because it corresponds to the generative pathway. Our framework is identical over the entire unit square. Earlier work has demonstrated applications of this to style transfer and image editing [6], which could also be performed using our methodology. In the revision, we have updated the text with a remark describing this point.\n\n**Two-time Maps.** There are several reasons to train a two-time map rather than a single-time map $f_s(x) = X_{s, 1}(x)$, which corresponds to a consistency model.\n\nWith a single-time map, there is no principled way to sample in multiple steps that increases accuracy. Consistency models alternate between adding noise and applying the map, which only has a small effect on quality [2,4,5]. This fails because adding noise jumps between trajectories of the probability flow ODE, and hence has no convergent limit. In our approach, we converge to an integral of the underlying flow: \n\n$$\\lim\\_{k\\to\\infty} \\circ\\_{i=1}^k \\hat{X}\\_{t\\_{i-1}, t\\_i}(x\\_0) = x\\_0 + \\int\\_0^1 \\hat{v}\\_{\\tau,\\tau}(x(\\tau))d\\tau, \\qquad 0 = t_0 < t_1 < ... < t\\_k = 1,$$\n\nwhich enables us to trade inference-time compute for sample quality. As shown by our numerical results, this is observed in practice, with a higher number of steps leading to better quantitative results on all problems.\n\nThe one-time map is also more limited from a training perspective. Because $f_s$ has neither $t$ dependence nor a semigroup property, there is only an Eulerian representation and neither LSD nor PSD can be used.\n\n**Notation.** We agree that the notation in our submission was cumbersome. The major culprit was a typo that incorrectly used the spurious notation $\\mathcal{L}\\_{\\mathsf{SSD}}$ in place of the objective $\\mathcal{L}\\_{\\mathsf{PSD}}$. This typo clashed with the $\\mathcal{L}\\_{\\mathsf{SD}}$ self-distillation objective. In the revision, we have corrected this error, and we propose the following notation:\n\n- $\\mathcal{L}_{b}$: flow matching / interpolant objective.\n- $\\mathcal{L}\\_{\\mathsf{D}}$: distillation objective, taken in practice to either be $\\mathcal{L}\\_{\\mathsf{LSD}}$, $\\mathcal{L}\\_{\\mathsf{ESD}}$, or $\\mathcal{L}\\_{\\mathsf{PSD}}.$\n- $\\mathcal{L}\\_{\\mathsf{SD}} = \\mathcal{L}\\_{b} + \\mathcal{L}\\_{\\mathsf{D}}$: self-distillation objective instantiated by choice of $\\mathcal{L}\\_{\\mathsf{D}}$.\n\nWe have added a section to the appendix that provides a glossary on notation, as well as provides the recommended form of each objective function, including the use of $\\mathcal{L}_b$.\n\n**Algorithmic aspects.** We agree that the presentation of this section may have been confusing because the two future directions came first. We have re-ordered it so that the future directions are presented *last*, and we have added text emphasizing that they were not explored in the present work.\n\n**$\\eta$-sampling.** We have added the following discussion to the appendix, which quantifies each method's cost.\n\nEvaluating the interpolant loss $\\mathcal{L}_b$ on a single sample requires a single neural network evaluation $b_t(I_t)$, leading to $B$ network evaluations on a batch.\n\nThe LSD objective requires a single partial derivative evaluation $\\partial_t X_{s, t}(I_s)$ and two network evaluations -- one for $v_{t, t}$ and one for $X_{s, t}$ -- per sample. The time derivative is a constant factor $C \\approx 1.5$ more than a forward pass. The LSD objective thus requires $(2 + C)B$ network evaluations. Adding the diagonal and off-diagonal parts, we find a complexity of $((1-\\eta)(2 + C) + \\eta)B$ for the full self-distillation objective.\n\nThe PSD objective requires three neural network evaluations, so that its expense is $3B$. Combining this with the diagonal component, we have $(3(1-\\eta) + \\eta)B$ network evaluations.\n\nThe ESD objective requires a partial derivative evaluation $\\partial_s X_{s, t}(I_s)$, a neural network evaluation $v_{s, s}(I_s)$, and a Jacobian-vector product $\\nabla X_{s, t}(I_s)v_{s, s}(I_s)$, leading to a complexity of $(2C + 1)B$. Adding the diagonal component, we find $((2C+1)(1-\\eta) + \\eta)B$.\n\nThe factor $\\eta$ can be chosen based on the available computational budget to systematically trade off the relative amount of direct training and distillation per gradient step.\n\n**Training Cost.** With the addition of our new results on the CelebA-64 dataset, we feel that our text answers this question to a complete level. On the checkerboard, we fixed the total number of iterations, highlighting that when allowing for differences in computational expense we obtain better results with the cheaper LSD and PSD methods. On the image datasets CIFAR-10 and Celeb-A, we fix the total training time, allowing for different numbers of iterations. We find a similar pattern, where LSD and PSD outperform ESD. To make this more clear, we have added further empirical detail both to the main text and appendix.\n\n**Multiple models.** We have performed these ablations on MNIST and find high-quality visual results in both cases. In particular, we tried the higher-order / multiple-model representation, as well as the more general representation $X_{s, t}(x) = x + \\phi_{s, t}(x)$. In the camera-ready version, we will include these results in the appendix.\n\n**KL Computation.** We generate $10^6$ samples from the learned models and compute approximate densities via histogram. Because we have access to the analytical functional form of the checkerboard density, we integrate the definition of the KL divergence $\\mathsf{KL}(\\rho^* || \\hat{\\rho}) = \\int \\log(\\frac{\\rho^*(x)}{\\hat{\\rho}(x)})\\rho^\\*(x)dx$ via quadrature. In the camera-ready version, we will include error bars. We'll include details on this KL computation in the appendix.\n    \n**Fairness of comparison.** We would like to stress that our results show that no single objective is universally best; performance depends on the network and task. The Lagrangian and Eulerian losses more strongly enforce the dynamical equations for the map, but they require gradients through space and time derivatives, which can be high variance for complex models like UNets. On the checkerboard, where the network is a simple MLP, these stronger objectives excel. We have clarified this in the revision and proposed exploring architectures better suited to such losses as a future direction.\n\n**Shortcut Models.** We have included a new section in the appendix that gives the details of this mapping. Let $\\Delta = t-s$. If we re-write $$X_{s, t}(x) = X_{s, s+\\Delta}(x) = x + \\Delta \\, v_{s, s+\\Delta}(x),$$ use the PSD-M objective with PSD scaling, and place a stopgrad operator on the teacher, our loss function becomes $$\\mathcal{L}\\_{\\mathsf{PSD}}(\\hat{v}) = \\int\\_0^1 \\int\\_0^t \\mathbb{E}\\left[|\\hat{v}\\_{s, s+\\Delta}(I_s) - \\frac{1}{2}\\mathsf{sg}(\\hat{v}\\_{s, s+\\Delta/2}(I_s) + \\hat{v}\\_{s+\\Delta/2, s+\\Delta}(\\hat{X}\\_{s, s+\\Delta/2}(I\\_s)))|^2\\right]dsdt.$$ This is precisely the shortcut model objective function.\n\n\n## References\n[1] Geng et al, “Mean Flows for One-Step Generative Modeling.”\n\n[2] Song et al, “Consistency Models.”\n\n[3] Kim et al,  “Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion.”\n\n[4] Lu et al, “Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models.”\n\n[5] Song et al, “Improved Techniques for Training Consistency Models.”\n\n[6] Boffi et al, \"Flow Map Matching with Stochastic Interpolants: A Mathematical Framework for Consistency Models.\""}, {"review_id": "V2RrmYsuwy", "reviewer": "Reviewer_2Efm", "summary": "The manuscript focuses on learning flow map-based models and how they can be converted into self-distillation models. One major advantage of using flow map models over traditional flow models is the inference-time speedup, as they do not require ODE or SDE discretization. Most existing flow map-based methods follow a two-stage setup: first learning a flow (teacher), then learning its solution (student). In contrast, this manuscript centers on directly learning the flow map.", "questions": "Line 133: \"along trajectories of the probability flow\" — This is slightly misleading, as it does not follow the full trajectory. Rather, it only requires correct prediction of $X_t$ given $X_s$ (as also illustrated in Figure 1b).\n\nProposition 2.4 defines an $\\epsilon$ with respect to ESD, but the remark on line 201 references both ESD and LSD. Could you define an $\\epsilon$ for LSD as well?\n\nThe notation in Equation (17) is somewhat confusing, particularly the first exponential term.\n\nIn proposition 2.4, L is the constant for $v_{t,t}$, is it for all t?  For example, by taking the sup over all t ?", "limitations": "The limitations of the approach should be discussed in a concluding section or highlighted through experiments.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n\n- The manuscript is very well written and clear.\n\n- The introduction and background sections are concise yet informative.\n\n- The method is theoretically grounded. Under certain assumptions, the authors derive a bound on the approximation error (in terms of 2-Wasserstein distance), which depends on a Lipschitz constant. They also directly relate the error with the training loss. \n\nWeaknesses\n\nMy understanding is the approximation quality is inherently limited by how well a truncated Taylor polynomial can approximate the flow map. There is no a priori guarantee that the Taylor approximation will perform well.\n\nThe error bound depends on the Lipschitz constant. One way to achieve a smaller error is by inducing “straighter” flows. A known strategy for this is to use optimal transport (OT), or mini-batch OT, in flow matching. In that setting, the training remains the same, but the pairs $(x_0,x_1)$ are drawn from the OT coupling [1]. This could also be applied here, where in Algorithm 1, the authors sample $(x_0,x_1)$ from the OT coupling. Trying this might help alleviate the lack of control over the quality of the Taylor approximation.\n\nOverall, I find the experimental section insufficient, particularly due to the lack of comparisons. While I agree that achieving state-of-the-art performance is not necessary, at a minimum, the authors should include results for a naive two-stage self-distillation baseline. Even if the results are similar, the proposed method remains valuable because it only requires one training phase.\n\n[1] Tong, Alexander, et al. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint arXiv:2302.00482 (2023).", "was_revised": true, "final_justification": "I would like to thank the authors for the fruitful discussion. I will keep my score as it is. In my opinion, the experimental section is insufficient. The experiments primarily report FID, yet it is also argued that too much importance should not be placed on FID values. The experimental section should address one key question: How accurate are the methods at learning flow maps? With the current setup, this question is only partially answered by showing FID for different step counts. Alternative experimental designs could have better addressed this, for example, by reporting FID scores for the same step counts using a model not trained with the self-distillation objective.", "rebuttal": "We thank the reviewer for their careful reading and evaluation of our work. Below, we address the main questions and weaknesses raised. If our arguments are found convincing, we would greatly appreciate the reviewer raising their score. We would equally appreciate further suggestions for us to improve our work, and will happily provide further clarification on any aspects of our paper or our answers.\n\n## Weaknesses\n*My understanding is the approximation quality is inherently limited by how well a truncated Taylor polynomial can approximate the flow map. There is no a priori guarantee that the Taylor approximation will perform well.*\n\nWe thank the reviwer for this important question, which was also raised by Reviewer cfDL. While the parameterization that we exploit is reminiscent of an Euler step, and while we find it useful to think about it in this way, we emphasize that the ansatz $X_{s, t}(x) = x + (t-s)v_{s, t}(x)$ is without loss of generality and introduces no approximation error. To see this, we may observe that the function $v_{s, t}(x) = \\frac{1}{t-s}(X_{s, t}(x) - x)$ is simply a time-dependent shift and rescaling of the flow map. \n\nOur rationale for choosing this representation is computational: it eliminates the need to compute the limiting partial derivative $\\lim_{s\\to t}\\partial_t X_{s, t}(x)$ that appears in the tangent identity, which our approach exploits to derive an off-diagonal signal. As we show in the text, this limit is precisely given by $v_{t, t}(x)$, and hence reduces to a single neural network evaluation rather than a Jacobian-vector product. In the revision, we have added several lines of clarification around the introduction of this parameterization. In particular, we emphasize that there is no approximation or limitation in using it. \n\nWe would also like to emphasize that our algorithmic formulation is general and allows for an arbitrary parameterization of the map, though this choice is a natural one for the reasons just described.\n\n*The error bound depends on the Lipschitz constant. One way to achieve a smaller error is by inducing “straighter” flows. A known strategy for this is to use optimal transport (OT), or mini-batch OT, in flow matching. In that setting, the training remains the same, but the pairs  are drawn from the OT coupling [1]. This could also be applied here, where in Algorithm 1, the authors sample  from the OT coupling. Trying this might help alleviate the lack of control over the quality of the Taylor approximation.*\n\nWe thank the reviewer for this very interesting suggestion. In response, we have added a short discussion of the role of couplings (and in particular mini-batch OT) in the \"Algorithmic aspects\" section, emphasizing that they made significantly improve the training. Given the large amount of contributions already present in the paper, we feel that a thorough investigation of the role of couplings would be best left for future work.\n\n*Overall, I find the experimental section insufficient, particularly due to the lack of comparisons. While I agree that achieving state-of-the-art performance is not necessary, at a minimum, the authors should include results for a naive two-stage self-distillation baseline. Even if the results are similar, the proposed method remains valuable because it only requires one training phase.**\n\nWe thank the reviewer for suggesting this important baseline, as well as for suggesting that we improve the quality of the experimental section. To address this concern, we have completed an additional set of runs on the CelebA-64 dataset, which we feel significantly expands the scope of the present work.\n\nWe agree that a comparison to a distillation baseline could be a natural experiment to include. However, we would like to respectfully point out that it is somewhat unclear what this baseline should be. We are not aware of a distillation baseline for PSD outside of progressive distillation [1], which requires recursively training a hierarchy of models and as a result is a much more complicated pipeline than what we propose here. We show that ESD is unstable in comparison to PSD and LSD, which was also observed in prior work when used for distillation on image data, making a baseline comparison challenging [2]. While clever stopgrad-based tricks have recently been introduced in the literature to stabilize Eulerian methods [3], incorporating these approaches into our paradigm is out of scope for the present work: in particular, we expect similar techniques can be used to improve PSD and LSD, which we plan to investigate systematically in future projects. \n\nThese observations leave only the LSD algorithm, for which a natural distillation baseline would be the Lagrangian Map Distillation (LMD) approach [2]. If the reviewer would like to see the associated FID scores for comparison, we will gladly include them in the camera ready version. In the meantime, because of this difficulty with defining an appropriate baseline, we have chosen to focus on the new round of experiments on CelebA-64.\n\n*The limitations of the approach should be discussed in a concluding section or highlighted through experiments.*\n\nThe main limitation of our present contribution is that, while we analytically characterize the vast design space of consistency models, we are unable to systematically test each component empirically due to the large associated computational expense. Critical aspects of the pipeline that deserve further testing include ablations over the flow map parameterization, stabilization and stopgradient schemes, as well as annealing schemes. In the camera-ready version, we will include a section covering this important limitation.\n\n## Questions\n*Line 133: \"along trajectories of the probability flow\" — This is slightly misleading, as it does not follow the full trajectory. Rather, it only requires correct prediction of $X_t$ given $X_s$ (as also illustrated in Figure 1b).*\n\nWe agree with the reviewer that the flow map can take large jumps along trajectories, and hence it does not necessarily follow the full trajectory. What we mean by this statement is that the flow map always maps along a single trajectory by definition -- i.e., $X_s$ and $X_t$ lie on the same trajectory. In particular, in the limit of an infinite number of steps, we precisely recover a trajectory of the underlying flow. We will happily clarify this in the camera-ready version.\n\n*Proposition 2.4 defines an $\\epsilon$ with respect to ESD, but the remark on line 201 references both ESD and LSD. Could you define an  $\\epsilon$ for LSD as well?*\n\nWe agree with the reviewer's observation that the current statement of the theorem may be misleading. In the camera-ready version, we will provide an updated version of the theorem that clearly states the assumptions that $\\mathcal{L}_b + \\mathcal{L}\\_{\\mathsf{ESD}} \\leq \\epsilon$ or $\\mathcal{L}_b + \\mathcal{L}\\_{\\mathsf{LSD}} \\leq \\epsilon$.\n\n*The notation in Equation (17) is somewhat confusing, particularly the first exponential term.*\n\nThe expression is $2 e\\times (1 + e^{2\\hat{L}})$ -- i.e., there is no first exponential term, but it is the constant $e$. We agree that as-written this expression may be misleading and we will clarify it in the revision.\n\n*In proposition 2.4, $L$ is the constant for $\\hat{v}_{t, t}$, is it for all t? For example, by taking the sup over all $t$?*\n\nYes, we assume that $\\hat{v}_{t, t}$ is uniformly Lipschitz in time $t$. We will make this statement more clear in the revision.\n\n\n## References\n[1] Salimans et al, \"Progressive Distillation for Fast Sampling of Diffusion Models.\" arXiv:2202.00512.\n\n[2] Boffi et al, \"Flow Map Matching with Stochastic Interpolants: A Mathematical Framework for Consistency Models.\" arXiv:2406.07507.\n\n[3] Sabour et al, \"Align Your Flow: Scaling Continuous-Time Flow Map Distillation\", arXiv:2506.14603."}, {"review_id": "b3VfFslqS1", "reviewer": "Reviewer_ySYP", "summary": "This paper introduces a general framework for distilling ODE-based generative models. The core proposal is a \"flow map\" framework that unifies different distillation schemes, including consistency models and progressive flow map matching (PFMM), under a single theoretical framework. The authors show that this flow map is characterized by three key conditions (Lagrangian, Eulerian, and Semigroup) which correspond to different distillation objectives.", "questions": "Please check the weakness part.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "+ The theoretical work is solid, providing a clean, unifying perspective on various distillation methods from an ODE standpoint.\n+ The framework itself is elegant and helps clarify the relationships between different techniques.\n+ The explanation of the core methodology is clear and well-written, making the theoretical contributions easy to follow.\n\n\n- Missing Proof of Proposition 2.2: The proposition is central to the paper's theory, stating that a map satisfying certain equations becomes a valid flow map. However, a proof for this claim appears to be missing. Could the authors include a proof or point to one in the literature?\n- Clarity on Novelty vs. Existing Work: The paper states that the Eulerian (ESD) and Semigroup (PSD) conditions correspond to consistency models and progressive distillation, respectively. This raises a key question: are the training objectives used for ESD and PSD identical to these existing methods? A more explicit comparison of the loss functions would help clarify the precise contribution and novelty of the proposed distillation schemes.\n- Scope of Experiments: The experiments are limited to a toy dataset and CIFAR-10. To convincingly demonstrate the effectiveness and scalability of the proposed framework, it would be beneficial to include results on more complex, higher-resolution datasets (e.g., CelebA-64, LSUN or ImageNet-64).\n- Ablation Studies: The paper proposes several training strategies, including different loss formulations (e.g., a linear combination of objectives, as mentioned in Proposition 2.3) and sampling strategies. However, the ablation studies only seem to cover the sampling strategy. \n- Writing: The connection between the proposed framework and existing methods like consistency models is a major strength. Introducing these connections earlier in the main text, rather than later, would help readers better situate the work and appreciate its relationship to the current literature from the start.", "was_revised": true, "final_justification": "I thank the authors for their detailed and comprehensive rebuttal, which has successfully addressed all of my major concerns. \nThe authors clarified the relationship between their framework and existing methods. They also provided new, promising results on CelebA-64 and will include large-scale experiments on ImageNet-256 in the final version.\nI have raised my score to reflect these improvements.", "rebuttal": "We thank the reviewer for their careful reading and evaluation of our work. Below, we address the main questions and weaknesses raised. If our arguments are found convincing, we would greatly appreciate the reviewer raising their score. We would equally appreciate further suggestions for us to improve our work, and will happily provide further clarification on any aspects of our paper or our answers.\n\n## Weaknesses\n*Missing Proof of Proposition 2.2: The proposition is central to the paper's theory, stating that a map satisfying certain equations becomes a valid flow map. However, a proof for this claim appears to be missing. Could the authors include a proof or point to one in the literature?*\n\nWe thank the reviewer for pointing out this organizational aspect of the original submission. The proof of Proposition 2.2 was split between the proofs of Propositions C.2, C.3, and C.4 in the appendix, which break down Proposition 2.2 into smaller results. We agree that this may have made it challenging to quickly reference the proof, and that this structure was not properly indicated in the text. Following the reviewer's advice, we have organized the proofs in the appendix so that Proposition 2.2 is proven in a self-contained fashion, and so that Propositions C.2, C.3, and C.4 continue to stand on their own in proving further theoretical guarantees.\n\n*Clarity on Novelty vs. Existing Work: The paper states that the Eulerian (ESD) and Semigroup (PSD) conditions correspond to consistency models and progressive distillation, respectively. This raises a key question: are the training objectives used for ESD and PSD identical to these existing methods? A more explicit comparison of the loss functions would help clarify the precise contribution and novelty of the proposed distillation schemes.*\n\nWhen we use the language \"corresponds to\", we mean that these methods are based on the underlying properties that we identify (the Eulerian equation or the semigroup condition). This is a novel contribution in its own right, as these past methods were motivated in a more *ad-hoc* fashion. Nevertheless, our contributions go beyond this additional level of theoretical clarity, and there are several key differences between consistency training, progressive distillation, and the approaches that we introduce here. In particular: \n\n1. Our methodological framework is valid for two-time flow maps defined over general stochastic interpolants, while consistency models are single-time flow maps over diffusion processes. Restricting to this setting -- i.e., choosing the interpolant to correspond to a diffusion process, choosing to estimate only the noise to data map $f_s(x) = X_{s, 1}(x)$, and using Eulerian distillation on a pre-trained model -- recovers consistency distillation [1,2].\n1. Eulerian self-distillation is new, but is related to consistency training. In consistency training, the loss function corresponds to the squared Eulerian residual where the target flow model $b$ is replaced by the derivative of the interpolant $\\dot{I}\\_s$. Mathematically, $\\mathcal{L}_{\\mathsf{CT}}(\\hat{f}) = \\int\\_0^1 \\mathbb{E}\\left[|\\partial\\_s \\hat{f}\\_s(I\\_s) + \\nabla \\hat{f}\\_s(I\\_s)\\dot{I}\\_s|^2\\right]ds$. Because $b\\_s(x) = \\mathbb{E}[\\dot{I}\\_s | I\\_s = x]$, and because the consistency training objective does not incorporate this conditioning information, it introduces an uncontrolled bias that needs to be sidestepped with complex training schemes [3]. Our approach can be viewed as a more principled way to train that leverages the implicit flow defined by the tangent condition $\\hat{v}\\_{s, s}(I\\_s)$ instead of the derivative of the interpolant $\\dot{I}\\_s$.\n1. Progressive distillation [4] is most related to the progressive flow map matching algorithm introduced in [2], which systematically converts a given flow map into a hierarchy of flow maps that can take larger steps, starting from an original velocity model. Critically, each stage requires a new neural network that learns to emulate the previous one. The PSD approach we introduce is similar, but can be viewed as a continuous variant that does not require multiple models. As a result, the objective function is different and is much simpler to use.\n\nTo make these points more clear, we have added a new section to the appendix \"Connections to existing schemes\" that significantly clarifies the relation between our approaches and these pre-existing methods. In particular, we include two subsections -- \"Distillation\" and \"Direct training\" -- which highlight how our framework reduces to, or is related to, existing methods in specific instantiations (consistency distillation, progressive distillation, shortcut models, consistency trajectory models, etc.). Our methodology is the most general framework for training consistency models that we are aware of, and one of our primary contributions is that we are the first authors to draw rigorous connections between these existing approaches from a single common viewpoint.\n    \n    \n**Scope of Experiments: The experiments are limited to a toy dataset and CIFAR-10. To convincingly demonstrate the effectiveness and scalability of the proposed framework, it would be beneficial to include results on more complex, higher-resolution datasets (e.g., CelebA-64, LSUN or ImageNet-64).**\n\nWe thank the reviewer for raising this important point. Since receiving the reviews, we have completed a round of experiments on CelebA-64, which support our hypothesis that the optimal training scheme is dataset and network-dependent. Here, similar to CIFAR-10, we find that ESD is unstable, which we hypothesize originates from the spatial Jacobian of the UNet. PSD and LSD both perform well. Unlike CIFAR-10, here we find that LSD performs the best, matching our results from the checkerboard. Numerical FID values included in an updated version of the text are shown below.\n\n| Method | 1 step | 2 step | 4 step | 8 step | 16 step |\n|--------|--------|--------|--------|--------|---------|\n|   LSD    | **15.77**       | **8.53**       |  **4.70**      |  **3.20**      | **2.66**        |\n|   PSD-M     |  17.24      |   9.34     |   5.77     |  3.97      |  3.25       | \n|   PSD-U     |  17.94      |  9.19     |   5.55     |   4.11    |   3.55      | \n\n\nWe are currrently running another round of large-scale experiments using the SiT architecture [5] on ImageNet-256, which we intend to have completed for the camera-ready version.\n\n*Ablation Studies: The paper proposes several training strategies, including different loss formulations (e.g., a linear combination of objectives, as mentioned in Proposition 2.3) and sampling strategies. However, the ablation studies only seem to cover the sampling strategy.*\n\nWe agree with the reviewer that there are several suggestions made in the text -- such as combinations of objectives, multiple models, and higher-order parameterizations -- that we did not ablate over. In this work, our aim is to introduce a systematic mathematical framework that analytically characterizes the design space of consistency models. Due to the combinatorial volume and associated computational expense of ablation studies over all factors, we are unable to include experiments testing every combination. Instead, we have studied a representative subset, including the three key choices of the objective and the sampling strategy for PSD. In an updated revision, we have also included some experiments studying the role of the stopgrad operator, where we highlight that placing a stopgrad to make the loss convex in the student typically produces the best results in high-dimensional settings. By laying out the design space completely as we have done here, researchers who wish to deploy a consistency model in their application domain of interest will have a guidebook for the hyperparameters they can study to maximize performance.\n\n*Writing: The connection between the proposed framework and existing methods like consistency models is a major strength. Introducing these connections earlier in the main text, rather than later, would help readers better situate the work and appreciate its relationship to the current literature from the start.*\n\nWe agree with the reviewer that a strength of our contribution is its connection to existing methods, and we thank them for noticing this. In an updated revision, we have included further details on these connections up front. Moreover, we have added an additional appendix that includes all necessary derivations to map our approach onto existing methodologies.\n\n\n## References\n[1] Song et al, “Consistency Models.” arXiv:2303.01469.\n\n[2] Boffi et al, \"Flow Map Matching with Stochastic Interpolants: A Mathematical Framework for Consistency Models.\" arXiv:2406.07507.\n\n[3] Lu et al, “Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models.” arXiv:2410.11081.\n\n[4] Salimans et al, \"Progressive Distillation for Fast Sampling of Diffusion Models.\" arXiv:2202.00512. \n\n[5] Ma et al, \"SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers\", arXiv:2401.08740."}, {"review_id": "DsVb0kjGiL", "reviewer": "Reviewer_1nz6", "summary": "The paper tackles the problem of training flow maps from scratch, without relying on a pretrained diffusion or flow-matching model. The authors begin with an analytical study of flow maps, deriving conditions under which a flow map is optimal. For each condition, they propose a corresponding training objective, each sharing the same global minimizer which is the optimal flow map. The proposed methods are evaluated on 2D toy datasets and the CIFAR-10 image generation benchmark, focusing primarily on the training dynamics and sample quality rather than achieving state-of-the-art performance.", "questions": "* Could the authors provide some more theoretical motivation or intuition for applying the stopgradient operation in lines 267-273? \n* Besides computational resources, what are some of the challenges to scale up these techniques to more complicated datasets?\n* Currently, most state-of-the-art diffusion models use a form of guidance (either classifier-free guidance or some of the newer alternatives such as autoguidance) during inference to achieve optimal results. Given this, many of the current distillation approaches[1,2] also distill the guidance scale into the student model. Is there a way of performing something similar when training flow maps from scratch? \n\n**References**:\n\n[1] Lu, C., & Song, Y. (2024). Simplifying, stabilizing and scaling continuous-time consistency models. arXiv preprint arXiv:2410.11081.\n\n[2] Meng, C., Rombach, R., Gao, R., Kingma, D., Ermon, S., Ho, J., & Salimans, T. (2023). On distillation of guided diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 14297-14306).", "limitations": "Yes.", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths:**\n* The paper is very well written and easy to follow. \n* The derived algorithms and objectives are clearly motivated and presented. \n* The proposed framework unifies several existing approaches, including consistency models and shortcut models. \n* The paper offers new insights and intuitions into structure, behavior, and training dynamics of flow maps. \n\n**Weaknesses:**\n* The main limitation of this work is the narrow scope of the experimental evaluation. While the paper’s focus is primarily theoretical, the practical results are somewhat underwhelming. For example, the ESD approach, which is closely related in spirit to the state-of-the-art consistency model sCM [1], fails to converge in the reported experiments. Why does this approach fail where sCM succeeds? Comparing ESD, LSD, and PSD on toy problems may not reflect their true potential, especially when architectural choices are not carefully optimized. A more thorough empirical comparison under stronger settings could yield very different conclusions.\n* I appreciate the discussion on the choice of teacher (lines 267-273), but the motivation for applying a stop-gradient remains unclear. Unlike the rest of the paper, which is grounded in clear theory, the use of stop-gradient feels arbitrary and ad hoc. Is its application dependent on the dataset? The model architecture? Does it depend on the data modality or dimensionality? Or something else? Furthermore, it would be helpful to clarify whether applying stop-gradient has any implications for the theoretical arguments made in Proposition 2.2.\n\n**References**:\n\n[1] Lu, C., & Song, Y. (2024). Simplifying, stabilizing and scaling continuous-time consistency models. arXiv preprint arXiv:2410.11081.", "was_revised": true, "final_justification": "The paper is a well written one about training few-step flowmap models. The experiments section were a bit lacking (namely toy checkerboard data and CIFAR10) which has been partially addressed by adding additional CelebA-64 experiments, with ImageNet256 ones to follow for the final version. Conditioned on the presence of the ImageNet256 experiments in the final paper, I raised my rating from $4\\to5$ and recommend accepting this paper.", "rebuttal": "We thank the reviewer for their careful reading and evaluation of our work. Below, we address the main questions and weaknesses raised. If our arguments are found convincing, we would greatly appreciate the reviewer raising their score. We would equally appreciate further suggestions for us to improve our work, and will happily provide further clarification on any aspects of our paper or our answers.\n\n## Weaknesses\n*The main limitation of this work is the narrow scope of the experimental evaluation. While the paper’s focus is primarily theoretical, the practical results are somewhat underwhelming. For example, the ESD approach, which is closely related in spirit to the state-of-the-art consistency model sCM [1], fails to converge in the reported experiments. Why does this approach fail where sCM succeeds? Comparing ESD, LSD, and PSD on toy problems may not reflect their true potential, especially when architectural choices are not carefully optimized. A more thorough empirical comparison under stronger settings could yield very different conclusions.*\n\nWe thank the reviewer for bringing up this important point, and we agree that more experiments would strengthen the contribution. To account for this, we have run a round of experiments on the CelebA-64 dataset, for which FID results are reported below.\n\n\n| Method | 1 step | 2 step | 4 step | 8 step | 16 step |\n|--------|--------|--------|--------|--------|---------|\n|   LSD    | **15.77**       | **8.53**       |  **4.70**      |  **3.20**      | **2.66**        |\n|   PSD-M     |  17.24      |   9.34     |   5.77     |  3.97      |  3.25       | \n|   PSD-U     |  17.94      |  9.19     |   5.55     |   4.11    |   3.55      | \n\n\nThe above results are similar to what we saw on CIFAR-10, with ESD demonstrating instability and LSD and PSD performing well. Unlike on CIFAR-10, in this case, we find that LSD performs the best, supporting our hypothesis that the optimal training scheme is dataset and architecture-dependent. In addition to these new results on CIFAR-10, we are currently running a round of experiments using the SiT architecture on ImageNet256 [1], which will be completed in time for the camera ready version.\n\n**Connection to consistency models.** We agree with the reviewer that ESD is similar to consistency models. However, many recent works have highlighted difficulties with training consistency models in continuous time [2,3], which is very similar to what we observe. While these instabilities can be avoided with clever, hand-designed normalization schemes and architectural modifications, our aim here is to study the performance of our three introduced training schemes with minimal engineering. We anticipate that all algorithms can be engineered significantly to maximize performance, and believe it is a very interesting direction of future work to incorporate these recent stabilization schemes into ESD, LSD, and PSD. Nevertheless, we find it particularly intriguing that LSD and PSD are stable without these engineering efforts, which suggests that they may be preferred in practice.\n\n*I appreciate the discussion on the choice of teacher (lines 267-273), but the motivation for applying a stop-gradient remains unclear. Unlike the rest of the paper, which is grounded in clear theory, the use of stop-gradient feels arbitrary and ad hoc. Is its application dependent on the dataset? The model architecture? Does it depend on the data modality or dimensionality? Or something else? Furthermore, it would be helpful to clarify whether applying stop-gradient has any implications for the theoretical arguments made in Proposition 2.2.*\n\nOur algorithms are based on the concept of self-distillation, which converts a distillation scheme into a direct training scheme by indentifying an implicit source of signal. Our aim in this section is to draw an analogy to more standard distillation schemes, in which there is a frozen teacher model that is converted into a more performant model. From this perspective, it is intuitive that it may be beneficial to leverage a stopgradient operator to \"freeze\" the implicit teacher when optimizing the student. Mathematically, one way to motivate this is that the objective functions we introduce are nonconvex in the model $\\hat{v}$ in a functional sense -- even before the parametric nonlinearity introduced by neural networks. This means that there may be spurious local minima, and optimizing the objective without using a stopgradient operator could reduce its value by changing the teacher to match the student, rather than vice-versa. In the revised paper, we significantly expanded this section to make these points more clear. Moreover, we have included some ablation studies over the choice of the stopgradient operator for LSD, PSD, and ESD on the checkerboard, CIFAR-10, and CelebA-64 datasets. In general, we find that using intuition from the distillation setting and applying a stopgradient to freeze the implicit teacher leads to the best results, though we do expect these conclusions to depend at least somewhat on the dataset and architecture. Our recommendation would be to ablate over these choices for a new dataset or problem domain to determine the right choice. We emphasize that this does not affect our theoretical arguments, which study properties of the loss value and the minimizers of the loss, The choice of stopgradient only affects the optimization dynamics.\n\nWe agree with the reviewer that stopgradients are somewhat *ad-hoc*, and we believe that making this intuition more rigorous, as well as understanding the properties of the stopgradient from a mathematical level, is a very interesting direction of future research. \n\n## Questions\n*Besides computational resources, what are some of the challenges to scale up these techniques to more complicated datasets?*\n\nThe associated issues are mostly computational in nature, due to the combinatorial number of ablations needed to study each algorithm systematically. We are particularly interested, in future work, to leverage recent stabilization schemes for training continuous-time consistency models [2,3].\n\n*Currently, most state-of-the-art diffusion models use a form of guidance (either classifier-free guidance or some of the newer alternatives such as autoguidance) during inference to achieve optimal results. Given this, many of the current distillation approaches [1,2] also distill the guidance scale into the student model. Is there a way of performing something similar when training flow maps from scratch?*\n\nWe thank the reviewer for this excellent suggestion. In general, our flow maps can be made conditional by incorporating a conditioning argument $$X\\_{s, t}(x; c) = x + (t-s)\\hat{v}\\_{s, t}(x;c)$$ for a conditioning variable $c$. The flow $v\\_{t, t}(x; c)$ then corresponds to a conditional velocity field. Given this observation, it is clear how we can defined guided flow maps, as we now describe.\n\nLet $c = \\emptyset$ correspond to unconditional generation, and let $$q\\_t(x; \\alpha, c) = b\\_t(x; \\emptyset) + \\alpha(b\\_t(x; c) - b\\_t(x; \\emptyset))$$ correspond to the guided velocity field at guidance strength $\\alpha$. This velocity field has a flow map $$X\\_{s, t}(x; \\alpha, c) = x + (t-s)v\\_{s, t}(x; \\alpha, c)$$ satisfying $v\\_{t, t}(x; \\alpha, c) = q\\_t(x; \\alpha, c)$, which may be learned via self-distillation using any of our approaches by incorporating additional random sampling over the guidance scale in the self-distillation term. We have added a paragraph titled \"Classifier-free guidance\" under the \"Algorithmic aspects\" section of our revised paper describing this approach.\n\n\n## References\n[1] Ma et al, \"SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers\", arXiv:2401.08740.\n\n[2] Le et al, \"Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models\", arXiv:2410.11081. \n\n[3] Sabour et al, \"Align Your Flow: Scaling Continuous-Time Flow Map Distillation\", arXiv:2506.14603."}]

Author rebuttal:

Meta-review: The authors propose a few-step generative modeling framework that unifies a range of existing distillation methods.

Strengths:

- The paper draws a useful design space and theoretical characterization of large-step flow maps.
- Empirically, the proposed methods provide positive results.

Weaknesses:

- The original paper only provided results on smaller datasets (reviewers oUCA, 2Efm, ysYP, 1nz6). The authors added additional experiments on CelebA, again in the form of an ablation of methods as their results on CIFAR-10.
- The proposed methods falls short of state-of-the-art methods (AC, 2Efm). The authors attribute this to the extensive compute requirements for fully training such models.

The reviewers unanimously highlight the foundational work presented in the theoretical framework, warranting acceptance. With a more extensive empirical evaluation, this could easily have been a spotlight or oral paper. The authors promised to add LSUN-Bedrooms and ImageNet-256 experiments.
Decision: Accept (poster)