VCT: Training Consistency Models with Variational Noise Coupling
Authors: Gianluigi Silvestri
Luca Ambrogioni
Chieh-Hsin Lai
Yuhta Takida
Yuki Mitsufuji
Venue: ICML 2025 poster (2025)
Keywords: Consistency Models
Generative Models
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: Improving Consistency Training with a learned data-noise coupling.

Abstract: Consistency Training (CT) has recently emerged as a strong alternative to diffusion models for image generation. However, non-distillation CT often suffers from high variance and instability, motivating ongoing research into its training dynamics. We propose Variational Consistency Training (VCT), a flexible and effective framework compatible with various forward kernels, including those in flow matching. Its key innovation is a learned noise-data coupling scheme inspired by Variational Autoencoders, where a data-dependent encoder models noise emission. This enables VCT to adaptively learn noise-to-data pairings, reducing training variance relative to the fixed, unsorted pairings in classical CT.
 Experiments on multiple image datasets demonstrate significant improvements: our method surpasses baselines, achieves state-of-the-art FID among non-distillation CT approaches on CIFAR-10, and matches SoTA performance on ImageNet 64x64 with only two sampling steps. Code is available at https://github.com/sony/vct.

VCT: Training Consistency Models with Variational Noise Coupling
Gianluigi Silvestri
Luca Ambrogioni Chieh-Hsin Lai
Yuhta Takida" Yuki Mitsufuji 34
Abstract
Independent coupling
0.8
Gooonted Dain
Consistency Training (CT) has recently emerged as a strong alternative to diffusion models for im- age generation. However, non-distillation CT often suffers from high variance and instabil- ity, motivating ongoing research into its training dynamics. We propose Variational Consistency Training (VCT), flexible and effective frame- work compatible with various forward kernels, including those flow matching. Its key inno- vation is a learned noise-data coupling scheme inspired by Variational Autoencoders, where data- data-dependent encoder models noise emission. This enables VCT to adaptively learn noise-to- data pairings, reducing training variance relative to the fixed, unsorted pairings in classical CT. Ex- periments on multiple image datasets demonstrate significant improvements: our method surpasses baselines, achieves state-of-the-art FID among non-distillation CT approaches on CIFAR- 10, and matches SoTA performance on ImageNet 64 x 64 with only two sampling steps. Code is available at https::..... : //g thub. com/ sony sonyy/vee /vct.
0.4
0.2
-0.
-044
-0.8
0.0
0.5
Variacional coupling
0.8
Nota
Gepented Dain
0.2
-02
-04
-08
0.0
-0.5 0.5
Figure I. Comparison of 1-step generation on toy data for indepen- dent and variational coupling. The data is sampled from a 2-d mix- ture Gaussians with means μ4 = (0, 0.5) and μt = (0, -0.5). show the posterior probabil-
On the bottom plot (our approach), ities learned by the encoder (in blue and green) correspouding to p(z (4) and p( 422). and their cumulative sum approximately recovers the prior distrihution. The gray lines connect the samples from the trained models the corresponding input noise. More details about the toy experiments are given in Appendix E.
1. Introduction
Generative Models are deep learning algorithms designed to learn the underlying probability distribution of given dataset, in order to then generate samples coming from such distribution. Some widely used models are Generative Adversarial Networks (Goodfellow et al., 2014), Variational Autoencoders (VAE; Kingma & Welling, 2013; Rezende et al., 2014), and Normalizing Flows (Chen al., 2018; Papamakarios et al., 2021: Kobyzev et al., 2020). More
recently, Diffiusion Models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2021b) have achieved state-of-the art (SoTA) results in several domains, including images, videos, and audio (Dhariwal & Nichol, 2021; Rombach et al., 2022; Karras et al., 2022; 2024; Ho et al., 2022; Kong et al., 2021). However, a weakness of diffusion models is the need for an iterative sampling procedure, which can require hundreds of network evaluations. Therefore, sub- stantial effort has been made to develop methods that can maintain similar generation quality while requiring fewer sampling iterations (Song et al., 2021a; et al., 2021; Salimans & Ho, 2022; Liu et al., 2022; Lu et al., 2022). Among such methods, recent and promising di-
*Work done during an internship at Sony AI OOnePlanet Re- search Center, imec-the Netherlands, Wageningen, the Nether- lands *Donders Institute for Brain, Cognition and Behaviou, Ni- jmegen, the Netherlands 'Sony AI, Tokyo, Japan "Sony Group Corporation. Correspondence
Gianluigi Silvestri <gian- luigi.silvestri @imec.nl, dianlu.silvestn @ gom'.
Martinnaaa
Proceedings of the 4nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

VCT: Training Consistency Models with Variational Noise Coupling
rection is given by Consisteney Models (CMs) (Song et al., 2023). CMs, while sharing many similarities with DMs, use different training procedure as they directly learn the probability flow equations rather than the score function.
regularized with an additional Kullback-Leibler (KL) diver- gence term, VCT enables an end-to-end training procedure compatible with various forward kernels. We demonstrate that this learned coupling effectively improves generation performance and scales well to high-dimensional data. A simple and intuitive example is shown in Figure I, where with the same settings, the model trained with the learned coupling generates samples closer to the data distribution. From this figure, we can see how the learned coupling par- titions the data differently from how the standard CT does, effectively changing the frrm of the underlying ODE, likely resulting in an easier training objective as the prior noise is partitioned according to the learned the data-dependent coupling distribution.
CMs can be either trained by distilling the ODE trajectories of a pre-trained diffusion model (Consistency Distillation, CD), or completely from scratch through a bootstrap loss (Consistency Training, CT), which results in a novel genera- tive modeling framework. However, the CT objective can be subject to high variance, making it difficult to train. A follow-up work (Song& & Dhariwal, 2024) analyses the train- ing dynamics of CMs and proposes several improvements which result in a more stable CT procedure, achieving SoTA results in few-step image generation. Since then, several works have proposed additional strategies further improve CT (Geng et al., 2025; Wang et al., 2024; Lee et al., 2025; Yang et al., 2024).
The reminder of the paper organized as follows: we first discuss relevant related CT methods and flow-based works employing coupling strategies. We then formulate CT from the Flow Matching perspective, which is generalization of the diffusion framework and it offers a more natural way to introduce the noise-coupling distribution. Finally, we describe our method, deriving similarities with Variational Autoencoders, and report our experimental results on com- mon image benchmarks.
A possible source of instability of CT training comes from the fact that different noise masks are applied to the same data point, creating ambiguity the target corresponding to the given noisy state, especially early during training with coarse discretization steps. From a more mathematical per- spective, training can be destabilized by sharp boundaries in the ODE flow mapping, which can be hard to learn for the model and give raise to high variance in the stochastic gradient estimator. For example, the ODE flow for mixture of delta distributions is defined by a tessellation of the initial noise space, with discontinuities along all the borders. Since the standard CT approach with fixed forward process cannot alter the target ODE flow, the optimum of the standard CT training can potentially be highly singular. The existence of these singularities depend on topological reasons (i.e. non-injectivity of the ODE flow mapping at t + 0 (Cornish et al., 2020)). However, the issue can likely be ameliorated by altering the forward process during training, which can be used to change the location of the singularities.
2. Related Work
Since the introduction of Consistency Models in (Song et al., 2023; Song & Dhariwal, 2024), several strategies have been proposed to improve training stability. The work from (Geng et al., 2025) proposes Easy Consistency Models (ECM) novel training strategy where time steps are sampled in continuous fashion and the discretization step is adjusted during training, as opposed to the discrete time grid used in iCT. It further shows the benefits of initializing the net- work weights with the ones from a pretrained score model, achieving superior performance with smaller training bud- get. (Wang et al., 2024) builds on top of ECM, introducing additional improvements and framing consistency training as value estimation in Temporal Difference learning (Sutton, 2018). Truncated consistency models, introduced in (Lee et al., 2025), proposes to add a second training stage on top of ECM, to allow the model to focus its capacity on the later time steps, resulting in improved few-steps generation performance. Other recent contributions to the consistency model literature are works such as (Kim et al., 2024; Heek et al., 2024) where the focus is on improving multistep sam- ple quality, (Lee et al., 2024) which trains model with both consistency and score loss to reduce variance, and (Lu & Song, 2025), which proposes several improvements to the continuous-timg training of consistency models. Our work can be seen as parallel contribution to the aforementioned methods, as we focus on learning the data-noise coupling, which can be used as drop-in replacement to the standard
A way implement this approach is to sample noise using conditional coupling function between data and noise. The concept of using a coupling function to reduce variance during training was successfully used in Flow Matching, with works such as (Pooladian et al., 2023; Tong et al., 2024; Lee et al., 2023), with the main objective of obtaining straighter ODE trajectories for faster sampling. Forms of coupling in CMs were proposed in works such as (Dou et al., 2024; Issenhuth et al., 2025), but their formulations do not match the performance of standard CMs.
In this work, we propose Variational Consistency Training (VCT), which introduces a variational formulation for train- ing the forward transition kernel via a coupling function between data and noise, a technique we term Variational Coupling (VC). This yields a loss function analogous that of VAEs. By learning data-dependent noise distribution
2

---

VCT: Training Consistency Models with Variational Noise Coupling
independent coupling.
et al., 2023), and 1/e (26; #1) 0; = 2 + la1 in (Kar- ras et al., 2022). This conditional flow is analogous the formal solution kernel of the forward process at time i in the generative diffusion framework. While it is difficult to di- rectly obtain the flow function 4 from the conditional flow wh (zy; T1), it is possible to give a formal expression for the resulting vector field:
There are several works showing the benefit of using cou- pling in Flow Matching (Pooladian et al., 2023; Tong et al., 2024; Liu et al., 2023; Lee et al., 2023; Albergo et al., 2024; Kim et al., 2025) and Diffusion Models (Bartosh et al., 2024a;b; Nielsen et al., 2024). Among these, (Lee et al., 2023) shares the most similarities with our method, as they also use an encoder to learn a probability distribution over the noise conditioned on the data. Their method results in improved performance compared to equívalent Flow Match- ing models, while requiring less function evaluations. Our method consists of a similar procedure but applied to CT, resulting in improved few-steps generation performance and confirming the effectiveness of learning the data-noise cou- pling. A different coupling strategy for CT is proposed in (Issenhuth et al., 2025), where the data-noise coupling is extracted directly from the prediction of the consistency model during training. Compared to our method, they do not need the additional encoder to learn the coupling, but their generator-induced coupling needs to be alternated with the standard independent coupling to avoid instabilities. The Flow Matching formulation in Consistency Models with lin- ear interpolation kernel was previously used in (Dou al., 2024; Yang et al., 2024), where the former also explores the use of minibatch OT coupling, while the latter trains the model to learn the velocity field and adds a regularization term to enforce constant velocity. In our work, we use the Flow Matching formulation, but keep most of the CT build- ing blocks, resulting simpler formulation with superior performance.
Ut (zt) = Eng (x1;
(I)
where ( 21) = di
d 4/, (xp; 21) is the vector field that
generates the conditional flow (2) In the case of the simple interpolation conditional flows (1((
(1 ()a0 + far,, the formula specializes follows:
(x,) Ez, my [x| -xol
(2)
Readers who are familiar with generative diffusion will immediately recognize that this expression is directly related
to the standard expression for the score function. From this connection, it is clear that the conditional vector field can be estimated with a regression objective
=r, [ x 1), ur uu (1b. 2));; 11
3.1. Noise coupling
An advantage of the Flow Matching formalism over SDE diffusion is that, as shown in (Pooladian et al., 2023; Tong et al., 2024), it is straightforward to introduce a probabilistic coupling (x01 xo) between the data and the noise distribu- tion. In this case, we require that f <<x1 xo do)d should follow standard normal distribution, at least approximately.
3. Background
The use of non-trivial noise coupling does not alter the form of the conditional velocity fields un (2y; x1) as far as the coupling is time-independent. However, it does alter the total velocity field (x:)) since it affects the conditional distribution p(æ1, æx). which determines the expectation in Eq. 1.
Flow Matching provides a general framework that gen- eralizes diffusion and score-matching models (Lipman et al., 2023; Albergo & Vanden-Eijnden, 2023). In the Flow Matching formalism, a deterministic flow function Sb, with initial condition X, is used to build interpolat- ing map between two distributions abs (3t) It, such that the data distribution Po (x is mapped into a distribution Pi ( commonly chosen to be Gaussian noise distribution Pi ( == WZ 0, I), by the pushforward operator. From this quantity, we can define the vector field (xx.)) as the infinitesimal generator of th;
4. Continuous consistency models from a Flow Matching perspective
As explained above, the flow function (x0) maps a noise- less state 2g to the noisy state æ;, Its inverse v/ (ax) can then be interpreted as denoiser, as maps each noisy state to a uniquely defined noiseless state ap- This function is often referred to as consistency map, and it follows the identity:
d
dt (xo) ux(x)) = (1b
In a diffusion model, the flow (2o) is the inverse of the probability ODE flow determined by the forward SDE. In- stead, in standard Flow Matching, the mapping is specified as a conditional flow (20o 21), which is typically taken as a simple linear interpolation between samples from the two densities po( and P1 (æ1). Some common examples are (Z 21) ) = Ft = (1 - t) + tay as seen in (Lipman
d Ti (15,(()) = 0, on tE [0, 1].
(3)
together with the boundary condition क, (xx) = zg- Eq. 3 is a consequence of the fact that all noisy states in an ODE trajectory a/e (2ro) share the same initial point æo. which
3

---

VCT: Training Consistency Models with Variational Noise Coupling
implies that s5, (zr)) is constant along the trajectory. This property can be used to define continuous loss for network fo (ace, t), trained to approximate the inverse flow vo (xz):
where #1 We are sampled according to the noise-coupling r(2 mo)) In this expression, O denotes frozen copy of the parameters which does not require gradients. This loss is in fact unbiased for At - 0, as it was shown in (Song et al., 2023). The boundary condition can be enforced through the parametrization introduced in (Karras et al., 2022):
com (0) Exzo L( x(t)
d dt fo(()(()))))) 8)
| de (4)
fa () Cakip (4)ac + Cout (o)Fo (2:
where A(t) is positive-valued function that weights the loss for different time points. This loss should be used together with the identity boundary condition foo. 0) = Zo. which we will discuss later. In distillation training, the determin- istic trajectories x; (20) are obtained by integrating the ODE flow obtained from pre-trained diffusion flow matching model. Alternatively, the consistency network can be trained directly by re-writing the total derivative in terms of the conditional flow:
where F6 is a neural network and Cakip and Cout are specified such that Cskip (0) = and Cout (0) = 0.
6. Consistency Models with Variational Coupling
Our method consists in learning a conditional coupling (0(x1 xo) with a neural network gø(x) parametrized by Φ. which we refer to as the encoder in the following given its analogy with VAEs. During training, we can sample noise conditionally from π (# mo) = (0 Po (2a)) instead of the independent noise commonly used in CT, and obtain noisy states for given time step t as follows:
d df 4/ (vb)(z))) V4s (ax) dr
+ deeb (xx))
-v (xx))nrlan ( x1)) |1 + 0,05, (2))
=
(5)
-Bmjla (vr," (ax (z) x1) + 0, (xx)]
Dt = (30 211);
[ の, (sby (20: (Too:)))) = -B
dt
x ~ (x)
ma) = N (2(:: (²
From this equality, together with the fact that the squared Euclidean norm is convex function, it follows that
where we express the corresponding coupled noise æn using the Gaussian reparameterization formula:
Ctot (0) S roond (0). with
I1 = -g((zo) + (6( e, C~N(E
L'' A(t)
dt where we moved the expectation outside of the squared norm using Jensen's inequality. Therefore, we can optimize the tractable "conditional loss" Lcond (0) instead of Cox (f), which contains the unknown flow function Tb, (xx)).
(oond (0) E,r,, ant
de so (ab m1). (1))
Here, we restricted our attention to linear forward models of the form thy x2) = ag + byag, which encompasses most models used in the diffusion and flow-matching liter- ature. Moreover (2o) and gg (zo) denote the mean and scale output of the encoder, which define the signal-noics coupling (see Appendix F for a visual representation). Both gy (2o) and gg(2)) preserve the same dimensionality of the input signal. The encoder network that produces the coupling distribution (x xa) can be trained end-to and-to-end alongside the consistency model, as shown in Algorithm 1. This results in joint optimization where the consistency network adjusts its constancy to minimize its total derivative along the trajectories while the encoder implicitly moves the trajectories towards the space of constancy of the model. In fact, the velocity field u, (a:)) depends on the coupling, since πT π( xx) affects the expectation in Eq. (1). This formulation results in viable generative model as long as the noise at time remains approximately P1 (x1), since se- vere deviation from the prior induced by the coupling would result in improper initialization for the sampling procedure and consequently in reduced sample quality. We therefore add a KL divergence as a regularizer. DKL (gol Ip)), result- ing in a loss resembling the Evidence Lower Bound loss of
5. Discretized consistency training
The continuous loss can be directly minimized in expecta- tion by sampling the time from uniform distribution and by computing the total derivative d ( x)), t) by automatic differentiation. However, in practice it is often more convenient to instead optimize a time-discretized loss with a finite difference approximation for the total deriva- tive:
di
Coond (0)
Σ
[I]
with
Afe fo x1), - fo (1h1, (xo: t.).
4

---

VCT: Training Consistency Models with Variational Noise Coupling
Algorithm I Variational Consistency Training (VCT)
Given that the KL terms in Eq. and the VAE loss identical, our loss function serves the loss of VAE with encoder (g (æo), 9g (2o)) and prior xx1 O, I). Since the VAE loss represents an evi- dence lower bound, it follows that the consistency loss also provides a lower bound on the model evidence:
Input: data distribution pdatas initial model parameter 0, initial encoder parameter Φ, learning rate 77. EMA rate
upper bound on
14, distance function d(.;:) ,), consistency weighting (-)(,),
KL weighting ÀKI.
θEMA + 0, ΦEMA & Φ and h +- 0
repeat
Proposition 6.1. The following upper bound negative log dendinstty holds:
Sample 2o ~ Pdata= ~ p(t), = At
Sample E (0,
I1 f (20) + g (2O)E
lloggn (20) S
1
Bug (20 (2o) (20)
fo( 1)12+ LkI(
t 0, + b
2a²
Dr t t0 + b
1
d E
So 10 /|] dt (kL(
Case (0, x) t da >a fo (at, 1), fo-((r,
$ Ju
20
dt
(KI( + DKL(N
) g (2())) (T)
C(0, ()) + faod (0, Φ) (o)
where LKL (Φ) == DKL (2 xx) ) (x
0 4 Voo[( ②))
Φ Φ (V.((0,
This proposition also establishes the connection between the minimization objective of CT-VC and that of the continuous- time CM as N + co (proof in Appendix B).
0EMA+ < stopgrad(uEEMA + (1-1100
ΦEMA f stopgrad(UgTMA + (1-1)))
k + h: +1
Compared to traditional VAEs, our method can be viewed as a time-dependent modification where the transition kernel smoothly interpolates between delta distributions centered at datapoints and a Gaussian distribution.
until convergence
Variational Autoencoders (Kingma& & Welling, 2013):
Ż(0, p) dise (0, x)
Varying 3. In both our model and VAE, the latent vec- tor needs to approximately follow normal distribution to avoid deviating from the prior. However, previous stud- ies (Hofffman & Johnson, 2016; Rosca et al., 2018; Aneja et al., 2021) have observed that VAE's aggregated posterior fails to match the prior. The same problem could occur in our model without additional tricks (see Sec. 7.2). To mitigate this prior-posterior mismatch, we introduce scalar hyperparameter 3 to control the strength of the KL regular- ization. This was first introduced by Higgins et al. (A-VAF 2022) for different purpose (inducing disentangled latent representation) in the VAE context.
+ Ea, (g (V))) I)].
While using an encoder to learn the noiae coupling requires additional computation during training, we empiri- cally find that relatively small encoder is enough learn an effective coupling, which results only in a minor increase of training time (see Appendix C). At sampling, the speed and computational requirements are identical to vanilla CMs for the one-step procedure, while for multistep sampling we need to account for additional forward passes of the encoder as shown Algorithm 2.
min Lise (0,4) 0.6
() + B (1oI [o))l 0, I))]
6.1. Connection with variational autoencoders
In this section, we will consider the special case with con- stant unit time weighting Aes Yt.
By carefully selecting the value of B, we can achieve a proper balance between flexibility and proximity to the prior, To understand the effect of [3 in our model, we present an alternative form of our objective function. Formally, can view the minimization of Eq. 9 as the relaxed Lagrangian problem of the following optimization problem:
ELBO perspective. First, we demonstrate the relationship between our model and VAE in terms of their objective functions. Specifically, our loss function in Eq. 7 serves as an upper bound on a standard VAE loss, where the latent vector 21 is regularized to be close to the prior P1. Using the triangle inequality, we can establish the following bound for goudd
min Cood (0, ())
0.0
s.t.
DKL xD) (N)]\ 0,T)) < 5,
fo(m1,1)11&
(8)
As 5 approaches 0, the coupling becomes イ( xp) = po (o(z))pz (a1), indicating that our model encompasses the standard CT. In VAE, selecting appropriate values of
N
|| | (200 21), fo- (4h,, 11), x1)11)
5

---

VCT: Training Consistency Models with Variational Noise Coupling
Algorithm 2 Multistep Variational Consistency Sampling
training, as it increases whenever the discretization scheme is changed, which reflects in Aty becoming smaller. For ECM models trained on ImageNet, which use the EDM- style weighting function
Input: Consistency model fo. encoder Jø, sequence of TN-1, initial noise Tr
time points 71 > 72
I < fo (0rr T)
for n = 1 to N- 1 do
Sample E ~ I) 1 f g0 (z) + gE(a)e
+ (ee = 8
data
T.. a,, xt + br. T1
we simply select a fixed 3 scalar for the whole training, as the discretization scheme does not affect the magnitude of the weights.
x So ( T,, Tn)
end for Output: x
7. Experiments
to achieve reasonable generation performance is generally challenging. Values too close to zero result in strong de- viation from the prior, while extremely large values cause over-smoothed decoders (Takida et al., 2022), leading to blurry samples. However, our model does not suffer from this over-regularization issue. While tuning 3 remains cru- cial in our model, as demonstrated in Section 7.2, unlike VAE, increasing values of 3 does not cause the oversmooth- ing problem but simply reduces our model to the standard CT. Consequently, the CT training objective enables sharp sample generation even when the posterior approximation is nearly a normal distribution.
In the following, we show that learning the data-noise cou- pling with our method is a simple yet effective improvement for CT. We adapt our Variational Coupling (VC) technique to two established baselines: improved Consistency Train- ing (iCT) (Song & Dhariwal, 2024) and Easy Consistency Tuning (ECM) (Geng et al., 2025). Note that for the latter, the model is initialized with the weights of pretrained score model from (Karras et al., 2022), while in the former the weights are initialized at random. More details about the baselines are provided in Appendix C.I. In this work, we consider only the framework of CT, where the unbiased vector field estimator ue (a:e) from equation is used to ap- proximate the noisy states d; during training, as opposed to Consistency Distillation that uses pretrained score model as a teacher. We evaluate the models on the image datasets MNNIST (Xiao al., 2017), CIFAR- 10 (Krizhevsky et al., 2009), FFHQ 64 64 (Karras et al., 2019) and (class- conditional) ImageNet 64 x 64 (Deng et al., 2009). To learn the coupling, we add a smaller version of the neural network used for CT, without time conditioning and with weights always initialized at random. For all the models, we use the variance exploding transition kernel (iCT-VE and ECM-VE) used in (Karras et al., 2022) and (Song & Dhariwal, 2024), with a 1, b, = f, and the linear inter- polation kernel (iCT-LI and ECM-LI) commonly used in Flow Matching (Lipman et al., 2023), with = 1-t/omman and by = t t/omax (details in Appendix A). For both kernels, we set amin = 0.002 and 0max 80. We report more experimental details in Appendix C. comparison with other models in Appendix D, while samples obtained with our best models are shown in Appendix G.
6.2. Choosing the 3 parameter
The most common weighting function (t) for CMs is an adaptive weighting scheme that changes over training based on how fine grained the discretization is, or equivalently the distance between consecutive time steps in ECM. At given training iteration, for two adjacent time steps t, and f we have:
Aet (l++1)
A = -411-
For the models trained with such an adaptive weighting function, we found it hard to tune 3 to a single scalar value. The magnitude of the weights increases during training as the discretization scheme becomes more fine-grained and A... becomes smaller, changing the balance between consistency loss and KL regularization, resulting in a very strong regularization at the early stages of training, or too weak one at the later stages. A simple yet effective solution is to use adaptive scaling for the KL regularization that the discretization scheme. To do so, we take as a reference the weighting of the consistency loss at the last step IN = Tmax, and define the adaptive KL weighting as:
changes according to
7.1. Baselines and models
As baselines, we re-implement the iCT and ECM models, corresponding to our iCT-VE and ECM-VE. As addi- tional model, we add CT with the minibatch Optimal Trans- port Coupling (-OT) proposed in (Pooladian et al., 2023; Tong et al., 2024) and used in (Issenhuth et al., 2025), to compare the effectiveness and scalability of the OT coupling with the learned one. Finally, we combine the baselines with
B AIx
AKL = (t)
This way, we only need to specify the scalar hyperparameter 9, and will have consistent regularization strength over
6

---

VCT: Training Consistency Models with Variational Noise Coupling
our proposed Variational Coupling (-VC). For the models with learned coupling, use gradient clipping with large value (200 in all the experiments) to avoid instabilities at the early stages of training.
Model
CIFARIO
FPHQ (64 61)
ImageNet (61 x 64)
ECM-VE"
3.60/.11
5,51 /3.18
ECM-VE
3,68/2.14
5499
5,26/3.22
ECM-LI
3.65/2,14
6.42 /4.7
5.13/3.20
ECM-VE-OT
3,46/2.13
6.11 /4.68
6.02/4.22
ECM-LI-OT
3.49/2.13
6.19
5.63 56.3/4.09
ECM-VE-VC (VCT)
3.26/2.02
5,47 /416
5.08/3.15
7.2. Ablation for different |3
ECM-LI-VC (VCT)
3.39/2.09
5.57 /4.29
4.93/3.07
To see the effect of 3 the generation performance, we compare the results for different values on the FashionM- NIST dataset for iCT-VE-VC in Table 1. As expected, with small values of B. the coupling distribution deviates from the sampling distribution and the performance degenerates, while increasing / high values reduces the benefits of the learned coupling. More detailes on how we tune /3 are reported in Appendix C.3.
Table 3. Comparison of FID (lower is better, reported as 1-step / 2- step performance) for different models hased F.M. The model marked with a = is the baseline as reported in (Geng et al., 2025). All the other models are from our re- i-mplementation. The best entries are highlighted in bold. For ImageNet, the results marked with f are obtained with models trained for 100k iterations, while the others use 200k iterations. Comprehensive comparisons with additional baselines are presented in Tables and 10.
FID for different
only use models based on iCT as there is no available pre- trained EDM model.
step
2 steps
8 8= 5
12.53
5.69
first benchmark to test the performance of iCT. On this dataset, we use small version of DDPM++ with 64 model chan- nels instead of 128 and no attention, and batch size 128, Our variant with Variational Coupling outperforms both iCT and iCT-OT, with best performance obtained with the LI transi- tion kernel, showing the benefit of the learned coupling. CIFAR-10: For all the CIFAR-IO experiments, we use the DDPM++ architecture from (Song et al., 2021b) as imple- mented in (Karras et al., 2022), with EMA rate 0.9999 as in (Geng et al., 2025). While this differs from the settings in (Song & Dhariwal, 2024), we found it to work better in our re-implementation. The remaining hyperparameters are the same as used in the respective baselines. From the results, we can see how using the learned coupling results in im- proved performance for both one and two steps generation, outperforming all the re-implemented baselines. The I-step result from the original iCT is superior to our model. How- ever, to the best of our knowledge, there is no open-source implementation that can reproduce the results reported in the paper. The learned coupling outperforms the minibatch OT coupling all cases, as is less affected by the effective (per device) batch size and the data dimensionality. Finally, our 2-step sampling performance for ECM-VE-VC is on par with the current SoTA achieved by other methods with similar settings (Wang et al., 2024; Lee et al., 2025: Lu & Song, 2025). We empirically compare the variance of the gradients for iCT-VE and iCT-VE-VC, and show in Figure how the resulting reduced variance corresponds to improved FID score. In particular, in early training the model with VC exhibits higher variance, which can be attributed to the fact that the encoder is still learning the coupling. As training continues, the coupling becomes effective at providing bet- ter data-noise pairs to the model, which results in reduced gradient variance and improved generative performance. FFHQ 64 x 64: We use FFHQ 64 64 as an additional
FashionMNIST: We choose FashionMNIST
/ 15
4.97
2.34
| 30
3.88
2.37
8 60
3.90
2.67
Table Comparison of FID performance (lower is better) for one and two sampling steps, for varying values of B. The models are iCT-VE-VC VC and trained on the FashionMNIST dataset with the same settings described in appendix C. Best entries in bold.
1-step / 2-step FID for iCT-based models
Model
FasSion-MNIST
CIFAR10
iCT-VE
2.83 2.46
ICT-VE 4.79/3.54
3.79
iCT-LI
4.75 /3446
3.81/2.87
iCT-VE-OT
2.82
3,28/2.66
iCT-LI-OT
4.91/
3.77
(CEVE-VC(VCT) 3.88/ / 2.37
2.86 2.33
(VCT) 3.62/
2.94 2.32
Table 2. Comparison of FID (lower is better, reported as 1-step/ / 2-step performance) for different models based on iCT. The model marked with is the baseline reported in (Song& & Dhariwal, 2024). All the other models are from our re-implementation. The hest entries are highlighted in bold.
7.3. Results
In tables 2 and 3 we report the and step sample quality evaluated with Frechet Inception Distance (FID) (Heusel et al., 2017), for both the results reported in the original papers and our re- ie-mpeementations. For high-dimensional data, we only use models based on ECM, as they require lower computational budget, while for FaShionMNIST we
7

---

VCT: Training Consistency Models with Variational Noise Coupling
ICT
Gracieat Varianes
10
101-3
(a) 1-step (FID=5.13 13, left) and 2-step (FID-3.20, right) samples from ECM-IJI
Trainnng Tremations (x 10P)
ICT ICT-VC
210
stop F1D
10
(b) 1-step (FID=4.93, left) and 2-step (FID=3.07, right) samples from ECM-LI-VC.
50
100
150
2001
150
300
350
000
Training
(x(10")
Figure 2. Visual comparison of generated class-conditional sam- ples ImageNet 64 64.
Figure 3. The top graph shows comparison of gradient variance during training for iCT-VE and iCT-VE-VC on CIFARIO. plot the variance for each epoch (shaded) and its exponential moving average with smoothing factor 0.9. Especially later during training, the model with learned coupling exhibits lower variance, which results in improved performance, shown in terms of 1-step FID in the bottom graph. For fair comparison, did not gradient clipping for iCT-VC in this run.
dataset to assess our method on higher-dimensional data. We reuse the same training settings used for CIFARIO, 10, without additional tuning, and with the same network architecture used in EDM. While the results are worse than current SOTA generative models (e.g, 2.39 FID from EDM), they confirm the benefit of using the learned coupling over the baselines. Moreover, the results highlight the limits of using the mini- batch OT coupling, which scales poorly with increased data dimensionality and in some cases performs worse than the independent coupling.
8. Conclusions
In this work, introduced novel approach Consistency Training (CT) by incorporating variational noise coupling mechanism. Our method leverages an encoder-basee cou- pling function to learn a data-dependent noise distribution, which results in improved generative performance. By fram- ing CT within the Flow Matching perspective, we provided principled way to introduce adaptive noise coupling while maintaining the efficiency of standard CT. Empirical re- sults multiple image benchmarks, demonstrate that our approach consistently outperforms baselines in one and two- step generation settings. Our findings highlight the potential of learned coupling in CT and suggest several promising directions for future work. These include exploring more ex- pressive posterior distributions, extending our method to the continuous-time CT formulation, and integrating variational consistency training with other recent CT improvements. We hope this work contributes to the broader understanding of the effect of coupling in CT and inspires further advance- ments in efficient generative sampling techniques.
ImageNet 64 x 64 (class conditional): As a baseline, we reuse the settings from ECM with the EDM2-S architecture and batch size 128. While the baseline is trained for 100k iteration, we found that models with Variational Cou- pling needed more time to converge properly, as the encoder weights are not pretrained and initialized at random. We therefore train our re-implemented baselines and models for 200k iterations instead, while we report the performance of our models trained for 100k iterations in Appendix C.5. In the 200k case, the models with Variational Coupling outper- form the other models, with the LI kernel obtaining the best overall FID, while the OT coupling performs poorly due to the small batch size and high data dimensionality. In Fig- ure 2 we compare samples from ECM-LI and ECM-LI-VC, where we can see how the images generated with VC are more clear and detailed.
8

---

VCT: Training Consistency Models with Variational Noise Coupling
Impact Statement
Dhariwal, P. and Nichol, A. Diffusion models beat gans
image synthesis. Advances in neural information processing systems, 34:8780-8794, 2021.
This paper presents work whose goal to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Dou, H., Lu, J., Du, J., Fu, C., Yao, W., qian Chen, X., and Deng, Y. A unified framework for consistency generative modeling. 2024.
Acknowledgment
Geng, Z., Pokle, A., Luo, W., Lin, J., and Kolter, J. Z. Con- sistency models made easy. In The Thirteenth Interna- tional Conference on Learning Representations (ICLR), 2025.
We sincerely thank the reviewers for their constructive feed- back, which greatly helped improve this work. We are also deeply grateful our colleague, Naoki Murata, for revising the manuscript and providing valuable comments.
Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial nets. Advances in neural informa- tion processing systems, 27, 2014.
References
Albergo, M. and Vanden-Bijnden, E. Building normaliz- ing flows with stochastic interpolants. In ICLR 2023 Conference, 2023.
Heek, J., Hoogeboom, E., and Salimans, T. Multistep con- sistency models. arXiv preprint arXiv:2403.0668 2024.
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Albergo, M. S., Goldstein, M., Boffi, N. M., Ranganath, R., and -inden- Stochastic interpolants with data-dependent couplings. In Forty-first International Conference on Machine Learning, 2024.
Higgins, I., Matthey, L., Pal, A., Burgess, C., Glorot, X., Botvinick, M., Mohamed, S., and Lerchner, A. beta- vae: Learning basic visual concepts with a constrained variational framework. In International Conference on Learning Representations, 2022.
Aneja, J., Schwing, A., Kautz, J., and Vahdat, A. A con- trastive learning approach for training variational autoen- coder priors. In Advances in neural information process- ing systems, volume 34, pp. 480 493,4933220 2021.
Bartosh, G., Vetrov, D., and Naesseth, C. A. Neural diffu- sion models. In International Conference on Machine Learning, pp. 5073-3095, PMLR, 2024a.
Ho, J., Jain, A., and Abbeel, Denoising diffusion proba- bilistic models. Advances in neural information process- ing systems, 33:6840-6855 2020.
Bartosh, G., Vetrov, D. P., and Andersson Naesseth, C. Neu- ral flow diffusion models: Learnable forward process for improved diffusion modelling. Advances in Neural Infor- mation Processing Systems, 37:73952-73985, 2024b.
Ho, J., Salimans, T., Gritsenko, A., Chan, W., Norouzi, M., and Fleet, D. Video diffusion models. Advances in Neural Information Processing Systems, 55:8633-8646 2022.
Hoffman, M. D. and Johnson, M. J. Elbo surgery: yet another way carve up the variational evidence lower bound. In Workshop in Advances in Approximate Bayesian Inference, NIPS, volume 1, 2016.
Berthelot, D., Autef, A., Lin, J., Yap. D. A., Zhai, S., Hu, S., Zheng, D., Talbott, W., and Gu, F. Tract: Denoising diffiusion models with transitive closure time-disilllaton. arXiv preprint arXiv:2303.0 04248, 2023.
Issenhuth, T., Santos, L. D., Franceschi, J.-Y., and Rako- tomamonjy, A. Improving consistency models with generator-augmenter augmented flows. In Forty-Second Interna- tional Conference on Machine Learning, 2025.
Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Cornish, R., Caterini, A., Deligiannidis, G., and Doucet, A. Relaxing bijectivity constraints with continuously indexed normalising flows. In International Conference Machine Learning, pp. 2133-214442 2020.
Milcoouur-Martnnaau A., Li, K., Taillef,,,,,,er, R., Kach- man, T., and Mitliagkas, I. Gotta go fast when gen- erating data with score-based models. arXiv preprint 1rXiv:2105.14088 2021.
Deng, J., Dong, W., Socher, R., Li, L.-J,, I, Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. leee, 2009.
Karras, T., Laine, S., and Aila, T. A style-based generator architecture for generative adversarial networks. In Pro- ceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 4401 4410, 2019.

---

VCT: Training Consistency Models with Variational Noise Coupling
Karras, T., Aittala, M., Aila, T., and Laine, S. Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems, 35: 2022.
Liu, L.,, Ren, Y., Lin, Z., and Zhao, Z. Pseudo numerical methods for diffusion models on manifolds. In Interna- tional Conference Learning Representations, 2022.
Liu, X., Gong, C., and Liu, Q. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations (ICLR), 2023.
Karras, T., Aittala, M., Lehtinen, J., Hellsten, J., Aila, T., and Laine, S. Analyzing and improving the training dynamics of diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 24184,24184, 2024.
Lu, C. and Song, Y. Simplifying, stabilizing and scaling continuous-time consistency models. In The Thirteenth International Conference on Learning Representations (ICLR), 2025.
Kim, B., Hsich, Y.-G., Klein, M., Cuturi, M., Ye, J. C., Kawar, B., and Thornton, J. Simple reflow: Improved techniques for fast flow models. In The Thirteenth Inter- national Conference Learning Representations (ICLR), 2025.
Lu, C., Zhou, Y., Bao, F. Chen, J., Li, C., and Zhu, J. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances Neural Information Processing Systems, 35:5775-5787 2022.
Kim, D., Lai, C.-H., Liao, W.-H., Murata, N., Takida, Y., Ue- saka, T., He, Y., Mitsufuji, Y., and Ermon, S. Consistency trajectory models: Learning probability flow ode trajec- tory of diffusion. In The Twelfth International Conference on Learning Representations, 2024.
Luhman, E. and Luhman, T. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.02388 2021.
Luo, W., Hu, T., Zhang, S., Sun, J., Li, Z., and Zhang, Z. Diff-instruct A universal approach for transferring knowledge from pre-trained diffusion models. Advances in Neural Information Processing Systems, 36:76525- 76546, 2023.
Kingma, D. P. and Welling, M. Auto-encoding variational bayes. International Conference on Learning Represen- tations, 2013.
Kobyzev, I., Prince, J., and Brubaker, M. Normalizing flows: An introduction and review of current methods. IEFE transactions on pattern analysis and machine intel- ligence, 43(11):396443979 2020.
Nielsen, B. M., Christensen, A., Dittadi, A.. and Winther, O. Diffene: Variational diffusion with learned encoder. In The Twelfth International Conference on Learning Representations, 2024.
Kong, Z., Ping, W., Huang, J.. Zhao, K., and Catanzaro, B. Diffwave: A versatile diffusion model for audio synthesis. In International Conference on Learning Represenations, 2021.
Papamakarios, G., Nalisnick, E., Rezende, D. J., Mohamed, S., and Lakshminarayanam, B. Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22(57):1-64, 2021.
Krizhevsky, A., Hinton, G., et al. Learning multiple layers of features from tiny images. 2009.
Pooladian, A.-A., Ben-Hamu, H., Domingo-Enrich, C., Amos, B., Lipman, Y., and Chen, R. T. Multisample flow matching: straightening flows with minibatch couplings. Proceedings of the 40th International Conference on Machine Learning, pp. 18100-28127, 2023.
Lee, J., Park, J., Yoon, J., and Lee, J. Stabilizing the training of consistency models with score guidance. In ICML 2024 Workshop on Structured Probabilistic Inference ((&) Generative Modeling. 2024.
Lee, S., Kim, B., and Ye, J. C. Minimizing trajectory curva- ture of ode-based generative models. In International Conference on Machine Learning, pp. 18957-18973. PMLR. 2023.
Rezende, D. J., Mohamed, S., and Wierstra, D. Stochastic backpropagation and approximate inference in deep gen- erative models. In International conference on machine PMLR, 2014.
learning, pp.
Lee, S., Xu, Y., Geffner, T., Fanti, G., Kreis, K., Vahdat, A., and Nie, W. Truncated consistency models. In The Thirteenth International Conference on Learning Repre- sentations (ICLR), 2025.
Rombach, R., Blattmann, A., Lorenz, D.,, Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF con- ference on computer vision and pattern recognition, pp. 10684-10695, 2022.
Lipman, Y., Chen, R. T., Ben-Hame H., Nickel, M., and Le, M. Flow matching for generative modeling. In The Eleventh International Conference on Learning Repre- sentations, 2023.
Rosca, M., Lakshminarayanan, B., and Mohamed, S. Distri- bution matching in variational inference. arXiv preprint arXiv: aeee:80020008 2018
10

---

VCT: Training Consistency Models with Variational Noise Coupling
Salimans, Т. and Ho, J. Progressive distillation for fast pling of diffusion models. In International Conference Learning Representations, 2022.
Xie, S., Xiao, Z., Kingma, D., Hou, T., Wu, Y. N., Murphy, K. P., Salimans, T., Poole, B., and Gao, R. Em distilla- tion for one-step diffusion models. Advances in Neural Information Processing Systems, 37:45073 -500734-5144, 2024.
Salimans, T., Mensink, T., Heek, J., and Hoogeboom, E. Multistep distillation of diffusion models via moment matching. Advances in Neural Information Processing Systems, 37:36046-360700 2024.
Xu, Y., Liu, Z., Tian, Y., Tong, S., Tegmark, M., and Jaakkola, T. Pfgm++: Unlocking the potential of physics- inspired generative models. In International Conference on Machine Learning, pp. 38566-38591. PMLR, 2023.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. International conference on machine learning, pp. 2256-2265, PMLR, 2015.
Yang, L., Zhang, Z., Zhang, Z., Liu, X., Xu, M., Zhang, W., Meng, C., Ermon, S., and Cui, B. Consistency flow matching: Defining straight flows with velocity consis- tency. arXiv preprint
2024.
Song. J., Meng, C., and Ermon, S. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021a.
Yin, T. Gharbi, M., Zhang, R., Shechtman, E., Durand, F., Freeman, W. T,, and Park, T. One-step diffusion with distribution matching distillation. In Proceedings the IEEF/CVF conference on computer vision and pattern recognition, pp. 6613-6623, 2024.
Song. Y. and Dhariwal, P. Improved techniques for train- ing consistency models. In The Twelfth International Conference on Learning Representations, 2024.
Zheng, H., Nie, W., Vahdat, A., Azizzadenesheli, K., and Anandkumar, A. Fast sampling of diffusion models via operator learning. In International conference machine learning, 42390-42402. PMLR, 2023.
Song, Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021b.
Zhou, M., Zheng, H., Wang. Z., Yin, M., and Huang. H.
Song, Y., Dhariwal, P., Chen, M., and Sutskever, Consis- tency models. In International Conference on Machine Learning, pp. 32211 1 -32222.. PMLR, 2023.
Score identity distillation: Exponentially fast distillation of pretrained diffusion models for one-step generation. In Forty-first International Conference on Machine Learn- ing, 2024.
Sutton, R. S. Reinforcement learning: An introduction. A Bradford Book, 2018.
Takida, Y., Liao, W.-H., Lai, C.-H., Uesaka, T., Takahashi, S., and Mitsufuji, Y. Preventing oversmoothing in via generalized variance parameterization. Neurocomputing, 509:137-156, 2022.
Tong, A., FATRAS, K., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Wolf, G., and Bengio, Y. Improving and generalizing flow-based generative models with mini- batch optimal transport. Transactions Machine Leam- ing Research, 2024. ISSN 2835-8856. URL https: / openreviess net. / forum? d=CDDSnc pert Certification.
Ex-
Vahdat, A.,
Kreis, K., and Kautz, J. Score-based generative modeling in latent space. Advances in neural information processing systems, 34: 1211118811113 1302, 2021.
Wang, F-Y.,, Geng, Z., and Li, H. Stable consistency tuning: Understanding and improving consistency models. arXiv preprint arXiv:2400. 0ee::2440..9999 18958, 2024.
Xiao, H., Rasul, K., and Vollgraf, R. Fashion-mnise a novel image dataset for benchmarking machine learning algorithms, 2017.

---

VCT: Training Consistency Models with Variational Noise Coupling
A. Consistency Models with Linear Interpolation Kernel
addition to the variance exploding forward process commonly used in CT, here we propose to use the linear interpolation kernel commonly used in Flow Matching:
(1 1) + la
(10)
reuse all of building blocks from iCT and ECM and make only the necessary adjustments. Accounting for boundary conditions, the transition kernel becomes:
2o
21 maxx
(1I)
mmax
mmxx
Other crucial components for stability during training are the scaling factors Cins Cakip and Cout and derive them for the linear interpolation kernel following the same procedure used in (Karras et al., 2022), also accounting for the boundary conditions when min * 0:
(a) =
(12)
V Vo2 data ) o2
Jantal1 (1 Omas Omin
omin
(o) (a mmin) data
(13)
Omin
Tanx
mmnn
(o) = (a mmin ddata
(14)
A.1. Derivations
We report the derivations for the scaling factors used for the linear interpolation transition kernel. We follow the same derivations from (Karras al., 2022) (appendix B.6), where the score matching objective is written
Do + n; - y
(15)
Where y data sampled from the data distribution with standard deviation data and is sample from noise distribution with standard deviation Given this objective, they propose to derive the scaling factors Gia (a), (O), Cout (o) as follows:
1
(an(a) VVary, + n
(16)
Cout (g) Varg Cskin (a (o) n)]
(17)
Cakip (a) argmin cop Cour (a)².
(18)
In our formulation, we only need to rescale by amaz and perform the same derivations. For simplicity, we define a a_ (omitting the dependence on σ), and proceed as follows:
Gin (a) Var az + n2
(19)
Cout (a) = Varg,n (y Cskip (o) n))
(20)
Catip (7) argmin (oun
(21)
The factor (7) simply becomes:
(n) = Volata 2 o
(22)
12

---

VCT: Training Consistency Models with Variational Noise Coupling
To derive (o) we can proceed as:
Cout cou = Varg,n(y - (a) + n)) Cout (o)² = Vary.n (1 acakip (0) yy + Cakip (r)m]
(23)
(24)
Cout (o)²= (1 aCskip (o)) odaaa Cakip(o) 00²
(25)
We can use this result to solve for Cakip (o)):
0 = ( (c)s)/css)pp(()
(26)
0 d[(1 (o)a (o))) data Cskip (0)) (a)
(27)
0 d(I acskip (o) //c))// (0) + o² ( /d)ck/ (o)
(28)
= dantal l cakp((o) 2a] + o (2ckkp(a))
(29)
0 = (o² + ² dd Cakip (0) ad
(30)
Cakip (a) addutal (0" + a od
(31)
Finally, we can compute Cout Cout(o))
Cout (o)² = (1 - 11--acskipll (Cskip(o)) (o)) dataa + cakiplo) 0²
(32)
a=002 odtan
2
ao 2
2
Cout (0)² =
(02
a²02
data
data
+
(02 +
a² 2
data
data
0²
(33)
Pdata
2
ao data to
2
Cout (o)²² =
+
(34)
(a² + azo data
(o2 a² 220
Cout 2a) (o² (o² )2 + (0o0 data
(35)
(02 + 2 data
Cout(or² (od )z + (a dasta + +
(36)
(02 ( +220 a²o? data )2
(oodnta)2
Cout (o)² (02 + a
(37)
Oodata
Cout (o) V02 + a data
(38)
If we want to use the boundary conditions for miin * O, then we can modify Cakip (a) and Cout (r) as:
(o) = 2 (1 a mmin
Omax Omin
(39)
(a mmin)² ) Jdatal (1
Emin Teax Omin
(o) = (a Tmin ) ddatt
(40)
which satisfy the condition Cakip(omin) = 1 and Cout (Fmin) = 0.
B. Consistency Lower Bound
Proposition B.1. The following upper bound for negative log-density holds:
- log po(oo)
<
20
(2:
(20)
|00 - 1012 + DKL gol2 mx) p(z)
<
1
E
So the
11
dl DrCs (asl= wo) plz)).
20
JO
dl
13

---

VCT: Training Consistency Models with Variational Noise Coupling
This proposition establishes the connection between the minimization objective of CT-VC and that of the continuous-time
CM.
Proof. For the triangle inequality, we have:
au fo(01, 1)ll
(41)
x1 fo (1); t2+1) fo- (1,, (2o: 21). te)]l
We square both sides, obtaining:
100 fo( 1)11 <
(42)
Σ fe (xo; x1), 1441) fo-(ooo (2br x1 ()))
Now, for the Cauchy-Schwarz Inequality, we can write the right hand side as:
fo ( (2o) x1 (20) t(+1) fo (200 20 (1); -
(43)
N
(2o: t++1) - fo- (zb1, (000: t4),
1=0
Since:
log po x1) x
20 100 f0 1)]1?
(44)
write lower bound the log density
log (xro 2
Ey
fo(a,, 1)122
DKL. (go(0s (ro) p(a1))
(45)
20
In summary, the ELBO bound for the Gaussian VAE is given by
20 Bas (wx) |a: So (æ), 1)11 DRL. (z( | xx)) || /(ali)) L(xo: 0, Φ).
1
log pappoo(2) (2o) <
(46)
Combining the inequalities we derived in Eq. 42 and 43, below we consider the case when N + co. First, we define:
SN = N Σ ((( [ll (1/ (20; 1--)) fo( (111
(2o:
(47)
We assume that t; is a partition of the interval [0, 1] of:
A = 4:41 0
(48)
For small At O (+), we approximate the squared difference in function values using first-order Taylor expansion
d
( (441) fo(W,, t.)11 ~
- dt
Fo
t)
At2
(49)
1=
Herre, we assume Vee4 - Vt. O(t++1 -t)
14

---

VCT: Training Consistency Models with Variational Noise Coupling
Since Al O (
substitute:
Il (11))) Sa 1il1
1 N2
di S0
d
(50)
Multiplying by N:
12
SN N
E (M4p))))))))))) folve, +.11
Σ E
fol
(51)
As N + 00,
d so(ves))
1 +
d
0
So
(52)
N
di
1=0
Thus, multiplying by N, we obtain:
di solvss 0
lim
SN =
(53)
+00
Combining the above limit with the ELBO bound:
Po (x)) x
1
Eg (un
xo
fo(01, 1)11 + DeL (go(= xo) p(z))
(54)
20
1
1
So (the 11/
+ DRI. (oo(= mo) p(z))
(55)
2x
N
t=
d
fo
wee t)
+ DEL go(z xo) p(z)) as N - oc.
(56)
20
Jo
dl
C. Experimental details
C.1. Baselines
Here we recap the detials about the two baselines used in this work, the Improved Consistency Training from (Song Dhariwal, 2024) and Easy Consistency Models from (Geng et al., 2025).
iCT: the training procedure equation from (Karras et al., 2022):
discretization of time steps between two values m 0.002 and dmax 80, with the
(Qui i- T N(k) -
(o - max
) where i E [[1,
min)
(57)
where p= 7 and N(k) is scheduler that defines the number of discretization steps at the k-th training iteration. N(k) chosen be an exponential schedule which starts from so 10 steps and reaches $1 1280 steps the end of the training. and is defined as:
K
N(k) m2k/2 .1) +1, K'
(58)
log, ($1/ so
During training, time steps & (or equivalently a;) are sampled following a discrete lognormal distribution:
p(o;) x erf (log(1441)
Pinean
log(os)
P P
erf
(59)
V2Pud
/2P
15

---

VCT: Training Consistency Models with Variational Noise Coupling
with Pa Pimaa-- -1.1 and Pad- 2.0. Then, the steps 1 and 4411 are used in the loss:
(e, Φ) t det (et(t))d((F ( (2 t + 1), fo- (II,, t4)),
(60)
whith the time dependent weighting function Aet (4:) = 1141-12" and d(...) is the Pseudo-Huber loss:
d(ær, v) = v 1|0 -l11 + c² -
(61)
ECM: ECM aims simplify and improve the training procedure from iCT. We report here the main differences. Instead of using a discretized grid of time steps, it samples time steps & from a continuous lognormal distribution with P'mean -1.1 and Pau 2.0 (-0.8 and 1.6 for ImageNet). The second time step r used in the discretized training objective is then obtained with mapping function
p(r iters) = - (t) = [itersdd] n(():
(62)
9
where n(t) = + kobb-)) = 1- If o.) is the sigmoid function, iter is the current training iteration, A: = 8, b 1, and q 2 for all the models but ImageNet, where q 4. The discretization step is made smaller for eight times over training (four times for ImageNet). The loss function is generalization of the Pseudo-Huber loss, which consists of the L2 loss and an adaptive weighting function wA) The models are initialized with the weights of pretrained diffusion models, which is shown to greatly improve stability during training and generation performance.
C.2. Training details
We report the training details for our models in Tables 4 and Note that the baselines are the ones from our reimplementation. The models have the same number of parameters and training hyperparameters regardless of the transition kernel used. In the following, we report additional information important for reproducing out experiments:
ECM-LI: In BCM, the time steps ar sampled from lognormal distribution, done in (Karras et al., 2022). This means that time steps > maax can be sampled during training. While this works well when using the variance exploding Kernel, in the linear interpolantion case the time step & cannot exceed maxx, and we therefore clip / to be at most maxx
Random seeds: All the training runs are initialized with random seed 42, For sampling and FID computation, we always set the random seed to 32, which was randomly chosen. This differs from what commonly done in EDM, where three different seeds are used to evaluate FID and the best result reported. While our evaluation can lead to slightly worse results, the evaluation consistent between our models and reimplemented baselines.
2-steps generation: Like in the original iCT baseline, all the models use t 0.821 for CIFAR1O and all the other datasets but ImageNet, where i = 1.526 is used insetad.
Data augmentation: We scale all the images to have values between -1 and 1. For CIFARIO we apply random horizontal flip with 50% probability.
Differences for ImageNet: The training procedure for ECM on ImageNet differs slightly from the one for the other datasets. The Adam optimizer is used instead of RAdam, with betas (0.9, 0.99), and inverse square root learning rate decay defined as function of the current training iteration i:
a(1)
Ore
(63)
V max(i/ 1)"
with Oirer = 0.001 (the initial learning rate) and iyef = 2000 iterations. The Exponential Moving Average uses the power function averaging profile introduced in (Karras et al., 2024). In ECM, three different EMA profiles are tracked during training, with rates 0.01, 0.05, and 0. In our reimplementation, we only use the rate 0.1. The number of times in which the discretization interval changes is reduced from 8 to 4, and the loss constant is set to 0.06.
C.3. Tuning
In our experiments, we tuned 99 with a coarse grid search with different values with a gap of 10. For iCT on CIFARIO, we initially tested the values B = [10, 20, 30, 40), of which B = 30 gave the best performance, then tested also for [25, 35]
16

---

VCT: Training Consistency Models with Variational Noise Coupling
Model Setups
MNSST
CIFARIO
Model Architecture
+
+
Model Channels
64
128
Nº of ResBlocks
4
4
Attention Resolution
16
Channel multiplyer
2, 2, 2] 13.6M
[2, 2, 2] 55.7M
Model capacity
Training Details
Minibatch size
128
1024
Batch per device
128
512
Iterations
400k
400k
Dropout probability
30%
30%
Optimizer
RAdam
RAdam
Learning rate
0.0001
0.0001
EMA rate
0.9999
0.9999
Training Cost
Number of GPUs
2
GPU types
H100
H100
Training time (hours)
28
92
Training time with OT (hours)
29
95
Encoder Details
Model Architecture
DDPM++
DDPM++
Model Channels
32
32
No of ResBlocks
Attention Resolution
Channel multiplyer
22 2, 2]
[2,
3 regularizer
30
30
Encoder Params
1.5M
1.6M
Training time with Encoder (hours)
34
102
Table 4. Model Configurations and Training Details for iCT FashionMNIST and CIFARIO
which did not improve the performance. Similarly, for ECM on CIFARIO, we tested for 3 = [10, 20, 30, 40J. and after achieving the best performance with 39 = 10, we tested 3 = [5, 15) which did not improve the performance. The tuning was done with the VE kernel and the best values were used also for the LI kernel. We used the best values of 3 also in FashinMNIST and FFHQ without additional tuning. For ImageNet we observed on early runs that / needed to be much 60, larger, so initially tuned for 3 [30, 60, 90, 120). After achieving the best results with the VE kernel for 90,
further tuned for [70, S0, 90, 100, 110] for both VE and LI kernels, and found 3 100 to be the best for VE and B = 90 for LI.
C.4. The effeet of gradient clipping
As method uses gradient clipping for stability, test the effect of gradient clipping on the baselines for a fair comparison. We report the results in tables 6 and for iCT and ECM respectively. While the results generally improve slightly in iCT based methods, it is cleare that the main performance gain in our method is given by the improved coupling.
C.5. Imagenet training with 100k iterations
We report here the ImageNet experiments with 100k iterations. For ECM-LI-VC we increased 33 to 3 100 as 3 = 90 diverged during training. Note that for runs with 100k iterations, we sometimes encountered divergences for our models with small 3 and sometimes also for the baselines, while it seems to be solved when training for 200k iterations. The 1 1-step/2-step FID results are reported in Table 8. For the settings with 100k iterations, our method performs similarly or slightly worse than the baseline. We believe this due to the fact that the encoder for our model requires more iterations to
17

---

VCT: Training Consistency Models with Variational Noise Coupling
Model Setups
CIFARIO
FFHQ 64 x 46
ImageNet 64 x 64
Model Architecture
+
+
EDM2-S
Model Channels
64
128
192
No of ResBlocks
4
4
Attention Resolution
[16]
[16]
[16, 8]
Channel multiplyer
[2, 2,
2,
[1. 3,
Model capacity
55.7M
61.8M
280M
Training Details
Minibatch size
128
128
128
Batch per device
128
128
128
Iterations
400k
400k
200k
Dropout probability
20%
20%
40% (res S 16)
Optimizer
RAdam
RAdam
Adam
Learning rate
0.0001
0.0001
0.001
EMA rale
0.9999
0.9999
0.1
Training Cost
Number of GPUs
1
GPU types
H100
H100
H100
Training time (hours)
37
95
51
Training time with OT (hours)
38
96
52
Encoder Details
Model Architecture
DDPM+
DDPM+
EDM2-S
Model Channels
32
32
32
No of ResBlocks
1
Attention Resolution
[16]
[16]
[16, 8]
Channel multiplyer
[2, 2, 2]
2,,2
[1. 2, 3,4|
] regularizer
10
100 (VE), 90 (LI)
Encoder Params
1.6M 49
1.6M
6M
Training time with Encoder (hours)
110
58
Table 5. Model Configurations Training Details ECM on CIFARIO, FFHQ x 46 ImageNet 64
learn the coupling, as demonstrated by the improved results for 200k iterations.
18

---

VCT: Training Consistency Models with Variational Noise Coupling
Method
1-step
2-step
Method
1-step
2-step
VCT-VE (w/o ge)
3.61
2.79
ECM-VE (wlo ge)
3.68
2.14
iCT-VE (w/ ge)
3.49
2.57
iCT-VE (w/ gc)
3.70
2.12
iCT-LI (w/o ge)
3.81
2.87
iCT-LI (w/o ge)
3.65
2.14
iCT-LI (w/ gc)
3.54
2.66
iCT-LI (w/ gc)
3.76
2.17
iCT-VE-OT (w/o ge)
3.28
2.66
iCT-VE- OT (w/o ge)
3,46
2.13
iCT-VE-OT (w/ ge)
3.21
2.56
iCT-VE-OT (w/ gc)
3.45
2.12
iCT-LI-OT (w/o gc)
3,42
2.77
iCT-LI-OT (w/o gc)
3.49
2.13
iCT-LI-OT (w/ ge)
3.18
2.63
iCT-LI-OT (w/ gc)
3.50
2.12
iCT-VE-VC (w/ ge) (VCT)
2.86
2.32
iCT-VE-VC (w/ ge) (VCT)
3.26
2.02
iCT-VE-LI (w/ ge) (VCT)
2.94
2.32
iCT-VE-LI (w/ ge) (VCT)
2.39
2.09
Table 6. FID performance with/without ge CIFARIO for iCT
Table 7. FID performance with/without on CIFAR10 for ECM
based methods.
based methods.
Method
1-Step FID
2-Step FID
ECM-VE
5.66
3.78
ECM-LI
5.63
3.48
ECM-VE-OT
6.74
4.71
ECM-LI-OT
6.65
4.64
ECM-VE-VC (VCT)
5.67
3.67
ECM-LI-VC (VCT)
6.34
3.77
Table Results our models trained on class conditional LageNet 64
100k iterations with batch size 128.
19

---

VCT: Training Consistency Models with Variational Noise Coupling
D. Comprehensive comparison
In this section we report a comprehensive comparison with other methods on CIFARIO and ImageNet 64 x 64, in tables 9 and 10.
Table 9. FID, NFE and # param. CIFAR-10. Bold indicates the best result for each category and NFE.
Method
NFE
FID
# param. (M)
Diffusson models
EDM (Karras al., 2022)
35
1.97
55.7
PFGM++ (Xu et al., 2023)
35
1.91
55.7
DDPM (Ho et al., 2020)
1000
3.17
35.7
LSGM (Vahdat et al., 2021)
147
2.10
475
Consistency models
1-step
iCT (Song & Dhariwal, 2024)
2.83
56.4
iCT-deep (Song & Dhariwal, 2024)
2.51
112
CTM (Kim et al., 2024) (w/o GAN)
5.19
55.7
ECM (Geng et al., 2025)
3.60
55.7
TCM (Lee al., 2025)
2.46
55.7
sCT (Lu & Song, 2025)
2.85
iCT-VE-VC (VCT)
2.86
57.3
iCT-LI-VC (VCT)
2.94
57.3
ECM-VE-VC (VCT)
3.26
57.3
ECM-LI-VC (VCT)
3.39
57.3
2-step
iCT (Song Dhariwal, 2024)
2.46
56.4
iCT-deep (Song & Dhariwal, 2024)
2.24
112
ECM (Geng et al., 2025)
2.11
55.7
TCM (Lee et al., 2025)
2
2.05
55.7
sCT (Lu & Song, 2025)
2
2.06
iCT-VE-VC (VCT)
2.32
57.3
iCT-LI-VC (VCT)
2.32
57.3
ECM-VE-VC (VCT)
2.02
57.3
ECM-LI-VC (VCT)
2.09
57.3
Variational score distillation
DMD (Yin et al., 2024)
3.77
55.7
Diff-Instruce (Luo et al., 2023)
4.53
55.7
SiD (Zhou al., 2024)
1.92
55.7
Knowledge distillation
KD (Luhman & Luhman, 2021)
9.36
35.7
DSNO (Zheng et al., 2023)
3.78
65.8
TRACT (Berthelot al., 2023)
3.78
55.7
3.32
55.7
PD (Salimans & Ho, 2022)
9.12
60.0
2
4.51
60.0
CD (LPIPS) (Song et al., 2023)
3.55
2.93
sCD (Lu & Song, 2025)
3,66
2.52
20

---

VCT: Training Consistency Models with Variational Noise Coupling
Table FID, and # param. on ImageNet 64. Dotted lines separate results by # param. Bold indicates best result for each calegory and NFE. Our models, marked with ", use batch size 128 while equivalent models based ECM use batch size 1024.
Method
NFE
FID
# param. (M)
Diffusion models
EDM2-S (Karras et al., 2024)
63
1.58
280
EDM2-XL XL (Karras et al., 2024)
63
1.33
1119
Consistency models
1-step
iCT (Song & Dhariwal, 2024)
4.02
296
iCT-deep (Song & Dhariwal, 2024)
3.25
592
ECM (Geng et al.. 2025) (EDM2-S)
4.05
280
TCM (Lee al., 2025)EDMM2S)
2.88
280
sCT (Lu & Song, 2025) (EDM2-S)
3.23
280
ECM (Geng et al., 2025) (EDM2-XL)
2.49
1119
TCM (Lee et al., 2025) (EDM2-XL)
2.20
1119
SCT (Lu & Song, 2025) (EDM2-XL)
2.04
1119
ECM-VE-VC (Ours, EDM2-S)"
1
5.08
286
ECM-LI-VC (Ours, EDM2-S)"
4.93
286
2-step
iCT (Song & Dhariwal, 2024)
2
3.20
296
iCT-deep (Song & Dhariwal, 2024)
2
2.77
592
ECM (Geng et al., 2025) (EDM2-S)
2.79
280
TCM (Lee al., 2025XEDM22S)
2
2.31
280
sCT (Lu & Song, 2025) (EDM2-S)
2
2.93
280
ECM (Geng et al., (0025)EDD2-XX))
2
1.67
1119
TCM (Lee et al., (EDM2XX))
2
1.62
1119
SCT (Lu & Song, 2025) (EDM2-XL)
2
1.48
1119
ECM-VE-VC (Ours, EDM2-S)*
2
3.15
286
ECM-LI-VC (Ours, EDM2-S)"
2
3.07
286
Variational score distillation
DMD2 w/o GAN (Yin et al., 2024)
1
2.60
296
Diff-Instruct (Luo et al., 2023)
5.57
296
EMD-16 (Xie et al., 2024)
2.20
296
Moment Matching (Salimans al., 2024)
3.00
400
2
3.86
400
SiD (Zhou al., 2024)
1
1.52
296
Knowledge distillation
DSNO (Zheng et al., 2023)
7.83
329
TRACT (Berthelot al., 2023)
1 7.43
296
2
4.97
296
PD (Salimans & Ho, 2022)
1
15.4
296
2
8.95
296
CD (LPIPS) (Song et al., 2023)
1
6.20
2
4.70
sCD (Lu & Song, 2025) (EDM2-XL)
1
2.44
1119
2
1.66
1119
21

---

VCT: Training Consistency Models with Variational Noise Coupling
E. Toy experiments
To gain a visual understanding of the benefits of variational coupling, we use the model to learn the distribution of mixture of two Gaussians, with means μ4 (0, 0.5) and #22 (0, 00.5), and standard deviation 0.05. We use iCT-LI so that the perturbed data reaches the prior even with small max. with min 0.002, max 0.1 and dasta 0.05. The models are trained for 40k iterations, with so 10 and = 80, and EMA rate 0.999. We use simple four-layers MLP with GeLU activation and Positional time embedding, with batch size 256 and learning rate le-1, For iCT-LI-VC we use 3 = 0.001. The results are shown in figures and 4 (one and two step generation respectively).
Independent compling
Nojse
Generaterl Data
0.6
0.4
0.2
0.0
0.2
0.4
0.8
-0.5
0.0
0.5
0.5
0.0
0.5
Variational coupling
Noise
Generated
0.6
Data
0.4
0.2
0.2
0.4
-066
-0.8
-0.5
0.0
0.5
U.5
0.0
0.3
Figure 4, 2-step generation result on the toy data, with t = 0.07.
22

---

VCT: Training Consistency Models with Variational Noise Coupling
F. Model diagram
figure 5 we show the difference between the forward process for standard Consistency Training and for our method with learned noise coupling, for given time step f and transition kernel characterized by the coefficients a and b,.
Independent Coupling
Variational Coupling
[~
10(,1)
(o(o))
(( [[((a
#t agdt + hyr
3; d,3ig + b,ri
flaatt
fo(zs.
Figure 5. Diagram Consistency Training with independent coupling (left) and variational coupling (right).
G. Qualitative Results
Here we report samples from our best models, iCT-LI-VC for FashionMNIST (figure 7), iCT-VE-VC and ECM-VE-VC VC on CIFAR1O (figures and 9), ECM-VE- VC on FFHQ 64 x 64 (figure 10) and ECM-LI-VC on class conditional Imagenet 64 x 64 (figure 11). In figure 6, we show the mean and standard deviation learned by the encoder for some images from the CIFAR 10 dataset. While the values are very close to standard Gaussian, the model still retains information from the original input.
23

---

VCT: Training Consistency Models with Variational Noise Coupling
Input image
Mean
Std
Figure 6. Visualization of the predicted mean and standard deviation for trained iCT-VE-VC model for different input images. For visualization purpose, we perform min-max rescaling for the predicted mean and standard deviation, as they tend to have most values close to zero and one respectively. We also the predicted 3 channels standard deviations to single channel with grayscale transform.
24

---

VCT: Training Consistency Models with Variational Noise Coupling
Figure 7. 1-step (FID=3.62, left) and 2-step (FID=2.22, right) generation from iCT-LI-VC trained FashionMNIST.
Figure 8. 1-step (FID=2.86, left) and 2-step (FID=2.32, right) generation from iCT-VE-VC trained CIFAR10.
25

---

VCT: Training Consistency Models with Variational Noise Coupling
Figure 9, 1-step (FID=3.26, left) and 2-step (FID=2.02, right) generation from ECM-VE-VC trained CIFARIO.
Figure 10. 1-step (FID=5.47, left) and 2-step (FID=4.1 16, right) generation from ECM-VE-VC trained FFHQ64 64.
26

---

VCT: Training Consistency Models with Variational Noise Coupling
Figure 11. 1-step (FID=1.93, left) and 2-step (FID=3.07, right) generation from ECM-LI-VC trained on class conditional ImageNet 64 64.
27


Reviews:
[{"review_id": "JtKWzxqFu1", "reviewer": "Reviewer_kRGJ", "summary": "This paper proposes a method that combines VAE and Consistency Model, specifically using the encoder to predict the noise corresponding to the data. The resulting data-noise coupling is used to train the consistency model. The authors claim that this approach can reduce the variance in consistency model training. Experimental results in the paper validate the effectiveness of this method in improving generation performance.\n\n## update after rebuttal\nI thank the author's response and I will maintain my score as Weak Accept. However, I still think that the comparison with CD is necessary, otherwise the importance of this paper will be reduced.", "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 reviewer for the comments and feedback. We address here some of the points and concerns raised.\n\n### **W1. Figure 3 Does Not Have Enough Support.**\n**R1.** We believe that the initial disadvantage of our method compared to the baseline is due to the fact that the encoder is still in early training, which results in an initial increase in variance. As training proceeds and the encoder’s quality improves, the variance reduces accordingly and the FID performance of our method surpasses the one of the baseline. While at first the improvement can seem only marginal, we would like to highlight that for models that already perform relatively well on a given dataset, small FID improvements can correspond to a significant image quality enhancement.\n\n### **W2. On Comparison with CD.**\n**R2.** We appreciate the reviewer’s suggestion that including additional baselines could strengthen the paper. Our focus is on training-from-scratch methods that do not assume access to a pre-trained teacher model. Thus, we primarily compare our approach with the CT counterpart. Nevertheless, we recognize the importance of a comprehensive evaluation and will include the FID results in a table for clarity in the camera-ready version.\n\n### **W3. $\\beta$-Weighting.**\n**R3.** We thank the reviewer for spotting the mistake. The correct formula for $\\lambda_{kl}$ is $\\beta \\lambda_{ct}(t_N)$ when using the adaptive loss, while simply $\\beta$ when using the weighting like in EDM. We will update the paper accordingly in the camera ready version.", "claims_and_evidence": "The paper attempts to demonstrate that the proposed method can reduce the training variance of CT, but the experimental results (Figure 3) do not provide sufficient direct support for this claim, and the improvement is not significant enough.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "This paper makes no new theoretical claims.", "experimental_designs_or_analyses": "It is recommended to add additional baselines such as consistency distillation in the experiment.", "relation_to_broader_scientific_literature": "This paper proposes a method to reduce the variance of consistency model training, which may have an impact on the field of generative models and visual generation.", "essential_references_not_discussed": "This article discusses relatively comprehensive related work.", "other_strengths_and_weaknesses": "Strengths\n+ The paper is clearly written and easy to understand. The proposed method is simple and straightforward, and looks promising.\n+ The robustness experiments with beta enhance the method's effectiveness.\n\nWeaknesses\n- In Figure 3, for gradient variance and FID metrics, iCT-VC does not show a clear advantage over iCT, with some intervals performing worse than iCT.\n- It is recommended to include a comparison with CD, as one of the disadvantages of CT compared to CD is training variance, and this method aims to alleviate that variance, making a comparison with CD meaningful.", "other_comments_or_suggestions": "It is recommended to conduct experiments on continuous consistency model training because its theoretical upper limit is higher.", "questions_for_authors": "In Line 300, it looks like beta does not affect weighting. Can the authors explain this?", "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": "0APaI3dUaJ", "reviewer": "Reviewer_chuH", "summary": "The paper aims to improve the training dynamics of Consistency Training (CT) by replacing the independent joint distribution between the source (data) and target (Gaussian noise) with a learned coupling. This coupling is parameterized as an encoder that maps each data point to a conditional noise distribution. Both the consistency model and the encoder are trained end-to-end using a mixture of consistency loss and a KL-divergence term between the outputs of the encoder and the prior. Experimental results show the benefits of this approach, with the learned coupling achieving improved sample quality, as measured by FID across multiple image datasets.\n\n\n### Update after rebuttal\nI thank the authors for their response. While I still have some reservations about the novelty of the core idea, the paper convincingly demonstrates the effectiveness of the proposed approach for consistency training through extensive experiments. The method also shows promise for integration with other consistency-based techniques. In light of this, I will raise my score.", "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 are grateful to the reviewer for the careful and thoughtful review. Below we address some of the points raised by the reviewer, especially the correctness of Eq. (8), how our results compare with the baselines, and the novelty of the method. \n### **W1. About Eq. (8).**\n**R1.** We thank the reviewer for pointing out the mistake in Eq. (8). The correct form of Eq. (8) should be:\n$$ ||x_0 - f_\\\\theta(x_1,1)||^2 \\leq \n N\\\\sum_{i=0}^N||f_\\\\theta(\\\\psi_{t_{i+1}}(x_0; x_1),t_{i+1})- f_{\\\\theta^-}(\\\\psi_{t_i}(x_0; x_1),t_{i})||^2.$$\nThis follows from the Cauchy–Schwarz inequality.  We notice that this modification does not affect our discussion. In particular, we demonstrate the connection between our training objective (an upper bound of the negative log-likelihood as in VAE) and the continuous-time Consistency Model as $N\\\\rightarrow\\\\infty$ in **R1.** to the **Reviewer zkoD**.  We will incorporate this correction and discussion in the camera-ready version. \n\n\n### **W2. On Related Baselines and Their Comparisons.**\n**R2.** We agree with the reviewer that having additional baselines can be beneficial for the paper. Upon camera-ready version, we plan to add a more comprehensive table, similarly to Table 1 from TCM [3], and including results from other relevant consistency model works such as SCT from [5], sCT from [4], and TCM.\n\n### **W3. On FID comparison.**\n**R3.** Regarding our FID results, on CIFAR-10 the aforementioned methods achieve 1-step/2-steps FID of 2.92/2.02 (SCT), 2.85/2.06 (sCT), and 2.46/2.05 (TCM), which are comparable to ours especially in the 2-steps regime. It is also important to consider that those improvements are orthogonal to ours and could in principle be combined. On Imagenet $64\\times 64$, our results are generally worse than the ones reported by the aforementioned methods, but it is important to take into account that we used minimal settings in terms of network size and training budget due to computational constraints. Regarding the iCT baseline, there is no official open source implementation available, and not being able to reproduce the exact results seems to be a common problem found in other papers too (see for example [1,2]).\n\n### **W4. On Novelty.**\n**R4.** We agree that the concept of coupling and its advantages have been explored before in other works. However, in the context of Flow Matching, it has been shown that coupling generally results in straighter trajectories which result in improved generation with less function evaluation, which does not entail that coupling would necessarily result in improved performance in CMs. While, as pointed by the reviewer, minibatch OT-coupling was already explored in CMs, our learned coupling shows improved scalability compared to minibatch OT with respect to data dimensionality and batch size, as we can see from our experiments on $64\\times64$ images.\n\n### **References**\n[1] Issenhuth, T., Santos, L. D., Franceschi, J.-Y., and Rakotomamonjy, A. Improving consistency models with generator-induced\ncoupling. arXiv preprint arXiv:2406.09570, 2024.\n\n[2] Lee, J., Park, J., Yoon, J., and Lee, J. Stabilizing the training of consistency models with score guidance. In ICML 2024\nWorkshop on Structured Probabilistic Inference & Generative Modeling, 2024a.\n\n[3] Lee, S., Xu, Y., Geffner, T., Fanti, G., Kreis, K., Vahdat, A., and Nie, W. Truncated consistency models. arXiv preprint\narXiv:2410.14895, 2024b.\n\n[4] Lu, C. and Song, Y. Simplifying, stabilizing and scaling continuous-time consistency models. arXiv preprint\narXiv:2410.11081, 2024.\n\n[5] Wang, F.-Y., Geng, Z., and Li, H. Stable consistency tuning: Understanding and improving consistency models. arXiv\npreprint arXiv:2410.18958, 2024.", "claims_and_evidence": "All claims made in the paper are properly explained and well supported, except for Eq. 8 which seems either incomplete or involves a typo. \n\nAccording to the triangle inequality, the left hand side (unsquared) is less than or equal to the summation on the right hand without squared terms. The inequality written is incorrect, and requires either removing the squared terms or including a constant multiplier $N$ on the right hand side to fix it.", "methods_and_evaluation_criteria": "The proposed method makes sense. The idea of changing the independent coupling is a fairly popular one and has been shown time and time again to lead to improvements in standard diffusion- and flow-based generative modelling. \n\nThe evaluation criteria is also reasonable with a good selection of datasets chosen. However, as this is a few-step generative modelling work, the selection of baselines compared against is quite small and including more methods will be better. Furthermore, the reported FID numbers are, at times, significantly worse compared to the baselines which raises concerns about the fairness of the evaluation, and how well the baseline was finetuned compared to the new method.", "theoretical_claims": "Only a single theoretical claim is made, that connects the proposed loss function with a VAE-style ELBO, which is sufficiently explain in Appendix B.", "experimental_designs_or_analyses": "The experiments are sound, if a bit lacking in the amount of baselines.", "relation_to_broader_scientific_literature": "As mentioned earlier, the main idea of learnt coupling between the source and target distributions is quite popular and has been applied in the context of diffusion and flow-based models before [4]. These ideas have also been used in the context of distillation, for example, minibatch OT has been shown to improve consistency models in [1]. However, the reported FID numbers in the experiments are significantly worse than more recent SOTA few-step consisteny training models such as sCT [2] and SCT [3], as well as distillation-based approaches. \n\n[1] Li, Yiheng, et al. \"Immiscible diffusion: Accelerating diffusion training with noise assignment.\" arXiv preprint arXiv:2406.12303 (2024).\n[2] Lu, Cheng, and Yang Song. \"Simplifying, stabilizing and scaling continuous-time consistency models.\" arXiv preprint arXiv:2410.11081 (2024).\n[3] Wang, Fu-Yun, Zhengyang Geng, and Hongsheng Li. \"Stable Consistency Tuning: Understanding and Improving Consistency Models.\" arXiv preprint arXiv:2410.18958 (2024).\n[4] Albergo, Michael S., et al. \"Stochastic interpolants with data-dependent couplings.\" arXiv preprint arXiv:2310.03725 (2023).", "essential_references_not_discussed": "The paper references the majority of prior works needed to understand the full context and all essential references are included.", "other_strengths_and_weaknesses": "**Strengths**:\n* The paper is very well written and nicely structured. \n* The task of improving CMs is an important one, and the idea of changing the independent coupling is an interesting one. \n* A good amount of ablations are performed that clarify the impact of each design decision. \n\n**Weaknesses**: \n\nThe paper lacks a bit of novelty as the main idea has already been shown to work in flow matching scenarios. Other similar ideas have shown to transfer over and also work well on consistency models, so it isn't incredibly surprising that this idea also falls in the same category. The improvements achieved over the baselines are quite marginal, and the final reported FID scores are mostly much worse than those of more recent consistency-based approaches.", "other_comments_or_suggestions": "No comments.", "questions_for_authors": "1) Is it possible to expand upon the baselines you compare against to also include more recent consistency works, as well as other few-step generative models?", "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": "J1iQ823mcm", "reviewer": "Reviewer_zkoD", "summary": "The authors propose an improved consistency training (CT) method by introducing a variational noise coupling scheme. The core idea involves training a data-dependent noise emission model using an encoder architecture inspired by Variational Autoencoders (VAEs). The method is theoretically linked to the VAE framework by deriving a loss function analogous to the Evidence Lower Bound (ELBO). Empirical evaluations on multiple image datasets demonstrate the superior of the proposed method.", "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 thank the reviewer for the insightful comments. \n### **W1. On Eq. (8), Typos, and Its Tightness**\n**R1.** The correct form of Eq. (8) should be:\n$$||x_0 - f_\\\\theta(x_1,1)||^2 \\leq \n N\\\\sum_{i=0}^N||f_\\\\theta(\\\\psi_{t_{i+1}}(x_0; x_1),t_{i+1})- f_{\\\\theta^-}(\\\\psi_{t_i}(x_0; x_1),t_{i})||^2.$$\nThis follows from the CS inequality. We establish its connection to the continuous-time CM as $ N\\\\rightarrow\\infty $ and analyze their optimality.  Applying a Taylor expansion under the assumption that $ \\\\Delta t := t_{i+1} - t_i = \\frac{1}{N} $ (which can be relaxed) and the above inequality, we obtain:\n$$- \\\\log p_{\\\\theta}(x_0) \\\\leq \\\\frac{1}{2\\\\sigma^2}\\\\, E_{q_{\\\\phi}(x_1| x_0)} \\Big\\Vert x_0 - f_\\\\theta(x_1, 1)\\Big\\Vert^2 +\\text{KL}(q_{\\phi}(z | x_0)||p(z)) + C$$\n$$\n\\leq \\frac{1}{2\\sigma^2} N\\\\sum_{i=0}^N  \\Big\\Vert \\frac{d}{dt} f_\\theta(\\psi_t, t) \\Big\\vert_{t=t_i} \\Big\\Vert^2 (1/N)^2 + \\text{KL}(q_{\\phi}(z | x_0) || p(z)) + C$$\n$$\\rightarrow \\frac{1}{2\\sigma^2} \\int_0^1 \\\\mathbb{E} \\left[ \\Big\\Vert \\frac{d}{dt} f_\\theta(\\psi_t, t) \\Big\\Vert^2 \\right] dt + \\text{KL}(q_{\\phi}(z | x_0) || p(z)) +C, \\quad \\text{as } N\\rightarrow\\infty, \n$$\nfor a constant $C$.\n\nTaking $ N\\\\rightarrow \\\\infty $ is reasonable even in practical scenarios, as CT and ECT propose designing the $N$ scheduler in a coarse-to-fine manner. Additionally, we observe that at the optimal values of $ \\\\theta $ and $ \\\\phi $, the reconstruction losses (i.e., the first term in the upper bounds) specifically recover the consistency function. Consequently, the bound becomes tight at the optimum. We will incorporate this discussion in the camera-ready version.\n\n### **W2. On Gradient Clipping.**\n**R2.** To evaluate the impact of gradient clipping, we applied this technique to our baseline model, iCT-VE, on CIFAR10. In the following table, we report results for both baseline and out VC method with and without gradient clipping (GC):\n|Method|1-Step FID|2-Step FID|\n|----------------------|-----------|-----------|\n|iCT-VE (w/o GC) |3.61|2.79|\n|iCT-VE (w/ GC)|3.52 |2.57|\n|(Ours) iCT-VE-VC (w/o GC)|3.20|2.45|\n|(Ours) iCT-VE-VC (w/ GC)|2.86|2.32|\n\nFrom the table it is clear that the main improvement comes from the learned coupling. Interestingly, in this case using GC for the baseline resulted in a performance improvement, even though generally GC is not mentioned in other CM literature, and we originally added it to our method to prevent early training instabilities for the learned noise distribution. However, we agree that applying GC to the baselines ensures a fairer comparison. We will include this discussion and report the results for all baselines with GC.\n\n### **W3. On ImageNet.**\n**R3.**  We also run the ImageNet experiments with 100k iterations. For ECM-LI-VC we increased $\\beta$ to $\\beta=100$ as $\\beta=90$ diverged during training. Note that for runs with 100k iterations, we sometimes encountered divergences for our models with small $\\beta$ and sometimes also for the baselines, while it seems to be solved when training for 200k iterations. The 1-step/2-step FID results are as follows:\n| Method|1-Step FID| 2-Step FID|\n|-------------|-----------|-----------|\n|ECM-VE| 5.66| 3.78|\n|ECM-LI| 5.63| 3.48|\n|(Ours) ECM-VE-VC| 5.67| 3.67|\n|(Ours) ECM-LI-VC| 6.34| 3.77|\n\nFor the settings with 100k iterations, our method performs similarly or slightly worse than the baseline. We believe this is due to the fact that the encoder for our model requires more iterations to learn the coupling, as demonstrated by the improved results for 200k iterations. We agree with the reviewer that including these results in the paper is important. We will add them, along with the corresponding results with OT coupling, to the camera-ready version.\n\n### **W4. On $\\beta$-Selection.**\n**R4.** In our experiments, $\\beta$ was tuned with a coarse grid search with different values with a gap $10$. For iCT on CIFAR10, we initially tested the values $\\beta=[10, 20, 30, 40]$, of which $\\beta=30$ gave the best performance, then tested also for $\\beta=[25, 35]$ which did not improve the performance. Similarly, for ECM, we tested for $\\beta=[10, 20, 30, 40]$, and after achieving the best performance for $\\beta=10$, we tested $\\beta=[5, 15]$ which did not improve the performance. The tuning was done with the VE kernel and the best values were used also for the LI kernel. We used the best values of $\\beta$ also in FashinMNIST and FFHQ without additional tuning. For ImageNet we observed on early runs that $\\beta$ needed to be much bigger, so we initially tuned for $\\beta=[30, 60, 90, 120]$. After achieving the best results with the VE kernel for $\\beta=90$, we further tuned for $\\beta=[70, 80, 90, 100, 110]$ for both VE and LI kernels, and found $\\beta=100$ to be the best for VE and $\\beta=90$ for LI. We will add a discussion about this process in the camera ready version, as we agree with the reviewer that guidelines on how to choose $\\beta$ are important for the practitioners.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes. The theoretical derivation in Appendix B (consistency lower bound) is basically correct but contains minor error. See Other Comments Or Suggestions for details.", "experimental_designs_or_analyses": "Yes. Two issues should be addressed:\n\n1. Gradient clipping: The proposed method uses gradient clipping (clipping value=200), while baselines do not. It is unclear whether this technique alone contributes to performance gains.\n\n2. Training iterations for ImageNet: The authors increased baseline training iterations from 100k to 200k for fairness but did not report results at 100k. This obscures the true computational trade-offs. Including 100k results and discussing training costs would strengthen the comparison.", "relation_to_broader_scientific_literature": "The work builds on consistency training [1] and leverages VAE-inspired coupling [2]. The idea of enhancing noise-data coupling aligns with [3], but differs by learning the coupling via an encoder instead of relying on the prediction of the consistency model itself during training. This connection is appropriately discussed.\n\n[1] Song, Y., Dhariwal, P., Chen, M., and Sutskever, I. Consistency models. In International Conference on Machine Learning, 2023c. URL https://api.semanticscholar.org/CorpusID:257280191.\n\n[2] Kingma, D. P. Auto-encoding variational bayes. International Conference on Learning Representations, 2013.\n\n[3] Issenhuth, T., Santos, L. D., Franceschi, J.-Y., and Rakotomamonjy, A. Improving consistency models with generator-induced coupling. arXiv preprint arXiv:2406.09570, 2024.", "essential_references_not_discussed": "No essential references appear missing.", "other_strengths_and_weaknesses": "Strengths:\n\n- The method is intuitively reasonable and grounded in established frameworks (VAEs, Flow Matching).\n\n- Extensive experiments validate the approach across datasets.\n\nWeaknesses:\n\n- Theoretical justification: While the proposed loss is derived as an upper bound to the VAE loss, the tightness of this bound is not discussed. In VAEs, the ELBO is a tight bound that achieves equality when the variational posterior matches the true posterior (i.e., optimality). However, the paper does not clarify whether the proposed upper bound can similarly reach equality or under what conditions this would occur. This raises questions about the theoretical validity of the method compared to the VAE framework, as a loose upper bound might weaken the connection to the original ELBO’s guarantees.\n\n- (Minor) Marginal gains on ImageNet: The improvement (5.13 to 4.93 in 1-step FID) is modest, raising questions about scalability to higher resolutions and more complex datasets.", "other_comments_or_suggestions": "1. Typos: In Appendix B, Equations (43, 48, 49) should use $\\geq$ instead of $leq$.", "questions_for_authors": "1. $\\beta$ selection: Table 1 shows $\\beta$ significantly impacts performance. How was $\\beta$ chosen for each experiment? Can the authors provide guidelines for selecting β in practice?", "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: Authors proposed a method that combines VAE and Consistency Model, specifically using the encoder to predict the noise corresponding to the data. Two reviewers are positive with one initially negative. After rebuttal, the negative reviewer seems to be satisfied by the theoretical justification and evaluations on ImageNet.
Decision: Accept (poster)