Improving Consistency Models with Generator-Augmented Flows
Authors: Thibaut Issenhuth
Sangchul Lee
Ludovic Dos Santos
Jean-Yves Franceschi
Chansoo Kim
Alain Rakotomamonjy
Venue: ICML 2025 spotlightposter (2025)
Keywords: Consistency Models
Primary area: deep_learning->generative_models_and_autoencoders

Abstract: Consistency models imitate the multi-step sampling of score-based diffusion in a single forward pass of a neural network.
They can be learned in two ways: consistency distillation and consistency training. The former relies on the true velocity field of the corresponding differential equation, approximated by a pre-trained neural network.
In contrast, the latter uses a single-sample Monte Carlo estimate of this velocity field.
The related estimation error induces a discrepancy between consistency distillation and training that, we show, still holds in the continuous-time limit.
To alleviate this issue, we propose a novel flow that transports noisy data towards their corresponding outputs derived from a consistency model.
We prove that this flow reduces the previously identified discrepancy and the noise-data transport cost.
Consequently, our method not only accelerates consistency training convergence but also enhances its overall performance. The code is available at https://github.com/thibautissenhuth/consistency_GC.

Improving Consistency Models with Generator-Augmented Flows
Thibaut Issenhuth
Sangchul Lee? ? Ludovic Dos Santos
Yves Franceschi Chansoo Kim
Alain Rakotomamonjy 14
Abstract
proach is consistency models (Song et al., 2023; Song and Dhariwal, 2024). Consistency models lead to high- one-step generators, that can be trained either by distilla- tion of a pre-trained velocity field (consistency distillation), or as standalone generative models (consistency training) by approximating the velocity field through a one-sample Monte Carlo estimate.
Consistency models imitate the multi-step sam- pling of score-based diffusion in a single forward pass of neural network. They can be learned in two ways: consistency distillation and consistency training. The former relies on the true velocity field of the corresponding differential equation, approximated by a pre-trained neural network. In contrast, the latter uses a single-sample Monte Carlo estimate of this velocity field. The related estimation error induces discrepancy between consistency distillation and training that, show, still holds in the continuous-time time limit. To allevi- ate this issue, we propose a novel flow that trans- ports noisy data towards their corresponding out- puts derived from a consistency model. We prove that this flow reduces the previously identified dis- crepancy and the noise-data transport cost. Con- sequently, our method not only accelerates con- sistency training convergence but also enhances its overall performance. The code is available at: GC.
The corresponding estimation error naturally induces discrepancy between consistency distillation and training. While Song et al. (2023) hinted that it would resolve in the continuous-time limit, we show that this discrepancy per- sists both the gradients and values of the loss functions. Interestingly, this discrepancy vanishes when the difference between the target velocity field and its Monte-Carle ap- proximation approaches zero. However, this is not the case with the independent coupling (IC) between data and noise used to construct the standard estimate. It is unclear how to improve this one-sample estimate without access to the true underlying diffusion model.
The approach we adopt in this paper alleviate this issue involves altering the velocity field thereby changing the target flow to reduce the variance of its one-sample es- timator. One possible solution to this problem is to resort to optimal transport (OT) to learn on a deterministic cou- pling. OT has been succesfully adopted in diffusion (Li et al., 2024), consistency (Dou et al., 2024), and flow match- ing (Pooladian et al., 2023) models. However, due to the prohibitive cubic complexity of OT solvers (e.g. Hungarian matching algorithm),, such methods need to be applied at the minibatch level. This incurs OT approximation error (Fa- tras et al., 2021; Sommerfeld et al., 2019) and stochasticity of the data-noise coupling, thus not solving the consistency training issuc
1. Introduction
large family of diffusion (Ho et al., 2020), score-based (Song al., 2021; Karras et al., 2022), and flow models (Liu et al., 2023; Lipman et al., 2023) have emerged as state- of-the-art generative models for image generation. Since they are costly to use at inference time requiring several neural function evaluations many distillation techniques have been explored (Salimans and Ho, 2022; Meng et al., 2023; Sauer et al., 2023). One of the most remarkable ap-
In our approach, we propose to use the consistency model, assumed to be an approximation of the target diffusion flow, to construct additional trajectories. The consistency model serves as a proxy to reduce the expected deviation between the velocity field and its estimator. More precisely. from an intermediate point computed from an IC, we let the consis- tency model predict the corresponding endpoint, supposedly close to the data distribution. This predicted endpoint cou- pled the same original noise vector, defining generator-
'Critco Al Lab, Paris, France ? AI, Information and Reasoning (AI/R) Laboratory, Korea Institute of Science and Technology 3 'Il and Robot Department, University of Science and Technol- ogy, Korea 'LITIS, Univ Rouen-Normandie. Correspondence to: Thibaut Issenhuth, Chansoo Kim <t.issenhuth @ com;; eau @ as.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Improving Consistency Models with Gencrator-Anggmentee Flows
2.0
2.0
1.5
1.5
10
1.0
0.5
0.5
0.0
0.0
005
0.5
-10
1.0
-1.-
1.5
2.0
2.0
0.4 (a) PF-ODE (IC).
1.0
1.2
0.2
0.8
1.4
(b) Generator-AAgmmeeted Flows (GC).
Figure 1. Comparison of the probability flow ODE (PF-ODE) and generator-augmentt flows (GC): target data is mixture of two Dirac delta functions, and GC is computed with closed-form generator. the background, we observe the density of probability paths. White arrows ODE trajectories associated the velocity field. Blue lines sample paths from IC in (a) and from GC in (b). Trajectories start from random intermediate points *. On this example, GC sample paths appear more aligned to the velocity field.
augmented coupling (GC). We show empirically that the resulting generator-augmentte flow presents compelling properties for training consistency models, in particular re- duced deviation between the velocity field and its estimator, and decreased transport costs supported by theoretical and empirical evidence. This can be observed in Figure 1. From this, we derive practical algorithms to train consis- tency models with cenerator-augmentee flows, leading to improved performance and faster convergence compared to standard and OT-based consistency models.
Notation. We consider an empirical data distribution and a noise distribution PN (e.g. Gaussian), both defined on R We denote by g a joint distribution of samples from p. and Pz. We equip R with the dot product (x,y) and write ||x|| (x, x) 1/2 for the Euclidean norm of x. We use distance function D: Rd x Rd [0, oo) measure the distance between two points from Rd. sg denotes the stop-gradient operator.
y
In consistency models, we consider diffusion processes of the form x; x, + a,2, where x, ~ Px: Pz. and at is monotonically increasing for t [J, where T R4, We denote the distribution of x, by p(x,), simply Pe- Conditional distributions butions of x s are denoted similarly. When considering discrete formulation with N intermediate timesteps, we denote the intermediate points as Xt, X, + oi, z, where t; is strictly increasing for E {O....,N} with to and fv T. The values of a and or are chosen to be sufficiently small and large, respectively, that po p. and pr se of
Let us summarize our contributions below.
We prove that in the continuous-time limit consistency training and consistency distillation loss function con- verge to different values and we provide a closed-form expression of this discrepancy.
finite-dimensional joint distri-
We propose a novel type of flows that we denote cenerator-augmentnn flows. It relies on generator- augmented coupling (GC) that can be used train consistency model.
We provide theoretical and empirical insights into the advantages of GC. We show that generator-augmennte flows have smaller discrepancy to consistency distilla- tion than IC consistency training, and that they reduce data-noise transport costs.
2. Consistency Distillation Versus Training
In this section, we provide the required background on dif- fusion and consistency models (Sections 1. and 2.2), then discuss the discrepancy between consistency distillation and consistency training (Section 2.3) which theoretically characterize in continuous-time.
We derive practical ways to train consistency models with GC. Our approach based on joint learning strat- egy leads to faster convergence and improves the per- formance compared to the base model and OT OT-ba approaches on image generation benchmarks.
2

---

Improving Consistency Models with Gencrator-Augmenteed Flows
2.1. Flow and Score-Based Diffusion Models
where (x,, 2) is sampled from the independent coupling 91 (n(x;2)) = p. (x(x))1 Jp: (z), is an index sampled uniformly at random from {0, .. - 1}. X1,41 X, + ot Z, and x is computed by discretizing the PF-ODE with the Fuler scheme as follows:
Score-based diffusion models (Ho et al., 2020; Song et al., 2021) can generate data from noise via multi-step process consisting in numerically solving either stochastic differen- tial equation (SDE), or equivalently an ordinary differential equation (ODE). Although SDE solvers generally exhibit superior sampling quality, ODEs have desirable properties. Most notably, they define a deterministic mapping from noise to data. Recently, Liu al. (2023) and Lipman al. (2023) generalize diffusion flow models, which are de- fined by the following probability flow ODE (PF-ODE):
+ (X li+1) = Xt,41 + (4: (141) (Ktaa1)
(5)
This loss can be used to distill a score model into fo-
In the case of consistency training, Song et al. (2023) circumvent the lack of a score function by noting that /
Vi.ya (x)
x]- In light of this, its single- sample Monte Carlo estimate S is used instead in Equa- tion (5) to replace the intractable x?, by Xt, = X,, + oz, in the consistency loss:
dx= (x( dt,
(1)
where (x) E[
= x] is the velocity field. Note that x; is defined as the random variable x [xx+o,2) o.Z, and is not to be confused with the time-derivative of the ODE, Va.
dr
Lcr(O) = 2x),) Nee4 (x.
[Mo,,) sg (fo(xx,, So Jo(Xe (ta)
In the context of consistency models (Song et al., 2023; Song and Dhariwal, 2024), the most common choice is (x) = o Vx logP:(x] (x) dt, in particular the EDM formulation (Karras et al., 2022) where 0g = f and thus (x) ~ loggoe (x). Here, Vx log Pt. a.k.a. the score function, can be approximated with neural network sx(xtt) t) (Vincent, 2011; I; Song and Ermon, 2019).
(6)
2.3. Discrepancy Between Consistency Training and Distillation and Velocity Field Estimation
Naturally, replacing V, by its single-sample estimate x makes consistency training deviate from consistency dis- tillation in discrete time. Still, Song et al. (2023, Theo- rems 2 and 6) suggest that this discrepancy disappears in continuous-time since Ler (0)= = LCD + o(+/N) and the corresponding gradients are equal in some cases. This equal- ity is then used in work of Lu and Song (2024), concurrent to ours, to train continuous-ting timoo consistency models at the cost of an elaborate architectural design. Without disproving these results, we find that scaling issues and lack of gener- ality soften the claim of closed gap between consistency training and distillation.
2.2. Consistency Models
Numerically solving an ODE is costly because it requires multiple expensive evaluations of the velocity function. To alleviate this issue, Song et al. (2023) propose training a consistency model fo. which learns the output map of the PF-ODE, i.e. its flow, such that:
fo o1) = =o,
(2)
for all (xt, 0x) E Rd x [on, Ør] that belong to the trajectory of the PF-ODE ending at (x).
Indeed, we provide in the following theorem a thorough theoretical comparison of Let and LCD. We first prove that they converge to different values in the continuous-time limit. The difference captured by a regularization term that depends the discrepancy between the velocity field and its estimate. Moreover, we show that the limits of the scaled gradients do not coincide in the general case, except when the (asymptotic) quadratic loss is used. The proof, and further discussion on why this discrepancy did not appear in Song et al. (2023), can be found in Appendix A.I.
Equation (2) is equivalent to (i) enforeing the boundary condition fo(Xo, on) Xo, and (ii) ensuring that fo has the same output for any two samples of single PF-ODE trajectory- - the consistency property. (i) is naturally satisfied by the following model parametrization:
fe (Xt,, o1,)) = Cskip Cakip((r. )x1, + Cour( ) Fx ot.): (3)
where Cssip (a) of o3+(0-00)+ (o) = (a-(0-00) variance of data, and Fg is neural network. This ensures (0) 1, Coun (0) O. (ii) is achieved by minimizing the distance between the outputs of two same-trajectory consecutive samples using the consistency loss:
o
v V0 og the
Theorem (Discrepancy between consistency distillation and consistency training objectives). Assume that the distance function is given by D(x, y) = ([[|x yl) for a continuous convex function yp : [0, 0c) + 10, 20) with xlxx) ~ Caaa as x - 0+ for some C > 0 and a 2 1, and that the timesteps are equally spaced, i.e., t, = T Furthermore, assume that the Jacobian ass dx does not vanish identically. Then the following assertions hold:
(OD(0) = Eg (x)(xx))) (x..2)
(4) ((oo.) (sg (f o.)), fo f ). (
3

---

Improving Consistency Models with Gencrator-Augmentted Flows
(i) The scaled consistency losses (0) and Na Na (0) converge as N + co. Moreover, the mini- mization objectives corresponding to these limiting scaled consistency losses are not equivalent, and their difference is given by:
the target flow to reduce the variance of its one-sample estimator. This approach is reminiscent of recent work tack- ling the data-noise coupling that we discuss in the following section.
3. Reducing the Discrepancy with Data-Noise Coupling
lim No (fc - LCD(O)) = CR(), N-40
(7)
where R(0) is defined by
Beyond independent coupling (IC). From Section 2.2, it appears that x, is computed through an IC q = P* (x))p=(z) of data and noise, in a similar fashion to flow matching (Lipman et al., 2023; Kingma and Gao, 2024). Making correlated choices of data and noise beyond IC could then help align &, and (x(),, thereby resolving the discrepancy from the previous section.
R(0)= 5.¹
T
X(o,) [loer - [e0c dt (8)
and satisfies R(0)= > 0, with
0fe
0f (x,, 0I) .1). Xt:
fo = Zo (x,, 0r) + 0x
o1)
(9)
The reliance on IC in consistency and flow models is increas- ingly recognized as a limiting factor. Recent advancements suggest that improved coupling mechanisms could enhance both training efficiency and the quality of generated sam- ples in flow matching (Liu et al., 2023; Pooladian al., 2023) and diffusion models (Li et al., 2024). By reducing the variance in gradient estimation, enhanced coupling can accelerate training. Additionally, improved coupling could decrease transport costs and straighten trajectories, yield- ing better-quality samples. In a different context, ReFlow (Liu et al., 2023) leverages couplings provided by the ODE solver in a flow framework, and demonstrates that it reduced transport costs. Moreover, Lee al. (2023) propose to learn encoder from data to noise, and use this encoder as a way to construct coupling when training flow model.
Ofe X1, 01)ar + Of f (x), vx(x)). (10)
ŌCD In particular, a 2, 2,
0x
R(0)= - 1/- T
Off (x,, (i) -x 12 dt.
vi(xi)) (11) (ii) The scaled gradient N VoLlV))) and Volc(() converge as N + co. Moreover, if a + 2, then their respective limits are not identical as functions of 0:
\o
0x
lim Na-1 Vo--Vo * lim N VOCCD(O). (12) N -0000
N-000
Couplings based on optimal transport (OT) solvers. OT is a particularly appealing solution for our alignment problem. Indeed, if we consider quadratic cost and distri- butions with bounded supports, OT is no-collision trans- port map (Nurbekyan et al., 2020), i.e. X, can be sampled by a unique pair of points (x., 2). Thus X, = vx,),, im- plying R()) = in Theorem 1. Several approaches have precisely targeted the reduction of transport cost in flow and consistency models.
This theorem reveals that the optimization problems of consistency training and distillation differ not only in dis- crete time but also in continuous-time. It even highlights a discrepancy between, firstly, the limiting gradients in continuous-time - although they are equal for a = 2 - and, secondly, the gradients of the limiting losses, which differ because of R( (0), even when a = 2.
This analysis shows the importance of employing probabil- ity paths whose sample path derivatives x; aligned with the velocity field (x.)). In particular, if a diffusion process X; satisfies X, = (x(),, we have R(0) = O and equal gradients for all a 2 1. Hence, for such x,, consistency training and consistency distillation would be reconciled both in discrete time and in the continuous-time limit.
Pooladian et al. (2023) have more directly explored OT coupling within the framework of flow matching models. They show that deterministic and non-crossing paths en- abled by OT with infinite batch size lowers the variance of gradient estimators. Experimentally, they assess the efficacy of OT solvers, such as Hungarian matching and Sinkhorn algorithms, in coupling batches of noise and data points, Dou et al. (2024) have successfully adopted this approach in consistency models, while Li et al. (2024) applied OT to diffusion models. However, due to the prohibitive cubic complexity of OT solvers, OT has to be applied by minibatch for matching samples (x..z) Besides an OT approxima- tion error, this incurs the loss of the no-collision property,
However, it is unclear how to directly improve the single- sample estimation X, of v, (x.). In partieular, increasing the number of samples per point x, to reduce its variance is not tractable, as it requires sampling from the inverse diffusion process ((x,)x,). Therefore, we adopt an alternative ap- proach to alleviate the discrepancy identified in this section, which involves altering the velocity field - thereby changing
4

---

Improving Consistency Models with Gencrator-Augmentted Flows
making R()) non-zero in real use-cases. Another line of
works using OT tools with score-based models relies on the
1000
Schrödinger Bridge formulation (De Bortoli et al., 2021; Shi et al., 2023; Korotin et al., 2024; Tong et al., 2024), which has mostly proven benefits on transfer tasks.
100
R,
10
Our approach. In this paper, we use a consistency model as a proxy of the flow of diffusion process to reduce transport costs. While not fully solving the alignment issue, we will show that our method present reduced transport costs and better alignment than dedicated OT-based methods.
1e-01
IC
Batch-OT (b=512)
GC
1c-02-
4. Consistency Models with Generator-Augmentee Flows
1e-02
1e-01
10
100
Timesteps a
Figure 2. Comparison of Ric, Roach and Roc on CIFAR-10. GC exhibits lower values of this quantity for all or.
Here, we introduce our method, denoted as generator- augmented flows, which relies on generator-augmentee coupling (GC). We capitalize on the true diffusion flow f (i.e. an ideal consistency model) to map noisy points to- wards the PF-ODE solution. We present theoretical and empirical evidences that GC not only reduces the data-noise transport cost but also narrows the gap between consistency distillation and consistency training. We will discuss how to train GC consistency models jointly with f in Section 5.
Generator-augmentte trajectories satisfy the boundary conditions of diffusion processes. We note the two fol- lowing important properties of the distribution of Ry:
p(x,)= = p(xo) = Pa, p(xx) ~ p(xx) = p(orz). (16)
The first property is achieved thanks to the boundary condi- tion of the consistency model (c.f. Section 2.1), and the see- ond property by construction of the diffusion process which ensures that the noise magnitude is significantly larger than R, for large t. However, for the timesteps (0, T)) the marginal distributions p(x:) and p(x,) do not necessarily coincide.
4.1. Generator-Augmenter Coupling (GC): Definition and Training Loss
The solution proposed in this work involves harnessing the diffusion flow, computed from a consistency model, to cre- ate a novel form of coupling. The idea to leverage the properties and accumulated knowledge within an ideal con- sistency model, f. construct pairs of points. To achieve this, we first sample intermediate point, which is done as usual by sampling x. Ps and pz using the IC be- tween the two distributions, and then predict the data point Sr, via the consistency model:
4.2. Properties of Generator-Augmenter Flows
Here, we present some properties of generator-augmente flows that motivate them for training consistency models.
4.2.1. REDUCING R(0) WITH GC
(x,, 2) ~ qr: X1, x. 01,2; &r, ==g(f(xx o.)).
In Theorem 1, we proved that the continuous-time consis- tency training objective decomposes into the sum of the consistency distillation objective and regularizer term: Ccr(0) (CD + R(0). Here, we study proxy term for R(0) that is easier to calculate:
(13)
Although &t depends on the timestep t., it is important to note that it (supposedly) follows the distribution Po- This &2 is coupled with z, thereby defining our generator- augmented coupling (GC) g, which we use construct the pair of points (Xe,, ))
R, = E X vx(
(17)
(X,, z) qi X1 = X;, + Oz
-X,, + Tta, X
This quantity measures the expected distance between the true velocity field and its one-sample Monte Carlo estimate. We study R.,IC, Roact--OT and R2,GG. They are the respec- tive proxy regularizer term for each type of probability path. Note that R, depends on the endpoint predictor, a con- sistency model, which impacts both probability paths and velocity fields. Our goal is to compare those proxy regular- izer terms, in order to demonstrate that GC does lead to a
(14)
These intermediate points can serve to define a new consis- tency training loss:
Loc(() Ey(se, 2).p( Sc.,. (2)
[Noo.))D( (Dss))(sssSsssiso o..)), Sol So(
(15)
5

---

Improving Consistency Models with Generator-Augmentic Flows
Transport Cost on CIFAR-10
(())
3800
|| [xx,, 2
=2²²]
c(t) Ex 25)
(19)
Mean quadratie transport
3700
3600
- gxx 2) |||x0 2 and c(() represent the transport costs of, respectively, IC and GC. We show below, with proofs in Appendix A.3, that c(t) is decreasing for d, close to zero and for large oz.
(0) =
3500
IC
3400
GC
Batch-OT (b=128)
Lemma 1 (Transport cost of GC coupling). Assume that I is a continuously differentiable function representing the ground-truth consistency model, i.e. the flow of the PF- ODE induced by the diffusion process Xt- Define We
3300
Batch-OT (b=2048)
le
02
lc
01
10
100
Timestep a
B|z|x,] = - E|8,
xxl). Then:
Figure 3. Comparison of transport costs between IC, batch-OT, and GC CIFAR-10.
( //of Ox
..o, We, W. )
-(f)=22or3
(20)
smaller discrepancy than IC. We further motivate the use of this proxy, in regards with Theorem 1, in Appendix A.4.
In the following theorem, proved in Appendix A.2, we show that R2 decays faster for GC than for IC.
Corollary 1 (Decreasing transport cost of GC coupling
in t - 0+). There exists a t, > such that for all E [0, (.). the derivative of c(t) takes the form (() = - 26,0; with at > O. Hence for ât positive, the cost is decreasing. In particular, in the EDM setting where or = t, c(t) decreasing for small t.
Theorem 2. Assume that the data distribution contains more than a single point. Also, assume that the generator- augmented coupling between the predicted data point & and noise z computed via an ideal consistency model f, i.e., the flow of the PF-ODE. Then, + DC.
The proof of this corollary proceeds by noting that for tort= 0. the consistency model f (x, t) is an identity function, its Jacobian is an identity matrix, and thus = [[]] 13]. Using the continuity of Jacobian elements and invoking intermediate value theorem on d, concludes the proof.
Ri < RI.C.
(18)
Empirical validation. Evaluating R, requires computing the difference between the sample path derivative xc and the velocity field (x.)... In the EDM setting, this difference can be approximated using a denoiser. Indeed, X; z and (x,) = [x] = E[z|x,] = E x: x. |x/] = (x, - (x,,,) ()) with an optimal denoiser D.. The opti- mal denoiser can be approximated by a denoiser network - (x-x(x--D ())
Corollary 2 (Decreasing transport cost of GC coupling in t Ri fmax)). Assume that the consistency model f f(x, o)) is a scaling function f(x, o1) = (()
or eux. Then, we have
20,00 E[]] wrll3]. In particular, c(t) is decreas- ing whenever of is increasing.
de
Da. Finally, we have: X (x:) R
Since IC, batch-OT, and GC define different pr 's and v. ss we train different denoiser D3 for each coupling. In Figure 2, we report the results from the comparison of the three proxy terms on CIFAR-10. We observe that RG < RR,baect-OT < RR,IC and that the gap increases with f, corroborating our theoretical findings (Theorem 2).
We note that, while the assumption of the consistency model being scaling function is strong, it nonetheless bears some degree of truth for & 2 fmax. see Lemma 3 of Appendix A.
Experimental validation. As stressed in Section 3, a line of work has brought evidence that reducing the transport cost between noise and data distributions could fasten the training and help produce better samples. We compare the quadratic transport costs involved in IC, batch-OT (Poola- dian et al., 2023; Dou et al., 2024), and GC (resp. c(0), cor (0), and c(t)). Results are presented in Figure 3. Inter- estingly, GC reduces transport cost more than batch-OT on CIFAR-10 because batch-OT is tied to the batch data points x₁ whereas our computed x; are not.
4.2.2. REDUCING TRANSPORT COST WITH GC
Here, we investigate the average transport cost between the noise z P> and the predicted data point & ~ p+ as a measure of the efficiency of the data-noise coupling. Recall that the diffusion process is given by x, xx + o,z. Then, knowing that the consistency model ff satisfying the boundary condition f (Xo, 0o) = Xg, we define the function
6

---

Improving Consistency Models with Gencralor-Augmeentee Flows
128
Minimum FID of CT
128
GC with pre-trained IC
iCT-IC
GC with pre-trained weak-IC
iCT-OT
64
iCT-GC (N 0.3)
FID IC predictor
iCT-GC (N 0.5)
FID
32
FID weak-IC predictor
FID
32
iCT-GC (n 0.7)
iCT-GC 1.)
16
8
200
400
600
800
1000
200
400
600
800
1000
Training iteration (x (x1e3)
Training iteration (x (e3)
Figure Performance of GC w.r.t the performance of the predic- tor on CIFAR-10.
Figure 5. Consistency models trained with with joint learning converges faster and outperforms consistency models trained with IC or minibatch-OT on CIFAR-10.
5. Training With Generator-Augmented Flows for Image Generation
proximates the diffusion flow output. To train consistency model GC, thus rely separate endpoint pre- dictor pre-trained IC (iCT-IC): f = go (cf. Section 4.1). This network predicts the endpoint: & = g ot.). During the training of the consistency model on GC, gø is kept frozen and considered a proxy of the true flow, as our theoretical results. In Figure 4, we report the performance of consistency models on CIFAR- 10 trained with GC using two different 96: (i) 96 fully trained as standard iCT-IC with 100k training steps; (ii) a weak Is partially trained as iCT-IC with 20k training steps.
In this section, we present a methodology to train con- sistency models with GC on unconditional image gener- ation. To construct points drawn from GC trajectories (Kt,) our theory requires an optimal predictor f on intermediate points drawn from IC (x,,). Thus, this lets us two potential training strategies: (i) pre-train an IC generator, and lever- age construct GC trajectories that train GC model; (ii) joint learning strategy: train single consistency model from scratch with both types of trajectories. Note that in this second setting, the model is unique: f fa. The sec- ond option is more appealing, since it is a simple one-stage training. We demonstrate that the joint learning approach improves performance and accelerates convergence com- pared standard consistency models.
Finding 1. Using partially pre-trained and frozen end- point predictor, truined on IC trajectories, allows to train consistency model with GC and which converges faster. However, the performance of the GC model depends on the quality of the endpoint predictor on IC trajectories.
Our experiments done on the following datasets: CIFAR- 10 (Krizhevsky, 2009), ImageNet (Deng al., 2009), CelebA (Liu et al., 2015) and LSUN Church (Yu et al., 2015). For the evaluation metrics, we report the Fréchet Inception Distance (FID, Heusel et al. (2017)), Kernel Incep- tion Distance (KID, Binkkwski et al. (2018)), and Inception Score (IS, Salimans al. (2016)). Most our experiments based the improved training techniques for consis- tency models from Song and Dhariwal (2024), denoted iCT-CC. Moreover, we present some results in the setting of Easy Consistency Tuning (ECT, Geng et al. (2024)). De- tails are provided in Appendix D. The code shared in the supplementary material and will be open-sourced upon publication for reproducibility.
It is important to note that this setup is not practical, as requires pre-training standard consistency model. We aim for a training methodology that accelerates convergence and improves performance when training from scratch, without doubling the number of required training iterations.
5.2. GC from Scratch with Joint Learning
In this section, model on IC and GC trajectories from the start of the train- ing, i.e. f = sg( (fo) (ef. Section 4.1). Thereby, we combine the training of the ideal IC predictor with the training of GC model based on this predictor. We introduce a joint learning factor μc: at each training step, training pairs are drawn from GC with probability μe, while the remaining pairs are drawn from standard IC. The loss can be written on average as:
propose learn simultaneously single
5.1. GC with Pre-Trained Endpoint Predictor
Our theoretical results assume having access to an ideal generator on IC trajectories, meaning that the generator ap-
(0) = pLac (0) + (1 (1) (O)
(21)

---

Improving Consistency Models with Cencrator-Augmeented Flows
Table iCT-IC is the standard improved consisteney model with independent coupling (Song and Dhariwal, 2024); iCT-OT is iCT with minibatch optimal transport coupling (Pooladian al., 2023; Dou et al., 2024); iCT-GC (fr = 0.5) is our proposed GC with joint learning.
Dataset
Model
FID 4
KID (x (x10²) 4
IS f
iCT-IC
7.42 + 0.04
0.44 + 0.03
8.76 + 0.06
CIFAR-10
iCT-OT
6.75 + 0.04
0.36 + 0.04
8.86 + 0.09
iCT-GC (μ = 0.5)
5,95 + 0.05
0.26 + 0.02
9.10 + 0.05
iCT-IC
14.89 + 0.17
1.23 0.05
9.46 + 0.06
ImageNet (32 x 32)
iCT-OT
14.13 = 0.17
1.18 + 0.05
9.62 = 0.06
iCT-GC (μ 0.5)
13.99 = 0.28
1.13 + 0.03
9.77 = 0.07
iCT-IC
15.82 0.13
1.31 0.04
2.33 = 0.00
CelebA (64 x 64)
iCT-OT
13.63 0.13
1.09 0.03
2.40 0.01
iCT-GC (μ 0.5)
11.74 0.08
0.91 + 0.04
2.45 0.01
iCT-IC
10.58 - 0.11
0.73 0.03
1.99 0.01
LSUN Church (64 x 64)
iCT-OT
9.71 + 0.13
0.64 0.03
2.00 +.0++.011 0.01
iCT-GC (μ 0.5)
9.88 + 0.07
0.66 0.04
2.14
We denote this joint learning procedure GC (p = -). Hence, GC (μ 0) corresponds to the standard IC proce- dure, while GC (μ = 1) corresponds to training only with GC points.Note that GC (μ = 1) is not expected to work, since our theoretical guarantees assume an optimal IC pre- dictor. The detailed algorithm presented in Algorithm in Appendix. We apply this joint learning to four image datasets, and include comparisons to iCT with batch-OT (Dou et al., 2024) an additional baseline. Results across multiple datasets and metrics are presented in Table 1, and visual examples are shown in Figure in Appendix.
Table 2. Performance of IC and GC consistency models trained in ECT setting (Geng et al., 2024). Short training: 4k iterations. Long training: 100k iterations.
Model
FID &
CIFAR-10 (Short Training)
ECT-IC
7.37 + 0.05 6.41 + 0.05
Joint learning of IC and GC trajectories consis- tently improves results compared to the base IC model and outperforms batch-OT in most cases.
Finding 2.
ECT-GC (μ1 = 0.3)
CIFAR-I0 (Long Training)
ECT-IC
4.11 + 0.03
As shown in Figure 5, observe interesting interpola- tion phenomenon between recover the steady FID) improvement typical of IC training. As μ increases, the convergence of the generative model accelerates. For 0.3 < μ S 0.7, on CIFAR-10, conver- gence speed and final FID are improved compared IC and batch-OT models. For μ the FID score decreases faster than other configurations early in the training process, but it soons increases training progresses further. It is explained by the poor performance of the predictions on IC yielding deviation from the ideal IC predictor from Section 4. For the other datasets, we simply chose f 0.5 and report those results. We provide further detail on the sensitivity our results to the choice of ju Appendices C.I and D.
ECT-GC (μI 0.3)
3.74 + 0.04
and 1. For μ 0, we
FFHQ 64 64 (Short Training)
ECT-IC
13.29 + 0.10 1
ECT-GC (μ 0.3)
FFHQ x 64 (Long Training)
ECT-IC
9.68 + 0.06
ECT-GC (μ 0.3)
8.51 + 0.09
ImageNer 64 x 64 Cond. (Short Training)
ECT-IC
10.82 + 0.18
ECT-GC (μI = 0.3)
10.31 + 0.22
ImageNet 64 x 64 Cond. (Long Training)
ECT-IC
5.84 + 0.21
ECT-GC (μ = 0.3)
6.39 + 0.20
5.3. GC in the ECT Setting
As an additional experiment, we explore the recent ECT setting (Geng et al., 2024) on CIFAR-10, where consis-
8

---

Improving Consistency Models with Cenerator-Augmentted Flows
tency models are fine-tuned from a pre-trained diffusion model. This approach enables training high-quality consis- tency models in one GPU-hour, though it requires an already trained diffusion model.
cations to score-based generative modeling. In Advances in Neural Information Processing Systems, volume 34, pages 17695-17709. Curran Associates, Inc., 2021.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 248-255, 2009.
We compare IC and GC trajectories in this setting, with both short (approximately one GPU-hour) and long (100k steps, 1 GPU-day) training times. Using the referenced hyper- parameters selected by Geng et al. (2024), we observe consistent advantage for GC, with optimal j value of 0.3. These preliminary results, summarized in Table 2, align with our previous findings on the iCT setting, further supporting the effectiveness of GC.
Hongkun Dou, Junzhe Lu, Jinyang Du, Chengwei Fu, Wen
Yao, Xiao qian Chen, and Yue Deng. A unified framework for consistency generative modeling, 2024. URL https: / oopeneeview net/ forum? Sncrrm?idee d=d
Don Fallis. The epistemic threat of deepfaes.. Philosophy & Technology, 246633-633 2021.
6. Conclusion
In this paper, we identify discrepancy between consis- tency training and consistency distillation. Building this theoretical analysis, we introduce generator-augmentee flows and show that they reduce a proxy term measuring this discrepancy. Additionally, decrease the data-to-noise transport cost, as demonstrated by theory and experiments. Finally, we derive practical algorithms for training consistency models using generator- augmented flows and demonstrate improved empirical per- formance.
Kilian Fatras, Younes Zine, Szymon Majewski, Rémi Fla- mary, Rémi Gribonval, and Nicolas Courty. Minibatch optimal transport distances; analysis and applications. arXiv preprint arXiv:2101.0179: 2021.
flows
Zhengyang Geng, Ashwini Pokle, William Luo, Justin Lin, and J Zico Kolter. Consistency models made easy. arXiv preprint arXiv:2406. 14548, 2024.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. In Advances in Neural Information Pro- cessing Systems, volume 30, pages 6629-6640. Curran Associates, Inc., 2017.
Impact Statement
If
used large-scale generative models, notably in text-to- image models. this work may increase potential negative impacts of deep generative models such deepfakes (Fallis, 2021).
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, volume 33, pages 6840- 6851. Curran Associates, Inc., 2020.
Acknowledgements. This research was funded by grant Nos. 2023-00262155, 2024-00460980 and 2025-02304717 17 (IIT)) funded by the Korea government (the Ministry of Science and ICT).
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based genera- tive models. In Sanmi Koyejo, Shakir Mohamed, Alekh Agarwal, Danielle Belgrave, Kyunghyun Cho, and Alice Oh, editors, Advances Neural Information Processing Systems, volume 35, pages (8565-265777. Curran Asso- ciates, Inc., 2022.
References
Mikolaj Binkowski, Danica J Sutherland, Michael Arbel, and Arthur Gretton. Demystifying MMD GANs. In International Conference on Learning Representations, 2018
Diederik Kingma and Ruiqi Gao. Understanding diffusion objectives as the elbo with simple data augmentation. Advances in Neural Information Processing Systems, 36, 2024.
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, and Quoe V Le. Symbolic discovery of optimization algorithms. In Advances in Neu- ral Information Processing Systems, volume 36, pages 49205-49233, Curran Associates, Inc., 2023.
Alexander Korotin, Nikita Gushchin, and Evgeny Burnaev. Light Schrödinger bridge. In International Conference on Learning Representations, 2024.
Alex Krizhevsky. features from tiny images, 2009. URL https://www cs toronto.edii edu/-kriz/ learning-
Learning multiple layers
Valentin De Bortoli, James Thortoo,, Jeremy Heng, and Amaud Doucet. Diffusion Schrödinger bridge with appli-
9

---

Improving Consistency Models with Gencrator-Augmentted Flows
Sangyun Lee, Beomsu Kim, and Jong Chul Ye. Minimizing trajectory curvature of ODE-based generative models. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 Proceedings of Machine Leurning Research, pages 18957-18973. PMLR, 23-29 Jul 2023.
of the 40th International Conference on Machine Learn- ing, volume 202 of Proceedings of Machine Learning Research, pages 28100-28127. PMLR, July 2023.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models, In International Conference on Learning Representations, 2022.
Tim Salimans, Ian Goodfellow, Wojeiech Zaremba, Vicki Cheung, Alec Radford, Xi Chen, and Xi Chen. Improved techniques for training GANs. In Advances in Neural Information Processing Systems, volume 29, pages 2234- 2242. Curran Associates, Inc., 2016.
Yiheng Li, Heyang Jiang, Akio Kodaira, Masayoshi Tomizuka, Kurt Keutzer, and Chenfeng Xu. Immisci- ble diffusion: Accelerating diffusion training with noise assignment. In The Thirty-eighth Anmual Conference Neural Information Processing Systems, 2024.
Axel Sauer, Dominik Lorenz, Andreas Blattmann, and Robin Rombach. Adversarial diffusion distillation. arXiv preprint arXiv:2311. 2023.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hame Hamu, Maxim- ilian Nickel, and Matthew Le. Flow matching for genera- tive modeling. In International Conference on Learning Representations, 2023.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion Schrödinger bridge matching. In Advances in Neural Information Processing Systems, volume 36, pages 62183-62223, Curran Associates, Inc., 2023.
Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In International Conference on Learning Representations, 2023.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In IEEE Inter- national Conference on Computer Vision (ICCV), 2015.
Nicki Skafte Detlefsen, Jiri Borovec, Justus Schock, Ananya Harsh, Teddy Koker, Luca Di Liello, Daniel Stancl, Changsheng Quan, Maxim Grechkin, and William Fal- con. TorchMetrics measuring reproducibility in Py- Torch, February 2022. URL https: / / //g//gtttttb/tt github.com. Lig Lightning-AI/ http:/ttcchhtrriicc/cc
Cheng Lu and Yang Song. Simplifying, stabilizing and scal- ing continuous-time consistency models. arXiv preprint arXiv:2410.I 1081, 2024.
Max Sommerfeld, Jörn Schrieber, Yoav Zemel, and Axel Munk. Optimal transport: Fast probabilistic approxima- tion with exact solvers. Journal of Machine Learning Research, 20(105):1-23, 2019.
Chenlin Meng, Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Sali- mans. On distillation of guided diffusion models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, 2023.
Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. In International Conference on Learning Representations, 2024.
Levon Nurbekyan, Alexander Iannantuono, and Adam M Oberman. No-collision transportation maps. Journal of Scientific Computing, 82(2):45, 2020.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances Neural Information Processing Systems, volume 32, pages 11918-11930. Curran Associates, Inc., 2019.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zem- ing Lin, Natalia Gimelshein, Luca Antiga, Alban Des- maison, Andreas Köpf, Edward Yang, Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chin- tala. PyToreh: An imperative style, nigh-performance deep learning library. In Advances in Newral Informa- tion Processing Systems, volume 32, pages 8026-8037. Curran Associates, Inc., 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score- based generative modeling through stochastic differential equations. In International Conference on Learning Rep- resentations, 2021.
Yang Song, Prafulla Dhariwal, Mark Chen, and llya Sutskever. Consistency models. In Proceedings of the 40th International Conference on Machine Learning, vol- ume 202 of Proceedings of Machine Learning Research, pages 3221 122211-330 1-2222,, PMLR, July 2023.
Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky T. Q. Chen. Multisample flow matching: Straight- ening flows with minibatch couplings. Proceedings
10

---

Improving Consistency Models with Gencrator-Augmentted Flows
Alexander Y. Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. Simulation-free schrödinger bridges via score and flow matching. In International Conference on Artificial Intelligence and Statistics, pages 1279-1287. PMLR, 2024.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23 (7), 2011.
Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song, Thomas Funkhouser, and Jianxiong Xiao. LSUN: Construction of a large-scale image dataset using deep learning with humans in the loop. arXiv preprint arXiv: 2015.

---

Improving Consistency Models with Gencrator-Auggeented Flows
A. Proofs
A.1. Continuous- Time Consistency Objectives
Theorem 1 (Discrepancy between consistency distillation and consistency training objectives). Assume that the distance function is given by (xx, y) = < - -yll) for a continuous convex function yp [0, 00) - [0, oo) with (o(x)) Ca as 0+ for some C and 2 1, and that the timesteps are equally spaced, i.e., 4 Furthermore, assume that the Jacobian afa does not vanish identically. Then the following assertions hold: 8x
ii) The scaled consistency losses No LCD and No (or converge N co. Moreover, the minimization objectives corresponding to these limiting scaled consistency losses are not equivalent, and their difference is given by:
lim No [Fcr(0) (cD())) CRO)) N
(22)
where R()) is defined by
R(0) /"" {(o,)B [|ocTfell - 110cDfoll dt
(23)
and satisfies R(0) with
OcTSo = 0fe
üfo
do (xt, 01)6, +
(x,, 0z) 81
(24)
0x
0fa
0fe
Sc
(xt, 01) +
(xt, 0z) ve(xt).
(25)
do
In particular, if a = 2,
R(0)= /'' A
0fo
(x,, 01) (x,
de
(26)
Ox
(ii) The scaled gradient Na- (0) and Na-1 Vo converge as cc. Moreover, a + 2, then their respective limits are identical functions of A:
lim Volc())) # lim Vo-DooL) N
(27)
N
Proof. (i) Note that fCD.f and Ocr satisfy:
θ for = मy. fe(x
Je So (x)
@cD
(28)
Here, the second equality follows by noting that (x,, = B[&, and all the other terms in the expansion of 8 are completely determined once the value of x; is known.
on fx σ,)
Now, we use Taylor's theorem to expand the difference between fo( )aya) and fx o..) in the consistency distillation loss, Equation (4). Together with the definition of x Equation (5), this yields:
fo(Xuui 0tm)) fo(xt, Oc.)
0.o
0fe,
= do (Xirr) (ot.. - -t.)+ 0x
(ean:)(ps)) (Xti1 - x(,)) + ooltt1 - (4)
(29)
OCD fcD (ti+1 - t.) + o(t+41 - t.).
(30)
Similarly, by expanding the difference between fo (x and (x(,,, in Equation (6),
fo(Xu41, fo(Xe,)
Öfe (XL (x411,04 1): (or; -00)) +
0fe
(un1:ou11) (X86.11 x1,) + o(t:41 +t)
(31)
do
0x
-Corfel (Xt,111 or,))) (li+1 14) + o(t:41 - tz).
(32)
12

---

Improving Consistency Models with Cencrator-Auggeented Flows
Therefore, for each (CD,CT)
No (0) = No
1
AB (x on o.l (1 [[][])]] (l141 4:)
(33)
N
N
CT
A(or,)
[10 o (1 (1+) (l141 - ti)
(34)
CT
ABo 0x) dt
(35)
in the continuous-time limit as N + For simplicity of notation, we write
(0)
lim N-00
No (0)
(36)
for each {CD, CT). Then, from the formula for the limiting losses L (0). Equation (35), we immediately obtain
(or (0) CTa-1 //" X(o,) fo (x,, )l (xx, dt.
(37)
Now, we specialize in the case a and invoke the general observation that, for any random vectors x and y, the following identity holds:
E ||x| ||[x|||]|
E ||x E[x
(38)
This can be easily proved by expanding the squared Euclidean as the inner product and applying the law of iterated expectations. Plugging in di fo(x,, 0z) and y <- Xx, and noting that ocD = (OcT x. by Equation (28), it follows that
(0)= = CT {' A(
-0)
())) -cDDf(xx, 0z)1 de
(39)
= CT f Mor))
T
Of 0x
1 v.(xx)) vo(xx)))] dt.
=CT
(xx, 01) (Xe: -
(40)
Next, we establish the positivity of R().. To this end, note that || || is convex function for 2 1. By invoking the conditional Jensen's inequality, we find that the expectation inside the limiting scaled consistency training losses, Equation (35) satisfy:
E fer
E
8
(xx,,
or)
E
E
i fo (x.
or)
||| xx,
xx]
(41)
Ot
|"]
E
E
Ot fo fo σl)
E fxc,
(42)
Integrating both sides with respect to A(o;) dt, we obtain the desired inequality. The Jensen's inequality also tells that the 0te ox) (x, Exxx, xx]) 0.
equality holds precisely when @ fo (x,, o,) Ej fx(x,, o,) x,] holds, or equivalently, 3x However, given the value of x,, the quantity x; can assume an arbitrary value in Rd because the conditional density of x, diz given x, is strictly positive everywhere. Consequently, the equality condition implies äfe 0. Since this contradicts the assumption of the theorem, the strict inequality between the two limiting losses must hold.
Finally, recall that the continuous-time consistency distillation loss, (0), is given by
(0) CTa-1 /"
Ofe (x, oe), + Öfe dx (x,, 01) vi(xx) dt.
A(o,)
(43)
do
13

---

Improving Consistency Models with Generator-Augmentic Flows
Similarly, the continuous-time consisteney training loss, Ler is given by
(0) = CTa L \(
Ofe do
Ofe 0x
}
1 dt.
(xt. o,)o, +
Xt. 0t) X
xe
(44)
Since va (x;)) E|x, x;] and [[||>; ][ (x)) - follows that (0) penalizes the Jacobian ff more strongly than LED (0) does. Therefore, the two limiting consistency losses do not define equivalent objectives.
ax
(ii) Using the convexity of s,, we can show that s (x: Ca -> O- Commining this with the vector calculus formula Vylly|| = y Tyll we get Vyp(Iyy]] x Cally|1a-2y for small From this, we can estimate the gradient of the distance between sg (xx, o1,)) and fx )e4t1) with respect to the model parameter as:
V Do (s r(,)): fo(x 01411)
0fo =
= =1 |10eDool 10 (OCDJo) 00
(f3-1 1:) tz)o-1
(45)
Here, the expression ||0c f fe) fe) U5e in the square bracket is evaluated at (Xeea O(ti). Similarly, the gradient of the distance between fo Xt,; o1,) and fo (Xt:11 (2a12) is estimated as:
DoD (X o.)) (x ((
0fo T 00
(46)
= (1 + (())Ca
[1ooToll
a- (OcTFo)
(1i41 - 11)a-1
Combining these two estimates, we now compute the limit of the scaled gradient Na VoC for each as:
{CD, CT}
No Na- (oL.
N-1
Ofe 00
CaTa-2
Σ
Mar jF (1 + a(1))
(
(1++1 - 1,)
(47)
C !'"
Ofe (xe, 0z) dt
ABa,)B fx) (3.5o
(48)
00
as . + oc. Finally, if a + 2, then the term 10.foll a- a.fō is nonlinear transformation of f.fo. This nonlinearity tells that, in general,
E forefol a-2 (Derfo)
Xx] * [10cDfoll (Bcofo)T.
(49)
Therefore, the scaled gradient limits are not identical as functions of 6, and in particular, their zero sets do not coincide.
Differences with Song et al. (2023)'s results. However, Song et al. (2023) provide equivalence results between consistency training and consistency distillation. The differences come from the following reasons.
The previous theorem states discrepancy between CT and CD objectives.
In Song et al. (2023, Theorem 2), is stated that Ler = LCD + o(AT). However, in this theorem, the o(AT) is actually too large compared to the other term, and consequently the result uninformative. Indeed has two the two following problems: (i) if the distance function decays faster than the norm does, i.e., D(x, y) = o[I vl)), then the of AT") term is actually too large compared to the magnitude of the two losses N - co; (ii) The C2 -eeggularyyy assumption on the distance function D is too restrictive, excluding many cases such as the distance function given by norm. For example, such a breakdown happens when D(x, y) is a metric induced by the norm, i.c., D(x, y) = || - y|l
8. D(x, y) appearing in the proof, are undefined along = y In this case, its partial derivatives, such as 0,D(x,, y) By
In Song et al. (2023, Theorem 6), the theorem about limiting gradient equality is stated with a general distance function 'D. However, the requirements on the Hessian of the distance function restrict the theorem's 's validity where the distance function is an (asymptotic) quadratic loss. Indeed, in their proof, it turns out that the Hessian can define a non-zero value only when D is an (asymptotic) quadratic loss. This coincides with our results in the case a 2.
14

---

Improving Consistency Models with Cencrator-Augmentted Flows
A.2. Proxy of the Regularizer
In this subsection, we establish theoretical result about the decay rate of the proxy of the regularizer. As preparation for the main result and for future use, introduce simple lemma that decomposes the forward flow generated by vector field into the sum of scaling term and correction term that is well-behaved
Lemma 2. Assume that Φ is the forward flow generated by the vector field Vu, meaning that it solves the characteristic equation:
① x 01) ve(0(xx
p(x,0u) =x.
(50)
Also, assume that V, is defined as
δt (x- o(x,)) vr(x) 0t
(51)
for some function D, which call "denoiser". Then Φ satisfies the following integral equation:
(xx σ,) + 00 /" 0
d D(6(x,0,)). ds.
(52)
dt
Proof. We first compute the derivative of Φ/0,:
0
((x,01)) 01)
1 o (o(x D(o(x, -((
o x 01)
(53)
0t
o, at
(54)
Integrating both sides with respect to f, follows that
Φ(x, σ.)
Φ(x, 0o) = 1/ o²
O()(x,0,), ds.
(55)
To
Rearranging and applying the initial condition (0) x, we obtain the desired equation.
As immediate consequence of this lemma, we obtain the following result about the asymptotic structure of trained consistency model:
Lemma 3. Assume that f is the consistency model generated by a bounded denoiser D, in the sense that f solves the transport equation
of (x. ox),+ + of do
(x, (x,0)) 02) ve(x)
(56)
0x
for a vector field ve defined as in Equation (51) with the denoiser D. Then
f(x,o,)
0o O(1)
(57)
at
uniformly in x and or- The implicit bound of the error term can be chosen to be the bound of D.
Proof. Let d be the forward flow generated by v, in Lemma 2. This & is precisely the inverse of the consistency model f. in the sense that (x) x holds. Then, replacing x in the equation of Lemma with f(x os), we get
To x + To ( à ot
f(x,0,)=
o1). o.). (() ds.
(58)
o²
Now let R be such that [x, o) < for any x e Rd and noise level Then, the integral term in Equation (58) bounded as:
ao [ / D(offix.))))))))))))) ds 5o
/'
o R ds ooR
1 R.
(59)
o σ
This proves the desired claim.
15

---

Improving Consistency Models with Generator-Augmennic Flows
Now the main result, which analyzes the asymptotic behavior of R2,Ic and R2, as oo:
Theorem Assume that the data distribution contains more than a single point. Also, assume that the augguettrr-uumeeeedddd coupling between the predicted data point X, and noise z is computed via an ideal consistency model f. i.e. the flow of the PF-ODE. Then, as 00,
Rr & RI.C.
(60)
Proof. We first investigate the asymptotic behavior of Ri in the limit of L + co. Recall that the diffusion process X is given by X = x, + o, for (x.. z) ~ 91. and note that
x; - (x.)) = 0,2 - 160,,/xx
- (x. D(x, ot)),
(61)
Tt.
where oix = -[xx,xx)) is the denoiser. Plugging this into the definition of Rr,IC, we get
R2,IC
E ||x,
(62)
Now, we claim that (x,,, oe) Ex,, [[ co. Intuitively, this is because x, oz for large and oz is
independent of x, More formally, note that the conditional distribution of x, given x, is p(x, (x.) = (xx; x,, of By Bayes' theorem, the conditional distribution of x. given X, is
exp
20
Xt x,1²)
p(xx)
pxx, JR p x.. (x(()) dx,,
(63)
Ji exp
207
X
xx1 p(xx)) dx',
As -> oo, have de 00, so the exponential terms converge 1. Consequently, p(xx|xx) p(xx) and hence [xx,|x]] + [[x.] claimed. Thus,
(or)
RI.CC-
E Xx -xx,11
(64)
Since the data distribution p. is assumed to have more than one point, the variance [[[[x, [[xx]]]]] is strictly positive. Therefore, R2,c decays rate asymptotically proportional to (&x
Next, we investigate the asymptotic behavior of R., Recall the consistency training loss for GC, Equation (15). Under the assumptions in Theorem 1, the scaled loss N Loc(O) converges to
= -0) S" \(o,)L
||
05o Do (xe, or Jr + Ox Ofa (X, 0z) 0zz 0020)
axal] dt.
(65)
Here, 8, = X2 + o,2 and 8, = fxe, oz), where f is the ideal consistency model for the flow associated with the diffusion process X,. The proof of this claim similar to that of Theorem 1, so we only highlight the necessary changes. Most importantly, the velocity term is not &, but o,z. This is due to how the discrete-time samples are constructed. Indeed, from Equation (14), we find that K X = (or- of. )z, which manifests as the velocity term diz in Equation (65). Consequently, the associated (average) velocity field v, is given by
ā _ [x)))
(xt)) Blo,Z|x,)
(66)
Therefore, Re reduces to
Rr, = ( 2
|
E | 1x.
(67)
Now, unlike in the IC case, we claim that E|x [2]] As X1 as t + dc. Heuristically, this is because both x; and X2 are almost deterministic functions of z; hence, the conditioning has negligible effect the limit.
16

---

Improving Consistency Models with Cencrator-Augmennted Flows
More precisely, let Φb be the forward flow generated by the PF-ODE vector field Vt- As in the proof of Lemma 2, integrating both sides of Equation (54) from & yields
((x00)) Φxxx00) 0z)
J. d D(o(x,0,),o,) ds.
(68)
Letting u. - oo, we claim that the right-hand side converges. Indeed, the empirical data distribution p. has compact support, meaning all the data points are confined in bounded region of R Since the values D are weighted averages the data points, it follows that D also bounded. Then the integrand o (( as). Os) is absolutely integrable on [t, oo), hence the convergence follows. Moreover, the limit does not depend l. Denote this limit by p(x):
0(x,00)) pxx)
á D()(x,0,),00 ds.
(69)
01
o
As shown in the previous part, we know that D(x, c d(1) as tstt + so with c E[x.]. Then, multiplying both sides of Equation (69) by and rearranging, we have, for large 1,
oz D(o(x, o.) ds
0() = oxp(x) +
(70)
4 o
| d, o²
-xp((x) (c (())) =xp(() +cl
ds
(71)
(72)
Since & a bijection, the above relation tells that p(x) is also bijection. Next, we replace x <- x, in the equation defining
p(x), Equation (69), to obtain:
p(8.)
x.
J
o.) ds.
(73)
0t
Since is invertible, applying p- to both sides yields
z + X,
- I D(o((,, ds)
po o
x1 p
(74)
Tt tt o
X; +
X,
- X4
o [n
D(o D(ol%,.0)) o.), ds
(75)
Since all of x., Nc, and D are bounded by the largest norm of the data point, they are all finite. Hence, the last line shows that &2 = p %x + ((2)) demonstrating that X, is almost a deterministic function of x, Therefore, E[x, [ = &, as required. Consequently, RG satisfies
0o 0
Ru,oo &
(76)
This proves that R. < Rr, required.
A.3. Transport Cost
As base for the two corollaries presented the paper, we will first derive a useful representation of the derivative the transport cost.
The main purpose of the lemma is to provide more tractable representation of c'()). the time derivative of the expected transport cost. We expect c(t) decrease with t. because the predicted data point f(x,, becomes more dependent the noise z as increases. However, directly analyzing f(x,, o.)) is challenging because the dependence of f (x o²) on not explicit. Therefore, the lemma aims to:
identify quantity that better captures the dependence between z and X;;
relate c(t) to this quantity.
17

---

Improving Consistency Models with Generator-Augmenter Flows
The proof proceeds by deriving key property of the ground-truth consisteney map f:: it satisfies the transport equation,
af do x
0f 0 Ōx (x) 01) ve(x)
(77)
This equation is equivalent to saying that the conditional expectation of the time derivative of f(x,, of) given x, is zero:
f(x),
E
X1 0.
(78)
By leveraging this property, can simplify (t( into expression involving w, z E/z xel. the residual between the true noise z and its prediction given x,. This residual captures the uncertainty in predicting based on X,, allowing us to relate (t) directly to the prediction accuracy of f.
Lemma (Transport cost of GC coupling). Assume that f is continuously differentiable function representing the ground-truth consistency model, iee. the flow of the PF-ODE induced by the diffusion process xt. Define w, --|[||x] de (&t- E[[x, x.)). Then:
@f 8x
We; wx) )
d(1) 2
0x)
(79)
Proof. Note that the inverse flow f (y, ox) transports the initial point y time O along the vector field v, up time Consequently, f-1 is flow with the corresponding vector field Vy:
f (y, (y, ot)) =. (y. at
(80)
By differentiating both sides the identity - f(z (y (r,,,), o,) with respect to t and applying the above observation, we get:
á ((f="())))),
(81)
Ot
of (y)
do ( at of 1 'y.
0 of (f''(y,o)),
(82)
of ox)
of Ox (x, vx(x),
(83)
do
where the substitution f-l σr) is used in the last step. Consequently,
01 [(x,, (r),, (1) )
<(()= 2E
(84)
0j (x,, (x),0,)00 + 05 (x,, 01) xr f(xx -
2E
(85)
Do
Ox
of
2E
(xr, (x), (xx
(86)
0x
a5
2.,E
(x), (z E(z)xx)) f(x,
(87)
ix
18

---

Improving Consistency Models with Generator-Augmentte Flows
where we used the relations x, x, + oz and (x(x) E[x, x,]. Now, let w, EEz xe). Then Elw, xx] hence application of the law of iterated expectations, ( 0 for essentially any function g TRd - R. Using this, we can further simplify the last line
0f
) -2
Of Ox
(xe, (x); Wes woswe)
(ooo))))..ww wi))
d(1)= -0,
(xs, We, 0x
(88)
proving the desired equality.
An immediate consequence of this lemma is that c(t) decreases for small t:
Corollary 1 (Decreasing transport cost of GC coupling in 0+). There exists a t, such that for all t t+). the derivative of c(t) takes the form d' (t) -o, with a, 0. Hence for o, positive, the cost is decreasing. In particular, in the EDM setting where o I, <(t) is decreasing for small I.
Proof. The proof of this corollary proceeds by noting that for t 0, the consistency model f(x, t) is an identity function, its Jacobian is identity matrix leading to a, B[] w, [|2] > 0 and by assumption, all the elements of the Jacobian continuous. By continuity of a, t, exists and invoking intermediate value theorem a, concludes the proof.
The next result is the statement about the asymptotic behavior of the transport cost c(t) in the large-t regime. Corollary (Decreasing transport cost of GC coupling in t fmax). Assume that the consistency model f(x, a)) is scaling function f(x, o²) = x. Then, we have ct(t) = 26,00 B[]l w/ In particular, c(t) is decreasing whenever og is increasing.
Proof. Under the assumption, we have of 50 I. Thus, by Lemma 1, o1
d(t) 26,B (
100 Wt m]
20,00
E [[[, w: [1]
(89)
This proves that d(t) < 0 whenever o, > 0.
Toy example. Let us consider a one-dimensional toy example where x, ~ ~NO, of) with o, 2 0 and z ~ N(0, 1) are independent. Also, we assume do = 0 for the sake of simplicity. this case, the marginal law of x, is also Gaussian with Pt (00 of+ o so the vector field for the diffuusio process x, is calculated v,(x) - Vx log lo ofto?
Then, the corresponding target diffusion flow and the transport cost function are:
G. X and ctt) = of +
20.00
01)
(90)
Vo
Voz +
We note that (x)) is indeed scaling function which is asymptotically proportional d. decreasing in t for > 0.
for large I, and c(l) is
Experimental validation. We validation the transport cost decrease in Figure 6, on toy dataset composed of two 2D-Diracs, and on CIFAR-10. Interestingly, we observe that when computing OT transport plans between batches instead of on the full data, GC allows to reduce transport cost more than batch-OT.
A.4. Proxy Term
In this part, we clarify the connection between the proxy term and the in the case of the quadratic loss (a = 2). Indeed, we can bound the regularization term with the proxy term thanks to the Jacobian's 's maximum singular value Smax 8m ), which bounded as typical networks are Lipschitz:
Off (xe. oz) (Xe ve(x)))
Off
Ofe Ox )|[x ve(x:)11
5
li+e - ve(ax)112 2
(91)
Ox
Ox
19

---

Improving Consistency Models with Gencralor-Augmentted Flows
Transport Cost One-Dimensional 2-Diracs
Transport Cost on CIFAR-10
1.2
Mean quadratic transport cost
3800
Mean quudratte traneeorrt
3700
1.0
Transport cost for IC
3600
Transport cost for GC
Transport for OT
3500
IC
0.6
34000
GO
Batch OT (b=128)
3300
Batch (b-2018)
10-2
10-1
1e-02
1e-01
100
Timestep
Timestop
Figure 6. Comparison of transport costs between IC, batch-OT, and GC on two 2D-Diracs (left) and CIFAR-10 (right).
Algorithm Training of consistency models with generator-augmentt
trajectories
Input:
Randomly initialized consistency model fa, number of timesteps N, noise schedule ot, loss weighting X(-).
learning rate 7), distance function D, noise distribution Pa, joint learning parameter μ. Output: Trained consistency model fa- while not converged do
x. P+: P= (patin p(oty))
(batch of real data and noise vectors}
{sampling timesteps {makk of dimension (batch size) with each m, ~ binomial (w)
~ binomial size=batch size)
x; x. 0z,2
{IC intermediate points fendpoint prediction from the model} { {mixing IC and GC trajectories) {GC intermediate points] fconsistency loss) (update model's weights)
x, - sg(fo(x(,,or))
&, - m S:; + (1 m) X
x - kt, + oz, 8i,+ t 8t, + ot., 2
( \(o,.) ((sg))
(xr, o.)). (D
(8 otui))
0 - v/ c(0)
end while
We could also use some assumptions on f. e.g. the fact that it close a scaling function for large (see Corrolary
fxx, 0x) ae x, then we would have:
Ofe Ox
(x,, or) or)
vv(x)))
(x,)11
(92)
B. Algorithm
We present the detailed algorithm for GC (μ -) in Algorithm 1.
C. Additional Results
C.I. Ablation Studies
Understanding why GC(p 1) fails. This experiment involves training a consistency model with GC(H 1). As shown in Figure 7(a), we observe that these models converge quickly but reach saturation early in the training process. When applying the timestep scheduling method with increasing number of timesteps from Song and Dhariwal (2024), the FID of the models worsens. Using a fixed number of timesteps prevents divergence of the FID, but it still plateaus at higher FID than iCT-IC.
In Figure 7(b), we plot the FID per timestep for three model/ traje//tyyyyttory pairs: GC(p = 1)-model on IC trajectories, GC(p = 1)-model on GC trajectories, and IC-model trajectories. Notably, we observe distribution shift between IC
20

---

Improving Consistency Models with Generator-Angmentce Flows
128
iCT-IC
25
ICT-GG on GC trajectories
ICT-GC
KCT-GG IC trajectories
JCT-GC, fixed steps
ICT-IC teajectories
64
20
FID
32
F1D
1.5
16
200
100
800
1000
10
Training iteration (x1e3)
Tinesteps (a)
(a) FID of IC vs GC models during training.
(b) FID of trained IC vs GC along trajectories.
Figure 7. Analysis consistency models trained only with GC on CIFAR-IO. (a) When trained only GC trajectories, consistency models does not reach the performance of the base model (iCT-IC). In (b), we show that is linked distribution shift problem: GC models are weak on IC trajectoires, thus are sub-optimal for predicting &, required their own training (Equation (13)).
Table 3. Analysis of performance with regards to some hyper- of iCT-GC (μ1 0.5) CIFAR- 10.
Model
FID
iCT-IC
7.42 + 0.04
iCT-GC (μ 0.5) iso-time
6.38 + 0.03
iCT-GC (μ 0.5)
5.95 + 0.05
iCT-GC(p 0.5) + dropout
7.77 + 0.04
iCT-GC (n= 0.5) -MA
6.73 + 0.05
and GC trajectories: the FID the GC-model on IC trajectories degrades at the intermediate timesteps of the diffusion process. This highlights why deviating from the theory and training model exclusively on GC trajectories is insufficient: to build x,, in Equation (13), the model is inferred on IC but trained on GC trajectories. IC and GC differ too much, the model cannot improve IC.
Iso wall-clock training time. As illustrated above, consistency models trained with GC converge faster than IC. However, each training step is more time-consuming, as it necessitates forward evaluation of the consistency model without gradient computation. Regarding wall-clock training time, the computational overhead of iCT-GC is approximately 20% of the iCT-IC. top part of Table 3, we report under "iCT-GC (μ = 0.5) iso-time" the results iCT-GC (μ = 0.5) trained with the same wall-clock duration as iCT-IC. Even when considering wall-clock training time, iCT-GC (jr 0.5) is still superior iCT-IC.
We evaluate the influence of two important papper-prramterss First, the dropout in the learned model. Second, whether to use or not the EMA to compute GC endpoints x. The results are presented in the bottom part of Table 3. Interestingly, the results on dropout are opposite to those found by Song and Dhariwal (2024), since using dropout lowers the performance of iCT-GC (p- = 0.5).
Hyper-parameters.
Analysis of y on ImageNet. We present further results of the joint learning procedure with varying μ ((0.3, 0.5, 0.7, on ImageNet-32 in Figure 9, For μ (0.3, 0.5), iCT-GC outperforms the base model iCT-IC.
C.2. Visual Results
We include in Figure 8 examples of generated images for considered baselines.
21

---

Improving Consistency Models with Cencrator-Augmentted Flows
(a) Trained with IC.
(b) Trained with OT.
(c) Trained with GC (μ = 0.5).
Figure Uncurated samples from consistency models trained on CelebA 64 for fixed noise vectors. Note that models trained with generator augmeeted trajectories tend to generate sharper images.
128
Minimum FID of CT
iCT-IC
64
iCT-GC (μ = 0.3)
FID
iCT-GC (μ 0.5)
ICT-GC (μ = 0.7)
32
iCT-GC (p=1)
500
1000
1500
Training iteration (x Ic3)
Figure Results varying y for iCT-GC on ImageNet-32.
22

---

Improving Consistency Models with Gencrator-Augmentted Flows
D. Experimental Details
The code is based on the PyTorch library (Paszke et al., 2019).
Scheduling functions and hy perparaameters from Song and Dhariwal (2024). The training of consistency models heavily rely on several scheduling functions. First, there is noise schedule {ar which is chosen as in Karras et al. (2022). Precisely, o; (0o + N (ax² - - o)) with loss, chosen as X(o:) = 0++1-0 Combined with the choice of noise schedule, it emphasizes to be consistent on timesteps with low noise. Then, Song et al. (2023) propose to progressively increase the number of timesteps N during training. Song and Dhariwal (2024) argue that good choice of dicretization schedule is an exponential one: N N(k)) where K' Togg =1 J, K is the total number of training steps, i is the current training step, So (respectively $1) the initial (respectively final) number of timesteps. Finally, Song and Dhariwal (2024) propose a discrete probability distribution the timesteps which mimics the continuous probability distribution recommended in the continuous training of score-based models by Karras et al. (2022). is defined as p(o.) x erf( log(d -imen erf( log(o:) Pume)) In practice, V2PM
7. Second, there is weighting function that affects the training
2oinss2l==> ,1)
K
/2PM
Song and Dhariwal (2024) recommend using: So 10, 1280, 7, P manan 1.1, Pad 2.0.
We use the lion optimizer (Chen et al., 2023) implemented from
We have selected based on the results from Figure 5, which presents a grid search for on CIFAR-IO. Given the bell-shaped relationship observed between μ and FID, we opted to retain the best performing value identified on CIFAR-10, μ = 0.5, for all subsequent experiments (Table 1), including those on other datasets, without further tuning. Importantly, even without exhaustive hyperparameter search, our method consistently outperforms baseline approaches. This choice is validated by the ablation study presented in Appendix C.I showing similar trend for another dataset, showing that the hyper-parameter μ is easy to tune.
Selection of syper-parameter
In the ECT setting, we found that μ < 0.5 leads to improved performance, while > 0.5 can degrade final performance. Overall, we recommend setting j to small values (around 0.3) since it leads to improved performance in all our experiments.
Details on neural networks architectures. We use the NCSN+ architecture (Song et al., 2021) and follow the implemen- tation from https://github.co https://ichhbbcco///aababeee
Evaluation metrics. We report the FID, KID and IS. For the three different metrics, we rely on the implementation from TorchMetrics (Skafte Detlefsen et al., 2022). For the three different metrics, we use the standard practice (e.g. (Song and Dhariwal, 2024)) of FID which is to compare sets of 50 000 generated versus training images. Confidence intervals reported in Table are averaged on five runs by sampling new sets of training images, and new sets of generated images from the same model.
Datasets. CIFAR-10 is dataset introduced in Krizhevsky (2009). ImageNet (Deng et al., 2009), CelebA (Liu et al., 2015), and LSUN Church (Yu et al., 2015) are used respectively 32 32, 64 x 64 and 64 x 64 resolutions. We preprocess these images by resizing smaller side to the desired value, center cropping, and linearly scaling pixel values to [-1,
Details on computational ressources As mentioned in the paper, the image dataset experiments have been conducted on NVIDIA A A100 40GB GPUs.
23

---

Improving Consistency Models with Gencrator-Augmentted Flows
Table 4. Hyperparameters for CIFAR-10. Arrays indicate quantities per resolution of the UNet model. (} indicate an hyper- parameter search performed for each type model (iCT. iCT-OT, iCT-GC (μ 0.5)).
Hyperparameter
Value
batch size image resolution
512
32
training steps
100
learning rate
{0.0001, 0.000033
optimizer
lion
so
10
81
1280
7
do
0.002
01
80
network architecture
SongUNet
(from (Karras et al., 2022) implementation)
model channels
128
dropout
0.3)
num blocks
3
embedding type
positional
channel multiplicative factor attn resolutions
[, 2]
0
Table 5. Hyperparameters for CelebA and LSUN Church. Arrays indicate quantities per resolution of the UNet model. { indicate an hyper- papemmprreette search performed for each type model (ICT, iCT-OT, iCT-GC (μu 0.5)).
Hyperparameter
Value
batch size image resolution training steps learning rate optimizer
128
64
150000
0.00008
lion
8o
10
81
1280
7
do
0.002
80
network architecture
SongUNet
(from (Karras et al., 2022) implementation)
model channels
128
dropout
[0., [0., 0.2
num blocks
[3,3, 3, 5] positional
embedding type
channel multiplicative factor
[1
attn resolutions
0
24

---

Improving Consistency Models with Gencrator-Anggmentee Flows
Table Hyperparameters for ImageNet- Ik. Arrays indicate quantities per resolution of the UNet model. search performed for each type model (iCT, iCT-OT, iCT-GC (μ- = 0.5)).
{ indicate an parameer
Hyperparameter
Value
batch size image resolution
512
32
training steps
150 000
learning rate
0.00008
optimizer
lion
80
10
81
1280
7
du
0.002
Φ1
80
network architecture
SongUNet
(from (Karras et al., 2022) implementation)
model channels
128
dropout
(0., 0., 0.2] [3,5,7]
num blocks
embedding type
positional
channel mult
[,]
attn resolutions
[16]
25


Reviews:
[{"review_id": "RJ8IV59Vyl", "reviewer": "Reviewer_1gLW", "summary": "This paper examines consistency models, a technique for achieving single-step (or few-step) sampling in diffusion-based generative modeling and proposes to modify the data-noise coupling used during training so as to reduce the discrepancy between consistency training and consistency distillation. Specifically, the work introduces generator-augmented flows, which leverage an auxiliary generator (i.e., a consistency model proxy) to form “better” couplings between data points and noise. The authors show that this approach can reduce the variance in the velocity-field estimation (i.e., bridging the gap between training and distillation) and also decrease the transport cost from noise to data. \n\nThe paper also clarifies that consistency training and consistency distillation do not coincide, even in the continuous-time limit, because the Monte Carlo estimate of the velocity field used in consistency training differs from the true velocity in distillation. The authors propose generator-augmented coupling (GC), in which an already-trained or concurrently-trained consistency model generates better endpoints for the intermediate noised data. This approach yields reduced discrepancy and lower transport cost compared to the standard “independent coupling” of data and noise.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their detailed and positive assessment. We address the raised questions/weaknesses below.\n\n---\n\n> **1.** The experiments cover moderate-scale. **Could the authors add an experiment on $256\\times256$**? *[emphasis ours]*\n\nWe cannot conduct experiments on larger resolutions from scratch for computational reasons. Moreover, the only pre-trained diffusion models available in the compute-efficient ECT setting are on ImageNet $512 \\times 512$, which is also too large given our resources. \n\nStill, we were able to obtain **new results on a larger-resolution ImageNet** ($64\\times 64$ conditional) and a **new dataset** (FFHQ $64\\times 64$) in the ECT setting that we will include in Section 5.3:\n\n### ECT Results (FID) on ImageNet-64 and FFHQ-64\n\n| Model                 | FFHQ-64 (Short) | FFHQ-64 (Long) | ImageNet-64 (Short) | ImageNet-64 (Long) |\n|-----------------------|----------------|---------------|----------------------|---------------------|\n| **ECT-IC**           | 13.29 ± 0.11    | 9.68 ± 0.08   | 10.82 ± 0.18        | **5.84** ± 0.21    |\n| **ECT-GC (μ=0.3)**   | **11.73** ± 0.09 | **8.51** ± 0.12 | **10.31** ± 0.22   | 6.39 ± 0.20        |\n\n---\n\n> **2.** The concept partly assumes an “ideal” or well-trained generator for the GC approach. The paper addresses this with “joint learning,” but theoretically that is more complicated, and some of the formal results rely on the generator’s accuracy. What do you think about this?\n\nIndeed, our theoretical assumption requires an ideal generator to construct GC. Our joint learning approach allows to approximate this ideal generator and, when the generator is close enough to an ideal one, joint learning succeeds. Interestingly, we show that when the generator is not trained enough on IC, then GC does not work anymore (see Section C.1 and Figure 7 in Appendix). This proves the effectiveness of the joint learning approach: the approximation is good enough to construct GC trajectories.\n\n---\n\n> **3.** Could the approach generalize easily to **other cost functions in the coupling**—e.g., $L_1$ or robust divergences? Does the preference for quadratic cost (in the “transport cost” sense) strictly matter for performance in your experiments? *[emphasis ours]*\n\nGC is **agnostic to all distance functions**, whether the one used in consistency model training (Eq. 4/6/15) or the transport cost used in the analysis in Section 4.2.2. Regarding the transport cost, the algorithm is not built to minimize this distance directly, as opposed to batch-OT. We did not study this further but intuitively expect GC to also minimize transport cost with regards to other cost functions. \n\nWe would be happy to provide further clarification, should we have misinterpreted to the reviewer's question.", "claims_and_evidence": "The paper’s theoretical proofs and experiments are well aligned.", "methods_and_evaluation_criteria": "The authors adopt mainstream metrics (FID, KID, IS) for generative image quality, consistent with typical evaluation practice in diffusion and consistency modeling. The theoretical framework is couched in well-established techniques (Wasserstein distances, “transport cost” definitions, analyzing velocity fields in PF-ODE form).\n\nAs for the datasets and benchmarks, the chosen datasets (CIFAR-10, CelebA, ImageNet $32\\times 32$, LSUN Church $64\\times 64$) are standard, ensuring results are comparable to prior consistency and diffusion papers.", "theoretical_claims": "The derivations about the discrepancy between training and distillation (Theorem 1) are well grounded. The arguments for reduced transport cost are fairly straightforward but well explained (they rely on the fact that $c(t)$ can be shown to decrease under certain conditions).", "experimental_designs_or_analyses": "The authors train on well-known tasks with standard settings for consistency or diffusion-based generative modeling (e.g., iCT from Song & Dhariwal 2024). They measure the relevant metrics (FID, KID, IS) using recognized library code. The authors also compared their results with iCT-IC, iCT-OT and examine short vs. long training schedules.", "relation_to_broader_scientific_literature": "The paper provides a refined theoretical view on the difference between consistency training and distillation, clarifying a claim about them “coinciding” in the continuous-time limit. This clarifies or partially corrects earlier suggestions by Song et al. that the discrepancy would vanish entirely. It builds on a line of research that tries to enhance data-noise coupling in generative modeling (OT-based or reflow-based methods). Its “generator-augmented flow” is a distinctive addition, conceptually simpler than full OT yet apparently effective. The authors have cited all the major literature, including consistency modeling (Song et al. 2023; Song & Dhariwal 2024) and prior works on alternative couplings (Pooladian et al., Dou et al.).", "essential_references_not_discussed": "I don't think there are major references missing.", "other_strengths_and_weaknesses": "Strengths:\n\nThe authors systematically analyze the limiting behavior of consistency losses, which is rarely spelled out in prior literature. Generator-augmented flows are easy to implement and require no heavy OT solvers. More importantly, on standard datasets, the approach yields consistent improvements, beyond batch-OT. \n\nWeaknesses:\n\n1. The results primarily cover moderate-scale $32\\times 32$ or $64\\times 64$ tasks. Additional experiments on $256\\times 256$ or higher could strengthen claims of broad applicability. Can the authors reply to this, or add experiments on this? \n2. The concept partly assumes an “ideal” or well-trained generator for the GC approach. The paper addresses this with “joint learning,” but theoretically that is more complicated, and some of the formal results rely on the generator’s accuracy. What do you think about this? \n3. Could the approach generalize easily to other cost functions in the coupling—e.g., $L_1$ or robust divergences? Does the preference for quadratic cost (in the “transport cost” sense) strictly matter for performance in your experiments?", "other_comments_or_suggestions": "Please refer to the \"weaknesses\" section.", "questions_for_authors": "Please refer to the \"weaknesses\" section.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FAwuyopLJj", "reviewer": "Reviewer_PJfJ", "summary": "The paper analyzes and discusses a discrepancy between consistency distillation and consistency training in consistency models, and proposes a novel consistency training procedure to ameliorate the problem by leveraging the solution of the probability flow ODE learned by the model during training.\n\n## Update after rebuttal\nI confirm that the approach seems promising and well motivated. However, the fact that the authors use different training settings makes it hard to compare the experimental results with the ones from related literature. Therefore I will keep my score as is.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their thorough and constructive review. We address the raised questions/weaknesses below.\n\n---\n\n> **1.** How would you explain that using **minibatch optimal transport on LSUN slightly outperforms** the generator coupling? *[emphasis ours]*\n\nOne of our contributions is to explain why using batch-OT in consistency models is a strong baseline via our theoretical analysis (cf. Sections 2.3 and 3). This further highlights the overall better performance of GC as suggested in Section 4.2.2 and confirmed in Table 1. We do not have a clear reason why LSUN is an exception. Note however that GC and OT are nearly in the same confidence interval of the FID on this dataset, while GC is significantly superior to batch-OT on CIFAR-10 and CelebA.\n\n---\n\n> **2.** As a weakness, the reported empirical results **differ significantly from the original baseline iCT**. One problem is surely due to the lack of an open source official implementation of iCT. I would suggest the authors to improve their iCT baseline, and to **include more experiments based on ECT**, as the experiments are relatively fast and can be done with small batch size. *[emphasis ours]*\n\nThe difference in FID values stems from the following differences: the original iCT baseline uses 400k iterations with a batch size of 1024, while we use 100k iterations with a batch size of 512. Still, **the comparison of all models in our paper is fair**: they are trained and evaluated in the same setting, yielding sound conclusions. Moreover, in the compute-efficient ECT setting, we are still superior to independent coupling. As recommended, we provide **new results in the ECT setting** on a larger-resolution ImageNet and FFQH (new dataset) that we will include in Section 5.3:\n\n### ECT Results (FID)\n\n| Model                 | FFHQ-64 (Short) | FFHQ-64 (Long) | ImageNet-64 (Short) | ImageNet-64 (Long) |\n|-----------------------|----------------|---------------|----------------------|---------------------|\n| **ECT-IC**           | 13.29 ± 0.11    | 9.68 ± 0.08   | 10.82 ± 0.18        | **5.84** ± 0.21    |\n| **ECT-GC (μ=0.3)**   | **11.73** ± 0.09 | **8.51** ± 0.12 | **10.31** ± 0.22   | 6.39 ± 0.20        |\n\n\nOverall, GC is superior to IC on all datasets for iCT (4/4), and on 5/6 settings for ECT. Note that in the ECT setting, we could run the experiments with the **same training time** than in the original paper, and thus report similar FID values. This validates the benefits of using GC in practice.\n\n---\n\n> **3.** I wonder if the authors have tried **annealing the parameter from $0$ to $1$** during training. Intuitively, I would expect that at early training iterations, the model would benefit from using mostly the independent coupling, as the learned solution is probably still quite off, and then as training progresses and the sample quality improves, the model can rely more and more on the generator. *[emphasis ours]*\n\nIndeed, this is an interesting question. We did experiment with this type of annealing which resulted in lower performance than with a fixed $\\mu$. One explanation is that it does not interact well with the scheduling on the number of timesteps that is periodically increased during training: the generator is not optimal on newly introduced timesteps.\n\n---\n\n> **4.** I would suggest adding \"Minimizing Trajectory Curvature of ODE-based Generative Models\" from Sangyun Lee, Beomsu Kim, and Jong Chul Ye, to the discussion of coupling methods in section 3.\n\nWe thank the reviewer for this interesting reference which we will include in our discussion. In this work, the authors propose to learn an encoder from data to noise, and use this encoder as a way to construct a coupling when training a flow model. In comparison to our method, their algorithm is based on a joint training of the velocity field and of the noise encoder which then requires an additional network.", "claims_and_evidence": "Yes, the claims are supported by empirical evaluation i sections 4 and 5, see below.", "methods_and_evaluation_criteria": "The experiments are suited to evaluate the proposed method. In particular, the datasets are standard benchmarks for conistency models, and the baselines iCT and ECT are the two most common baselines for consistency training. However, the experimental results in iCT are not convincing as they are quite far from the results reported in the original paper, and the model has been tested only on small settings for ECT. See 'Other strengths and weaknesses'.", "theoretical_claims": "I checked the correctness of the claims in section 4 and the proofs in appendix A.", "experimental_designs_or_analyses": "The authors provide extensive evaluation of the method compared to baselines using independent coupling and minibatch OT, showing the effective advantage of using the proposed method. The analysis in figures 2 and 3 confirm the claims made in section 4.", "relation_to_broader_scientific_literature": "Consistency Models are a promising recent generative models for one-step generation. The proposed method exploits an improved training procedure, which is relevant both for advancing the understanding of the models and for practical applications.", "essential_references_not_discussed": "I would suggest adding \"Minimizing Trajectory Curvature of ODE-based Generative Models\" from Sangyun Lee, Beomsu Kim, and Jong Chul Ye, to the discussion of coupling methods in section 3.", "other_strengths_and_weaknesses": "As a weakness, the reported empirical results differ significantly from the original baseline iCT, which makes it hard to asses the validity of the proposed method. One problem is surely due to the lack of an open source official implementation of iCT, but the differences are a bit too large. However, the authors do include some experiments on ECT, of which the original implementation is open sourced, and there the results are convincing. I would suggest the authors to improve their iCT baseline, and to include more experiments based on ECT, as when using a pretrained network, the experiments are relatively fast and can be done with small batch size.\n\nA clear strength of the proposed method is the fact that it does not add any extra model parameters, and maintains the same sampling speed of equivalent baselines with independent coupling, for only a minor overhead in training time due to the extra forward pass. The model can potentially scale to much higher data dimensions and smaller batch size compared ot minibatch optimal transport.", "other_comments_or_suggestions": "I wonder if the authors have tried annealing the $\\mu$ parameter from 0 to 1 during training. Intuitively, I would expect that at early training iterations, the model would benefit from using mostly the independent coupling, as the learned solution is probably still quite off, and then as training progresses and the sample quality improves, the model can rely more and more on the generator.", "questions_for_authors": "- 1) How would you explain that using minibatch optimal transport on LSUN slightly outperforms the generator coupling? I would expect the minibatch OT coupling to work worse given the small batch size.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7gukyjYCH2", "reviewer": "Reviewer_z7zm", "summary": "# Update\n\nMy criticism is minor, and I think the authors have adequately addressed it. As I have already given a positive review, I decided not to change my evaluation.\n\n# Old Summary\n\nConsistency model is a type of generative model that sends points along the sampling trajectory of a diffusion model to the last point (i.e., a noiseless data item) on that trajectory. There are well-known two training algorithms for training consistency models. The first is consistency distillation (CD), whichi involves distilling a teacher diffusion model. The second is consistency training (CT), which does not require a teacher. Models trained with CT generally underforms those trained by CD.\n\nThe paper points out a theoretical difference betweem CD and CT, which may help explain the discrepancy in observed performance. In particular, it shows that the loss functions for CD and CT are different even at the limit of the time step size approaching zero. Moreover, when the distance function used in the losses are not the L2 distance, the gradients of the losses are different. This implies that CD and CT would likely lead to different converged models.\n\nThe paper then proposes a way to improve CT by altering the points at which the model under training are evaluated during training. Let $f_\\theta(x_t, t)$ denote the model under training. In CT, we would sample $x_* \\sim p_{\\mathrm{data}}$ and $z \\sim \\mathcal{N}(0,I)$, and a time step $t_i$. Then, we would try to minimize the distance between $f_\\theta(x_* + t_i z, t_i)$ and $f_\\theta(x_* + t_{i+1} z, t_{i+1})$. Assuming access to idealized consistency model $\\overset{\\circ}{f}(x_t, t)$, the paper proposes computing the final point $\\tilde{x} = \\overset{\\circ}{f}(x_* + t_i z, t_i)$ and then minimizing the distance between $f_\\theta(\\tilde{x} + t_i z, t_i)$ and $f_\\theta(\\tilde{x} + t_{i+1} z, t_{i+1})$ instead. This modification is called \"generator-augmented coupling\" (GC).\n\nThe paper mathematcally shows that GC has two benefits. The first is that GAF reduces a \"proxy term\" that measures the discrepancy between the loss functions of CT and CD. The second is that it reduces the expected L2 distance between a Gaussian noise $z$ and the point $f_\\theta(z, 0)$ it is mapped to.\n\nHowever, GC as just stated is useful because, if we already have the idealized consistency model, we would not be interested in training a new consistency model in the first place. The paper offers two remedies. \n\nThe first remedy is to train a consistency model with CT and use it as an approximation of $\\overset{\\circ}{f}$ in the next round of CT, but now with GC. The paper finds that this results in faster convergence in the 2nd round. Still, this approach is not practical as it makes the training process twice longer.\n\nThe second remedy is to modify CT with probability $\\mu$, CT-with-GC is used instead of the normal CT loss calculation method, and the model under training $f_\\theta$ is used to approximate $\\overset{\\circ}{f}$ when doing GC. The paper found that, with $\\mu = 0.5$ yielded models that achieved better scores (FID, KID, and IS) on the CIFAR-10, ImageNet 32x32, CelebA 64x64, and LSUN Church 64x64 datasets.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their detailed and positive assessment. We address the raised questions/weaknesses below.\n\n---\n\n> The paper claims that [$\\tilde{\\mathcal{R}}_t$] is **\"a proxy term for  $\\mathcal{R}(\\theta)$\"**, [...] but does not provide any mathematical justification for this. [...] I encourage the authors to make this connection clearer. *[emphasis ours]*\n\nWe can make this connection clearer in the case of the quadratic loss ($\\alpha=2$). Indeed, we can **bound the regularization term with the proxy term** thanks to the Jacobian's maximum singular value $s_{\\text{max}}( \\frac{\\partial f_\\theta}{\\partial x} )$, which is bounded as typical networks are Lipschitz:\n\n$ \\left\\| \\frac{\\partial f\\_\\theta}{\\partial x}(x\\_t, \\sigma_t) \\left( \\dot{x}\\_t - v_t(x\\_t) \\right) \\right\\|^2 \\leq \\left\\| \\frac{\\partial f\\_\\theta}{\\partial x} \\right\\|^2 \\|\\dot{x}\\_t - v\\_t(x\\_t)  \\|^2 \\leq s^2\\_{\\text{max}} ( \\frac{\\partial f\\_\\theta}{\\partial x} )  \\|\\dot{x}\\_t - v\\_t(x\\_t)  \\|^2 $\n\nMoreover, under  some other assumptions on $f$, for example the fact that it is close to a scaling function for large $t$ (see Corrolary 2), if $f(x, \\sigma_t) = \\frac{\\sigma_0}{\\sigma_t} x$, then we would have:\n\n$\\left\\| \\frac{\\partial f_\\theta}{\\partial x}(x_t, \\sigma_t) \\left( \\dot{x}_t - v_t(x_t) \\right) \\right\\|^2 = (\\frac{\\sigma_0}{\\sigma_t})^2 \\|\\dot{x}_t - v_t(x_t)  \\|^2$.\n\nWe will include this discussion in the next version of our paper.\n\n---\n\n**Typos.** We thank the reviewer for all the suggestions and the careful reading. We will modify the paper accordingly.", "claims_and_evidence": "Experiment results were easy to understand. Improved FID scores on 4 datasets over vanilla CT and OT coupling methods clearly show that GC has merit.", "methods_and_evaluation_criteria": "The benchmark datasets and the metrics used to evaluate trained models seem reasonable. However, using datasets with higher resolution images would make the experimental results stronger. Baselines to compare GC against (i.e., iCT-IC and iCT-OT) are reasonable as well.", "theoretical_claims": "I think the theoretical results are clearly stated. I also briefly checked the proof, and I found no glaring issues.\n\nHowever, I think there might be an issue with the use of the proxy term $\\widetilde{\\mathcal{R}}\\_t = \\mathbb{E}[\\\\| \\dot{x}_t - v_t(x_t) \\\\|^2 ]$. The paper claims that it is \"a proxy term for $\\mathcal{R}(\\theta)$\" where $\\mathcal{R}(\\theta)$ captures the discrepancy between the CT loss and the CD loss but does not provide any mathematical justification for this. Honestly, I failed to find a clear connection between $\\widetilde{\\mathcal{R}}\\_t$ and $\\mathcal{R}(\\theta)$ by myself. What I noticed is that $\\widetilde{\\mathcal{R}}\\_t$ measures the variance of target velocity values in the conditional flow matching loss [1], and its being lower is better for training a flow matching model. However, we are training consistency models with with a different loss here, and it is unclear how target variance for the conditinal flow matching loss would relate to the consistency training loss. I encourage the authors to make this connection clearer.\n\n*Citation*\n* Lipman et al. \"Flow Matching for Generative Modeling.\" 2022.", "experimental_designs_or_analyses": "I think the experiments are sound.", "relation_to_broader_scientific_literature": "The paper provides new insights on why CT underforms CD, and provides a new way to improve the CT training algorithm. The authors rightly situate their approach among methods that use more sophisticated couplings between the data distributions and the noise distributions.", "essential_references_not_discussed": "I believe the references are adequate.", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "* \"A most remarkable\" -> \"The most remarkable\"\n* \"consists in\" -> \"involves\"\n* Dash in Latex is \"---\" instead of \"--\". Moreover, there should be no spaces before and after it. Hence, \"velocity field -- thereby changing the target flow -- to reduce\" should have been \"velocity field---therby changing the target flow---to reduce\"\n* \"Two Diracs\" -> \"Two Dirac delta funcitons\" or \"Two impulses\"\n* \"$p_T \\approx p(\\sigma_T z)$\" -> \"$p_T \\approx \\mathcal{N}(0, \\sigma_T^2 I)$\"\n* \"$\\sigma_d$ the variance of the data\"  -> \"\\sigma_d is the standard deviation of the data\"\n* \"leaving the (asymptotic) quadratic loss\" -> \"except the version of the loss where the L2 distance function is used.\"\n* \"While not fully solving the alignment issue either\" -> remove \"either\"\n* \"At $\\mu = 0$ -> \"For $\\mu = 0$\"\n* \"we observe fast convergence but early divergence\" -> \"the FID score decreases faster than other configurations early in the training process, but it soons increases as training progresses further.\"", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Jk9YcCCNZx", "reviewer": "Reviewer_khfs", "summary": "The paper studies some theoretical aspects of the Consistency models (in particular, the consistency training regime). As the practical contribution, the authors propose to augment consistency training with so-called Generator-Augmented flows. The idea is to substitute the basic independent/mini-batch OT coupling used in consistency training with more sophisticated coupling which (typically) depends on the trained consistency generator itself. The authors provide some theoretical reasoning behind their methodology and conduct practical benchmarking.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive and constructive review. We address the concerns below.\n\n---\n\n> **Theoretical Claims / 1.[C]** Lines 791-793: as I understand, $\\dot{\\tilde{x}}_t=\\dot{\\sigma}_t z$?\n\n$\\dot{\\tilde{x}}_t = \\frac{d(\\hat{x}_t)}{dt} + \\dot{\\sigma}_t z \\neq \\dot{\\sigma}_t z$ where $\\hat{x}\\_t = G\\_\\theta(x\\_t, \\sigma\\_t)$. However, a pair of training GC points starts from the same $\\hat{x}_t$, which sets the velocity term to $\\dot \\sigma_t z$.\n\n---\n\n> **Experimental Designs Or Analyses / 1.[Q]** I do not completely understand the experimental setup in sections 4.2.1 and 4.2.2. What was chosen as an \"ideal consistency model”?\n\nThe ideal consistency model in all experiments of Sections 4 and 5 is a consistency model trained in a standard manner on independent coupling (iCT-IC), as described in Section 5.1. We apologize for the confusion and will clarify this point in Section 4.\n\n---\n\n> **Experimental Designs Or Analyses / 2.[C]** Figures 2, 3, x-axis. What is “timesteps $\\sigma$”?\n\nIt is the diffusion noise's standard deviation $\\sigma_t$, as in $x_t = x_0 + \\sigma_t z$.\n\n---\n\n> **Experimental Designs Or Analyses / 3.[I]** The gap in $\\tilde{R}$’s values in Figure 2 does not seem to be significant. Do we really need to fight for such an improvement?\n\nNote that on large timesteps, there is an order of magnitude of difference between IC and GC. Then, the relationship between improvement in $\\tilde{R}$’s values and improvement in final sampling quality is unknown. We hypothesize that a small improvement of $\\tilde{R}$ could lead to significant change on sampling quality, as our analysis suggests.\n\n---\n\n> **Experimental Designs Or Analyses / 4.[I]** Table 1: In the original work about iCT [Song and Dhariwal, 2024], the authors report FID values ~ 2 smaller than yours. Why such a difference? Also, how many function evaluation steps was used to obtain the FID?\n\nWe report FID with 1 NFE. The difference with iCT stems from the following differences: iCT uses 400k iterations with batch size 1024, while we use 100k with batch size 512. For fairness, all models were trained and evaluated in the same setting. Note that our ECT experiments follow the original setting from the authors, and that we present new results (see rebuttal to Reviewer PJfJ).\n\n---\n\n> **Other Strengths And Weaknesses.** Does Generator flow worsens stability issues of consistency models?\n\nWe observed instabilities in two settings: with mixed precision in the iCT setting; with large pre-trained models in the short training setting of ECT. In both cases, instabilities arose for both IC and GC.\n\n---\n\n> **Other Comments Or Suggestions.** I do not understand *why the discrepancy between CT and CD losses matters*. [...] Why is it an issue that needs to be addressed, as outlined in the abstract and contributions? [...] I expect that the main problem (it is also mentioned in the text) is the variance of $L_{CT}$ estimator. *[emphasis ours]*\n\nCD provably approximates the target diffusion flow, ensuring convergence to the data distribution. We show that CT is not equivalent to CD due to a Jacobian regularization term (Theorem 1), which prevents CT from learning the true diffusion flow. This challenges widely accepted knowledge and suggests that the generator learned by CT may not converge to the data distribution. We hypothesize that addressing this discrepancy improves performance, as confirmed by our experiments.\n\n\n---\n\n> **Questions / 1.** How do the results of Theorem 1 correspond to those by [Song et. al]. *What was lost in the analysis by [Song et. al.]* that they “missed” the discrepancy between CT and CD? *[emphasis ours]*\n\nThe difference comes from two issues:\n- in their Theorem 2 ($L_{CT} = L_{CD}+o(\\Delta T)$), the $o(\\Delta T)$ is actually too large compared to the other term, and consequently the result is uninformative;\n- their Theorem 6 (limiting gradient equality) is stated with a general distance function, but the requirements on its Hessian restrict the theorem's validity to a quadratic loss. \n\nWe will clarify those differences in the next version of the paper.\n\n---\n\n> **Questions / 2.** Theorem 1, case $\\alpha=2$. How could it be the case that limiting gradients are equal, but the objectives themselves are not, and the difference between objectives depends on $\\theta$.\n\nThe case $\\alpha=2$ is particular: objectives are not equal, and gradients of the limit loss are not equal. However, limiting gradients are equal because of the interplay between the \"stop-gradient trick\" (whose effect disappears only in continuous time and not in the limit) and the quadratic loss (cf proof of Theorem 1, L668). This equality is in some sense a coincidence.\n\n---\n\n> **Minor / 1.** What is $\\sigma_d$?\n\nAs in EDM (Karras et al., 2022) and later works, $\\sigma_d=0.5$ in our experiments. \n\n---\n\n> **Minor / 2.** It is noted that $v_t(x) = \\mathbf{E}[\\dot{x}_t | x_t =x]$. [...] Some references? \n\nThis can be seen in Liu et al. (2023), Definition D.1.", "claims_and_evidence": "Everything more-or-less ok", "methods_and_evaluation_criteria": "The practical validation methodology seems to be ok", "theoretical_claims": "I went over the proof of Theorem 1 (not very attentive) - seems to be ok\n\nI only take a glance at the proof of Theorem 2, but (seems to) understand the idea behind it.\nOther theoretical results are not carefully checked.\n \nSome questions/issues/clarifications:\n1. [C] Proof of theorem 2, lines 791-793. “Most importantly, the velocity term is not $\\dot{\\tilde{x}}_t$, but $\\dot{\\sigma}_t z$.” But, as I understand, $\\dot{\\tilde{x}}_t = \\dot{\\sigma}_t z$?", "experimental_designs_or_analyses": "I think the authors have a good suit of experiments. In particular, the authors make several theoretical claims/insights about their method and try to support these claims/insights practically. At the end of the manuscript, the authors compare their approach with baselines on several benchmark datasets. This is a good experimental design.\n\nSome questions/issues/clarifications:\n1. [Q] I do not completely understand the experimental setup in sections 4.2.1 and 4.2.2. What was chosen as an  “ideal consistency model” ($\\overset{\\circ}{f}$) in this experiment? I didn’t find experimental details.\n2. [C] Figures 2, 3, x-axis. What is  “timesteps $\\sigma$”? Is it overall time T/ delta t?\n3. [I] The gap in $\\tilde{R}$’s values in Figure 2 does not seem to be significant. Do we really need to fight for such an improvement? \n4. [I] Table 1: You report the FID metric for iCT-IC baseline. In the original work about iCT [Song and Dhariwal, 2024], the authors also provide FID score, see Table 2 in their paper, and their obtained FID  values (~2.5-3). Why do you report ~ 2 times larger FID values? Also, how many function evaluation steps (1, 2, or more) was used to obtain the FID?", "relation_to_broader_scientific_literature": "To the best of my knowledge, the idea with Generator-Augmented flow and its theoretical investigation are novel (at least, in the context of Consistency Models). \n\nRegarding the survey of Schrodinger bridge-related papers  (lines 209-211, second column), some good references:\n1. Solving schrodinger bridges via maximum likelihood. (Entropy)\n2. Diffusion bridge mixture transports, schrodinger bridge problems and generative modeling. (JMLR)\n3. Adversarial Schrödinger Bridge Matching (NeurIPS’24)", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Consistency models are known to be not very stable, especially when $\\delta t$ step approaches zero and in continuous-time limit. I expect that the utilization of Generator flow should lead to even worse stability. Do the authors experience such problems in practice?", "other_comments_or_suggestions": "To be honest, I do not understand why the discrepancy between CT and CD losses matters. Ok, the losses $L_{CT}$ and $L_{CD}$ are a bit different (e.g., in continuous time-limit), but why is it an issue that needs to be addressed, as outlined in the abstract and contributions? As I understand, the aim of CT is to learn a generative model - why does the difference between $L_{CT}$ and $L_{CD}$ cause poorer performance of CT?  I expect that the main problem (it is also mentioned in the text) is the variance of $L_{CT}$ estimator. Please elaborate more on this point.", "questions_for_authors": "1. How do the results of Theorem 1 correspond to those by [Song et. al]. What was lost in the analysis by [Song et. al.] that they “missed” the discrepancy between CT and CD? Some more elaborate discussion should be done on the comparison of the theoretical results.\n2. Theorem 1, case $\\alpha = 2$. How could it be the case that limiting gradients are equal, but the objectives themselves are not, and the difference between objectives depends on $\\theta$. \n\n[Song et. al.] Consistency model. \n\nSome misprints/minors:\n1. What is variance of data, $\\sigma_d$, lines 147-148\n2. In lines 116-117 it is noted that $v_t(x) = E[\\dot{x}_t | x_t = x]$. From which it follows? Some references?\n3. Line 254-255, second column $\\tilde{R}_{batch-OT}$ should be with $t$ subscript.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This submission provides an in-depth theoretical analysis of consistency training (CT) and distillation (CD) objectives and how they differ. It additionally proposes a new "generator-augmented" training for CT that comes with lower transport cost between data and noise and reduced discrepancy between CD and CT. While there were some minor concerns on the evaluation setup, the paper overall provides substantial insights into consistency models, which are of high interest to the community. Therefore, I strongly support accepting the submission.
Decision: Accept (spotlight poster)