Inductive Moment Matching
Authors: Linqi Zhou
Stefano Ermon
Jiaming Song
Venue: ICML 2025 oral (2025)
Keywords: generative models
diffusion models
flow matching
moment matching
consistency models
Primary area: deep_learning->generative_models_and_autoencoders

Abstract: Diffusion models and Flow Matching generate high-quality samples but are slow at inference, and distilling them into few-step models often leads to instability and extensive tuning. To resolve these trade-offs, we propose Moment Matching Self-Distillation (MMSD), a new class of generative models for one- or few-step sampling with a single-stage training procedure. Unlike distillation, MMSD does not require pre-training initialization and optimization of two networks; and unlike Consistency Models, MMSD guarantees distribution-level convergence and remains stable under various hyperparameters and standard model architectures. MMSD surpasses diffusion models on ImageNet-256x256 with 2.13 FID using only 8 inference steps and achieves state-of-the-art 2-step FID of 2.05 on CIFAR-10 for a model trained from scratch.

Inductive Moment Matching
Linqi Zhou'
Stefano Ermon
Jiaming Song'
Figure 1. Generated samples on ImageNat-256x256 using 8 steps.
Abstract
1. Introduction
Generative models for continuous domains have enabled numerous applications in images (Rombach et al., 2022; Sa- haria et al., 2022; Esser et al., 2024), videos (Ho et al., 2022a; Blattmann et al., 2023; OpenAl, 2024), and au- dio (Chen et al.. 2020; Kong et al., 2020: Liu al., 2023), yet achieving high-fidelity outputs, efficient inference, and stable training remains core challenge continues to motivate research in this domain. Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2020b), one the leading techniques, require many inference steps for high-quality results, while step-reduction methods, such as diffusion distillation (Yin al., 2024; Sauer et al., 2025; Zhou al., 2024; Luo al., 2024a) and Consistency Models (Song et al., 2023; Geng al., 2024; Lu & Song, 2024; Kim et al., 2023), often risk training collapse without careful tuning and regularization (such as pre-generating data-noise pair and early stopping).
Diffusion models and Flow Matching generate high-quality samples but are slow at inference. and distilling them into few-step models often leads to instability and extensive tuning. To resolve these trade-offs, we propose Inductive Moment Matching (IMM), new class of gen- erative models for one- or few-step sampling with single-stage training procedure. Unlike distillation, IMM does not require pre-training initialization and optimization of two networks; and unlike Consistency Models, IMM guarantees distribution-level convergence and remains sta- ble under various hyperparameters and standard model architectures. IMM surpasses diffusion models 1556266 with 1.99 FID us- ing only 8 inference steps and achieves of the-art 2-step FID of 1.98 on CIFAR-10 for a model trained from scratch.
trilemma that
address the aforementioned trilemma, we introduce In- ductive Moment Matching (IMM), a stable, single-stage training procedure that learns generative models from scratch for single- or multi-step inference. IMM operates on the time-dependent marginal distributions of stochas- tic interpolants (Albergo et al., 2023) continuous-tinn time stochastic processes that connect two arbitrary probability
Luma AI "Stanford University. Correspondence Linqi Zhou <alexzhou lumalabs.ai>.
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Inductive Moment Matching
density functions (data at i = 0 and prior learning (stochastic deterministic) mapping from any marginal at time & to any marginal at time < f, it can naturally support one- or multi-step generation (Figure 2).
1). By
Data
Prior
(x
0- 0x(X,)
2: (xx)
IMM models can be trained efficiently from mathematical induction. For time s < r < L, we form two distributions at s by running one-step IMM from samples at r and We then minimize their divergence, enforeing that the distributions are independent of the starting time-steps. This construction by induction guarantees convergence to the data distribution. To help with training stability, we model IMM based on certain stochastic interpolants and optimize the objective with stable sample-base divergence estimators such as moment matching (Gretton et al., 2012). Notably, we prove that Consistency Models (CMs) are fingle-particle, first-moment matching special case of IMM, which partially explains the training instability of CMs.
Figure 2. Using an interpolation from data to prior, we define a one-step sampler that moves from any to < t. directly trans- forming qi (Xt) to ga (x.). This can be repeated by jumping to an intermediate t before moving to r.
in FM, the intermediate variable x; = ax + o,E becomes a deterministic interpolation and its interpolant velocity vi = a/x + afe reduces to FM velocity. Thus, its training and inference both reduce to that of FM. When e ~ NNO, T), stochastic interpolants reduce to u-prediction diffusion.
2.2. Maximum Mean Discrepancy
On 1m5geNet-256x256, IMM surpasses diffusion models and achieves 1.99 FID with only 8 inference steps using standard transformer architectures. On CIFAR-10. 10, IMM similarly achieves state-of -frt of 1.98 FID with 2-step generation for a model trained from scratch.
Maximum Mean Discrepancy (MMD, Gretton al. (2012)) between distribution p(x). g(y) for x E RD is an inte- gral probability metric (Müller, 1997) commonly defined on Reproducing Kernel Hilbert Space (RKHS) H with a positive definite kernel k RR D x RD
- IR as
2. Preliminaries
MMD :(y))
[x(, )) Ey (k(y, I11 (1)
2.1. Diffusion, Flow Matching, and Interpolants
where the norm in H. Choices such as the RBF kernel imply an inner product of infinte-dimensionall feature maps consisting of all moments of p(x) and g(y), i.e. E[x²) and E[y') for integer j 2 1 (Steinwart & Christmann, 2008).
For data distribution q(x). Mariance-Preserviin (VP) dif- fusion models (Ho et al., 2020; Song et al., 2020b) and Flow Matching (FM) (Lipman et al., 2022; Liu et al., 2022) construct augmented variables x; as an interpolation
3. Inductive Moment Matching
between data x q(x) and prior 6 (0, I) such that
X; 0;x OyE where ag 01 = 1, 01 do 0. VP diffusion commonly chooses 0₁ = cos 2 t) 02 sin sin( (gt) and FM chooses Oe 1 - L, = L. Both e-prediction diffusion (Salimans & Ho, 2022) and FM are trained by matching the conditional velocity V, = dix + ofE such that a neural network Gz(X,, approximates Ex,e [v][x.] Samples can then be generated via probability-fie ODE (PF-ODE) dx. Go(x,, t) starting from (),
We introduce Inductive Moment Matching (IMM). a method that trains a model of both high quality and sampling ef- ficiency in a single stage. To do so, we assume a time- augmented interpolation between data (distribution at 0) and prior (distribution at t 1) and propose learning an implicit one-step model (i.e. a one-step sampler) that trans- forms the distribution at time / to the distribution at time 8 for any s < t (Section 3.1). The model enables direct one- step sampling from = to s = O and few-step sampling via recursive application from any t to any r < t and then to any until O; this allows to learn the model from its own samples via bootstrapping (Section 3.2).
di
Stochastic interpolants. Unifying diffusion models and FM, stochastic interpolants (Albergo et al., 2023; Albergo & Handen-Eijnden, 2022) construct a conditional interpolation ge (xx|xx, E) [I x), I) between any data x ~ ((x) and prior E ~ p(E) and sets constraints [x e) = e, cx(x,c) x, and 71 %o = 0. Similar to FM, determin- istic sampler can be learned by explicitly matching the con- ditional interpolant velocity
3.1. Model Construction via Interpolants
Given data q(x) and prior ~ p(E). the time- augmented interpolation x; defined in Albergo et al. (2023) follows x, ~ (x, |x, c). This implies a marginal interpo- lating distribution
0, e) + iiz where z ~ (0, such that Gx t) s Ex.ez [v./x.]. Sam- dx: = G).
pling is performed following the PF-ODE de
similarly starting from prior 6 ~ p(e).
qe (xx) = 1ʃ q (xxlxx ()((())((((xxc)
When 7: = 0 and Ix(x, E) 0, + orE for at, detined
(2)
2

---

Inductive Moment Matching
We learn model distribution implicitly defined by a one- step sampler that transforms (xx)) into 2s (Xx) for some S f. This can be done via a special class of interpolants, which preserves the marginal distribution qs (x,) while in- terpolating between x and x,- We term these marginal- preserving interpolants among a class of generalized inter- polants. Formally, we define x, as generalized interpolant between x and X, if, for all s [0. fl, its distribution follows
3.2. Learning via Inductive Bootstrapping
While sound, the naive objective in Eq. (7) difficult to optimize in practice because when is far from 8, the in- put distribution (x(x,) can be far from the target q, (x,). Fortunately, our interpolant construction implies that the model definition in Eq. (6) satisfies boundary condition (x,)) = P%), (x.) regardless of 0 (see Lemma 4), which indicates that p (x.) R (x.) when is close to s. Fur- thermore, the interpolant enforces P))(( ~ Pa) for any < t close to t as long as the model is continuous around f. Therefore, we can construct an inductive learning algorithm for (x,) by using samples from (x)))
Ueje(Xx| x+) = (x xz); 1
(3)
and satisfies constraints Laje (x, x₂) = Xt, Loje (x, (x,x,) x() = x, Taje - "ool = O, and 9qu (x,) x, 6) = qe (x |x, E). When t t= I, it reduces to regular stochastic interpolants. Next, we define marginal-preserving interpolants.
For better analysis, we define a sequence number for parameter 0, and function r(s, t) where 8 r(s, < t such that - (xx) learns to match Psj- 0u-1 (x.). We omit r's arguments when context is clear and let r(s,, t) be a finite decrement from t but truncated at S t (see Appendix B.3 for well-conditioned r(s, t))-
Definition 1 (Marginal- Marginnl-Preserrvnng Interpolants). A general- ized interpolant x, is marginal-preseriie if for all t E [0,1] and for all 8 E [0, 4). the following equality holds:
[o(x.) = lʃ as (x x,) (x) (xx
(4)
where
General objective. With marginal-preserving interpolants and mapping r(s, t), we learn 0n in the following objective:
((xx))xx))))(x))))) de.
qe (x(x
(xx)
C()0)) - Ez,t w0s,, ()MMDD (xx).
That is, this class of interpolants has the same marginal at regardless of t. For all t E [o, 1], we define our noisy model distribution at & E [0, as
where w(s, t) is weighting function. We choose MMD as our objective due to its superior optimization stability and show that this objective learns the correct data distribution.
(x,) = /f g./o (x,/xx x: )p' (x(xx) (x) (xx)xx, (6)
Theorem 1. Assuming r(s, t) is cell-conditioned, the in- terpolant is marginal -rrgeeilprreeeviingg and 6 is minimizer of Eq. (8) for each n with infinite data and network capacity, for all t E [o, 1), s E [0,t].
where the interpolant is marginal preserving and p (x)xx) is our clean model distribution implicitly parameterized as a one-step sampler. This definition also enables multi- step sampling. To produce a clean sample x given x, go (xx) in two steps via an intermediate (I) we sample x ~ P², (x(xx) followed by x, ~ (x) |X, x,) and (2) if the marginal of x, matches g, (xx), we can obtain x by pols (x(x,). We are therefore motivated to minimize di- vergence between Eq. (4) and (6) using the objective below.
lim MMD2 (x. 12 120000
6% ....).......... 0.
(9)
In other words, 0,, eventually learns the target distribution qs (x,) by parameterizing a one-step sampler o0 (x(xx).
4. Simplified Formulation and Practice
Naïve objective. As one can easily draw samples from the model, it can be naively learned by directly minimizing
We present algorithmic and practical decisions below.
4.1. Algorithmic Considerations
[(0) = Ez,t [x.)) P. (xx.))
(7)
Despite theoretical soundness, it remains unclear how to empirically choose marginal-preserringg interpolant. First, we present a sufficient condition for marginal preservation.
with time distribution p(s, t) and a sample-based divergence metric D)..)) such as MMD or GAN (Goodfellow et al., 2020). If an interpolant x, is marginal-preserving minimum loss is 0 (see Lemma 3). One might also notice the similarity between right-hand sides of Eq. (4) and (6). However, 4, (x.) = p2),(x.)) does not necessarily imply ph (x|x,) = q ( In fact, the minimizer (x(x,) is not unique and under mild assumptions, a deterministic minimizer exists (see Section 4).
then the
Definition 2 (Self-Consistent Interpolants). Given 8, E [0, 1]. S t, an interpolant x, ~ qs (x, x,) is self consistent if for all r E [s, t]. the following holds:
g.te (xx(xx = (x,xxx xx) (x, x, dxx, (10)
Note that is different from optimization steps. Advancing from n to n can take arbitrary number of optimization steps.
3

---

Inductive Moment Matching
In other words, x, has the same distribution if one (1) di- rectly samples it by interpolating x and x, and (2) first sam- ples any x, (given x and x²) and then samples x, (given x and x,)). Furthermore, selfscoonistency implies marginal preservation (Lemma 5).
t=0 0
Paaliised
(x.,t)
(2575)
DDIM interpolant. Denoising Diffusion Implicit Mod- els (Song et al., 2020a) was introduced as a fast ODE sam- pler for diffusion models, defined as
(X,1
(X.
DDIM(X,,X,8.. 8,8, = a To
- 200
x
+
Ixx (II)
attract
repulse
repulse
and sample x, = (x,, [x(x 8, t) can be drawn when Ex [x|] is approximated by a network. We show in Appendix C.1 that DDIM as an interpolant, i.e. Velt = 0 and [ xt) (DIM(X,; t), is self-consistent. Moreover, with deterministic interpolants such as DDIM, there exists a deterministic minimizer 22 (x) (x) of Eq. (7).
attract
[x,,8]
(8...)
(X.)
Figure 3. With self-consistent interpolants, IMM uses M particle samples (M = 2 is shown above) for moment matching. Samples from p e (xx) are obtained by drawing from (x)((x)xx) followed by guie (X= |x, xt). Solid and dashed red lines indicate sampling with and without gradient propagation respectively. After M sam- ples are drawn, sample x. is repulsed by x,, and attracted towards samples of x, and x through kernel function k(.,-)
Proposition 1. (Informal) Ysse = 0 and Lafe x1) fies mild assumptions, there exists deterministic pig (x)xx) that attains 0 loss for Eq. (7).
xx) satis-
See Appendix B.6 for formal statement and proof. This allows us to define p (x(xx)) = 8(x ge s. t)) for a neural network (x s, t) with parameter 8 by default.
An empirical estimate of the above objective uses M particle samples approximate each distribution indexed by t. In practice, we divide batch of model output with size B into B/ M groups within which share the same (s, () sample, and the objective is approximated by instantiating B //M number of M x M matrices. Note that the number of model passes does not change with respect to M (see Appendix C.4). A M = 2 version is visualized in Figure and a simplified training algorithm is shown in Algorithm 1, A full training algorithm is shown in Appendix D.
Eliminating stochasticity. We use DDIM interpolant, de- terministic model, and prior p(e) (0, o where ad is the data standard deviation (Lu & Song, 2024). As result, one can draw X, from model via X, f (x)) == DDIM(X,, (X,, 8, f), s,t) t) where x, ~ (x.).
Re-using Xt for Xy. Inspecting Eq. (8) and (6), one re- quires x, ~ gr (xx) to generate samples from the target distribution. Instead of sampling x, given new (x,c) pair, we can reduce variance by reusing x, and x such that x, DDIM(X,,,, r, This is justified because x, de- rived from x₁ preserves the marginal distribution q, (x,) (see Appendix C.2).
4.2. Other Implementation Choices
We defer detailed analysis of each decision to Appendix C.
Flow trajectories. We investigate the two most used flow trajectories (Nichol & Dhariwal, 2021; Lipman et al., 2022),
Stop gradient.
We set n to optimization step number, i.e.
advancing from n - 1 to n is a single optimizer step where 0n is initialized from -.- Equivalently, we can omit n from On and write 0,- as the stop-gradient parameter 0-
Cosine. a = cos (t or = sin()
OT-EM. ag = 1- t, 02 =
Network go 0x,,, We set go s, f) = Cskp (t)x + Cou Go (I)Xe, Croise (8), Croise (4) with neural net- work Go. following EDM (Karras al., 2022). For all choices we let Gin (t) 1/Vož + oo//aa (Lu & Song, 2024). Listed below are valid choices for other coefficients.
Simplified objective. Let x,, x/ be i.i.d. random variables from (xx) and x,, x/. are variables obtained by reusing X2, xt respectively, the training objective can be derived from the MMD definition in Eq. (1) (see Appendix C.3) as
(u(s;)t Is (asd)Ya)) (12)
(f) = Exxxx Xxxxxxxxxxxx x.xx
Identity. (akip()) = 0, Come (ome(t) 1.
+ k6 (Verr k6 (yacn'ar)) k (Yst.
Simple-EDM (Lu & Song, 2024). Cskip(t) = as/(a? +
af), (t) = -a V +
where No,t = (x1), y'a.- = (x), Vo.r = for (x+), Yerr- = for (x))), k(..-) is a kernel function, and wrs t) is a prior weighting function.
Euler-FM. (tkip(t) = 1, Cour (t) = -ta.. This is specific to OT-FM schedule.
We show in Appendix C.5 that f2, (x:: similarly follows the
4

---

Inductive Moment Matching
EDM parameterization of the form (x,) Cskip 1)xx Coal (oa(s, Go (t)x,; Croise (8); Croisee (f))-
Algorithm 1 Training (see Appendix D for full version)
Input: parameter 0, DDIM(X,,
x, s.t), B. M,
Noise conditioning Cnoise (-). We choose Croise (t) = ct for some constant 2 1. We find our model convergence relatively insensitive to c but recommend using larger c, egg- 1000 (Song et al., 2020b; Peebles & Xie, 2023), because it enables sufficient distinction between nearby r and t.
Output: learned 0
while model not converged do
Sample data v, label and prior with batch size B and split into B/M groups. Each group shares a (s, r, t) sample.
For each group, 24 + DDIM(e, 26, 1) For each group, 24 + DDTM I., t).
Mapping function r(s, t). We find that r(s, f) via constant decrement in 7) 0r/ar, works well where the decrement is chosen in the form of (Nmax- //hinn/21 / for some appropriate k (details in Appendix C.7).
For each instance, set with prob. p. Minimize the empirical loss LIMM(()) in Eq. (67). end while
Kernel function. We use time-dependent Laplace kernels of the form k,, (2, y) = exp(-((sx 1) max([|æ y 1122 /) D) for y IR some e > avoid undefined gradients, and (0ss t) (s, t)l. We find Laplace kernels pro- vide better gradient signals than RBF kernels. (see Ap- pendix C.8).
A Algorithm 2 Pushforward Sampling (details in Appendix F)
Input: model fo (4)] (O,, I) (optional) w Output: Xto
Sample XN -N(0,
for i NN..., do
X f- fl -l (xt,) or f4-- to,
Weighting w(s,t) t) and distribution t). We follow VDM (Kingma et al., 2021; Kingma & Gao, 2024) and define p(t) U(c, T) and p(s]t) = U(c, t) for constants e, T,TE [0, 1]. Similarly, weighting is defined
end for
put. Similarly, we define our guided model as f. (x+) Ckip(s, x)xx t)G (x,, s, t, c) where GD (x1, 3, c) is as defined in Eq. (14) and we drop (mn(-) and Guaise for notational simplicity. We justify this decision in Appendix E. Similar to diffusion models, с is randomly dropped with probability p during training without special practices.
w(s 1) = 20 - 2x) (- di
(b-
(13)
+
where (-) is sigmoid function, X denotes log-SNR, and a E (1, 2], o.), b E IR. are constants (see Appendix C.9).
We present pushforward sampling in Algorithm and detail both samplers in Appendix F.
4.3. Sampling
Pushforward sampling. A sample x, can be obtained by directly pushing X ~ qe(x.) through f (x.). This can be iterated for an arbitrary number of steps starting from NO, T) until s O. We note that, by definition, one application of fa, (x;: is equivalent to one DDIM step using the learned network go(x,, as the prediction. This sampler can then be viewed as few-step sampler using DDIM where go s, t) outputs a realistic sample x instead of its expectation Ex [x| xx] as in diffusion models.
5. Connection with Prior Works
Our work is closely connected with many prior works. De- tailed analysis is found in Appendix G.
Consisteney Models. Consistency models (CMs) (Song et al., 2023: Song & Dhariwal, 2023; Lu & Song, 2024) uses network go (x,, t) that outputs clean data given noisy input x,. It requires point-wise consistency go (x(x,, t) ga (x,,r) for any r < | where x,, is obatined via an ODE solver from x, using either pretrained model or groundtruth data. Discrete-time CM must satisfy go (x 0) = Xg and trains via loss Ex.xx, [d(go (x,, t). go- (x,] where d(-.)) is commonly chosen Le LPIPS (Zhang et al., 2018).
Restart sampling. Similar to Xu et al. (2023); Song et al. (2023), one can introduce stochasticity during sampling by re-noising a sample
higher noise-level before sampling again. For example, a two-step restart sampler from x; requires E (0, t) for drawing sample x = = (x.) where X, ~ qo (x, (x,1 (x:)).
We show in the following Lemma that CM objective with Lg distance is single- pprticle estimate of IMM objective with energy kernel.
Classifier-free guidance. Given a data-label pair (x, c), during inference time, classifier-free guidance (Ho & Sal- imans, 2022) with weight w replaces conditional model output Go(x,. s,t.c c) by a reweighted model output via
Lemma 1. When Xt = =x,, X, = =x,, k(æ, v) = -1 y11 and s > O is a small constant, Eq. (12) reduces to CM loss Exxxx (u(t))lg((x,))) - ge (x., r))r for some valid mapping r(1) <t.
wG wG 3, t, c) + (1 - w) G))e,, s, Ø)
(14)
where Ø denotes the null-token indicating unconditional out-
This single-particle estimate ignores the repulsion force
5

---

Inductive Moment Matching
imposed by -(;-)) Energy kernel also only matches the first moment, ignoring all higher moments. These decisions can be significant contributors to training instability and performance degradation of CMs.
Stochastic interpolants (Albergo et al., 2023; Albergo & Vanden-Eijnden, 2022) extend these ideas by explicitly defining a stochastic path between data and prior, then matching its velocity to facilitate distribution transfer. While IMM builds on top of the interpolant construction, it di- rectly learns one-step mappings between any intermediate marginal distributions.
Improved CMs (Song& & Dhariwal, 2023) propose pseudo- huber loss as -).,-) which we justify in the Lemma below.
Lemma 2. Negative pseudo-hube loss ke(z,)) y) = c
Diffusion distillation. To resolve diffusion models' sam- pling inefficiency, recent methods (Salimans & Ho, 2022; Meng et al., 2023; Yin et al., 2024; Zhou et al., 2024; Luo et al., 2024a; Heek et al., 2024) focus on distilling one-step or few-step models from pre-trained diffusion models. This two-stage approach currently has the best generation qual- ity and efficiency. However, such methods usually require jointly optimizing two networks and training requires care- ful tuning in practice avoid mode collapse (Yin et al., 2024). Another recent work (Salimans et al., 2024) ex- plicitly matches the first moment of the data distribution available from pre-trained diffusion models. In contrast, our method implicitly matches all moments using MMD and can be trained from scratch with a single model.
V ||l=- + e² for <> 0 is conditionally positive def- inite kernel that matches all moments of x and y where weights on higher moments depend on c.
From a moment moment-matching perspective, the improved per- formance explained by the loss matching all moments of the distributions. In addition to pseudo-huber loss, many other kernels (Laplace, RBF, etc.) are all valid choices in the design space.
We also extend IMM loss to the differential limit by taking r(s,,) 1) -> i, the result of which subsumes the continuous- time CM (Lu & Song, 2024) as a single-particle estimate (Appendix H). We leave experiments for this to future work.
Diffusion GAN and Adversarial Consistency Distillation. Diffusion GAN (Xiao al., 2021) parameterizes the gener- ative distribution as 1/): (x. [xx) = [J (x,((x.|x,, xt,o)xxxx.)0 (x) (x z, (zz)dzxx for s as fixed decrement from and p(z) noise distribution. It defines the interpolant Zsjo (x, |x. x2) as the DDPM posterior distribution, which is self-consistent (see Appendix G.2) and introduces random- ness to the sampling process to match ge (x(xx) instead of the marginal. Both Diffusion GAN and Adversarial Consis- tency Distillation (Sauer et al., 2025) use GAN objective, which shares similarity to MMD in that MMD is defined as an integral probability metric where the optimal discrimina- tor is chosen in RKHS. This eliminates the need for explicit adversarial optimization of neural-network discriminator.
Few-step generative models from scratch. Early one- step generative models primarily relied GANs (Good- fellow et al., 2020; Karras et al., 2020; Brock, 2018) and MMD (Li et al., 2015; 2017) (or their combination) but scaling adversarial training remains challenging. Re- cent independent classes of few-step models, e.g. Consis- tency Models (CMs) (Song et al., 2023; Song & Dhariwal, 2023; Lu & Song. 2024). Consistency Trajectory Models (CTMs) (Kim et al., 2023; Heek et al., 2024) and Shorteut Models (SMs) (Frans et al., 2024) still face training instabil- ity and require specialized components (Lu & Song, 2024) (e.g., JVP for flash attention) or other special practices (e.g.. high weight decay for SMs, combined LPIPS (Zhang et al., 2018) and GAN losses for CTMs, and special training sched- ules (Geng et al., 2024)) to remain stable. In contrast, our method trains stably with single loss and achieves strong performance without special training practices.
Generative Moment Matching Network. GMMN (Li et al., 2015) directly applies MMD to train a generator Go(z) where z ~ I), to match the data distribu- tion. It is a special case of IMM in that when t and r(s, t) = our loss reduces naive GMMN objective.
7. Experiments
We evaluate IMM's empirical performance (Section 7.1), training stability (Section 7.2), sampling choices (Sec- tion 7.3), scaling behavior (Section 7.4) and ablate our prac- tical decisions (Section 7.5).
6. Related Works
Diffusion, Flow Matching, and stochastic interpolants. Diffusion models (Sohl-Dickstein et al., 2015; Song et al., 2020b: Ho al., 2020; Kingma et al., 2021) and Flow Matching (Lipman et al., 2022; Liu et al., 2022) are widely used generative frameworks that learn score velocity field of noising process from data into a simple prior. They have been scaled successfully for text-to-image (Rom- bach et al., 2022; Saharia et al., 2022; Podell et al., 2023; Chen et al., 2023; Esser et al., 2024) and text-to-video (Ho et al., 2022a; Blattmann et al., 2023; OpenAI, 2024) tasks.
7.1. Image Generation
We present FID (Heusel et al., 2017) results for uncondi- tional CIFAR-10 and class-conditional ImageNet-256x256 in Table 1 and 2. For CIFAR-10, we separate baselines into diffusion and flow models, distillation models, and few-step models from scratch. IMM belongs to the last category
6

---

Inductive Moment Matching
Family
Method
HD (1)
Steps (4)
Fumily
Mechod
[ID(1)
Steps (1)
#Params
DDPM (Ho et 2021))
3.17
10081
GAN
GigatAAN (Kang 2023)
BigGAN (Brock, 2018)
3.45
6.95
56MM
112M
DDPM (Song et al., 2020h)
3.16
1000
StyloCIANXX. (Karm all, 20200
2.30
166M
NCSN++ (Song al., 20206)
2.38
1000
Diffacion
Solver (Lu ct 2022)
4.70
10
VQGAN (Faser et 2022)221)
26.52
1024
227M
& Flow
IDPM (Nichol Dhariwal, 2020)
EDM (Karaa et al. 2022)
2.05
2.90
4000
35
Misked
&
MAR all. .0244
VAR-120 (Tian et sl. 2024a)
MISCGIT (Chung ot.al, 2022)
2.57
6.18
10
227M
6I0M
400M
Flow Matching (Lipman 2022)
6.35
142
VAR (Tiam et al., 202(a)
1.92
10
20
Roctified Flow (Liu et al., 2022)
2.58
127
ADM (Dhariwal & Nicholl 2021)
10.94
250
55
PD (Salimans & Ho, 2022)
4.51
2
CDM et at. 20226)
4.8%
S100
2-Roctified Flow (Salimuns & Ho, 2022)
4.85
SumDitf (Haogeboom et 2023)
2.77
512
20
DFNO (Zhong al., 2023)
3.78
LDML+GG (Rombach al. 2022)
3.60
250
400M
(1 Lulhmann Lahhaa, 2021)
9:36
Diflusion Flow
U-DIT-L. (Tisn 2004h)
TRACT (Berthelot et al., 2023)
3.32
2
U-TEH (Ruo 2023)
2.29
50IM
Few-Step Distillation
PID (LPIPS) (Tee et al. 2024)
Diff fist-listrettttt (Lso al., 2024a)
5.57
3.92
DIT-XL22 (w
XL/2 (or 1.0) (Peebles & Xic, 2023)
(Peables 2023)
2012)
3.22
9.62
250
675M
DMD (Yin 20244)
3.77
1,5)
2.27
675M
CD (LPIPS) (Song et al., 2023)
CTM GAN) (Kim at.,, 2023)
1,87
2.93
2
SIT 1.5) (Ma 2004)
S/T:X/22 (e- 1.00 (Ma et al., 2024)
2.15
9.55
250
250
675M
675M
SiD (Zhou et al., 2024)
1.92
1
(Song 2003)
34.24
675M
SIM (L.o0 all, 20241s
20.3
675M
sCD (I.u Song, 2024)
2.52
2
Shoricut (Fruns et al., 2024)
10.60
675M
7.80
675M
ICT (Song & Dhariwal, 2023)
2.83
3.60
128
675M
2.46
2
New Step from Seratch
IMM (ours) (XL/2, x 1.25)
7.77
675M
BCT (Geng al., 2024)
3.60
5.31
675M
Few Step from Scratch
2.11
2
3.66
675M
sCT (Lu & Song. 2024)
2.97
2.77
675M
2.06
2
IMM (ours) (XL/2, 32 1.5)
8,05
675M
IMM (ours)
3.20
3.99
675M
1.98
2
251
675M
1.99
675M
Table 1. CIFAR-I0 results trained without label conditions.
Table 2. Class-conditional 12a6e26 results.
Positional amb
Pushionward, Ureiam
Fourier emb 15
Puahforward EDM
Restart, -u
500.5kk
4.5
10000
Reatart, EDM
Trannig Stupe
Steps
Training Steps
150g
Figure 5. More particles indi- cate more stable training on 255geNet
Figure 6. ImageNet- 256x256 FID with different sampler types.
Figure 4. FID convergence for different embeddiing (left). CIFAR- 10 samples from Fourier embedding (scale- = 16) (right).
7.2. IMM Training is Stable
in which it achieves state-of-the-art performance of 1.98 using pushforward sampler. For ImageNet-256x2255 we
We show that IMM is stable and achieves reasonable perfor- mance across a range of parameterization choices.
use the popular DiT (Peebles & Xie, 2023) architecture be- cause of its scalability, and compare it with GANs, masked and autoregressive models, diffusion and flow models, and few-step models trained from scratch.
Positional vs. Fourier embedding. A known issue for CMs (Song et al., 2023) is its training instability when us- ing Fourier embedding with scale 16, which forces reliance on positional embeddings for stability. We find that IMM does not face this problem (see Figure 4). For Fourier em- bedding we use the standard NCSN+ (Song et al., 2020b) architecture and set embedding scale to 16; for positional embeddings, we adopt DDPM++ (Song et al., 2020b). Both embedding types converge reliably, and we include samples from the Fourier embedding model in Figure 4.
We observe decreasing FID with more steps and IMM achieves 1.99 FID with 8 steps (with = 1.5), surpass- ing DiT and SiT (Ma et al., 2024) using the same architec- ture except for trivially injecting time s (see Appendix I). Notably, we also achieve better 8-step FID than the 10- step VAR (Tian et al., 2024a) of comparable size. At 16 steps, IMM also achieves 1.90 FID outperforming VAR's 2B variant (see Appendix I.4). However, different from VAR, IMM grants flexibility of variable number of inference steps. Lastly, we similarly surpass Shortcut models' (Frans et al., 2024) best performance with only 8 steps?. We defer infer- ence details to Scction 7.3 and Appendix I.2.
Particle number. Particle number M for estimating MMD is an important parameter for empirical success (Gretton et al., 2012; Li et al., 2015), where the estimate is more accurate with larger M. In our case, naively increasing M can slow down convergence because we have a fixed batch size B in which the samples are grouped into B/ M groups
NFF is twice the number of steps.
7

---

Inductive Moment Matching
5/2
Correlation 0.92
N2
6/2
B/2
34
1/2
F10 00X
XL/2
KC22
L/3kmm
Trining Compute (GFLOPs)
Sampling Computs (GFLOPs)
Transformer GFLOPs
20.2
Figure 7. IMM scales with both training and inference compute, and exhibits strong correlation between model size and performance.
Increasing sampling steps
Increasitng modde sine
Figure 8. Sample visual quality increases with increase in both model size and sampling compute.
means fewer t's are sampled for each step, thus slowing con- vergence. A general rule of thumb is to use a large enough M for stability, but too large for slowed convergence.
log abs, meen
1000
Gaallt+
Goslff= 1t
Noise embedding Cuoise (-). We plot in Figure the log ab- solute mean difference of t and r(s, t) in the positional em- bedding space. Increasing e increases distinguishability of nearby distributions. We also observe similar convergence on MengeNet-256x255 across different c, demonstrating the insensitivity of our framework w.r.t. noise function.
Goalf 10601
04
66
64
Traking Steps
Figure 9. Log distance in embedding space for Canise() (1) cf (lefi). Similar 125geN convergence across different (right).
of M where each group shares the same t. The larger M means that fewer t's are sampled. On the other hand, using extremely small numbers of particles, e.g. M = 2. leads to training instability and performance degradation, espe- cially on a large scale with DiT architectures. We find that there exists a sweet spot where a few particles effectively help with training stability while further increasing M slows down convergence (see Figure 5). We that in ImageNet- 256x256, training collapses when M = 1 (which is CM) and M = 2, and achieves lowest FID under the same com- putation budget with M = 4. We hypothesize M < does not allow sufficient mixing between particles and larger M
7.3. Sampling
We investigate different sampling settings for best perfor- mance. One-step sampling is performed by simple push- forward from T to (concrete values in Appendix I.2). On CIFAR-I0 we use 2 steps and set intermediate time t such that 7,, 1.4, a choice we find to work well empir- ically. On ImageNet-256x
we go beyond 2 steps and, for simplicity, investigate (1) uniform decrement & and (2) EDM (Karras et al., 2024) schedule (detailed in Ap- pendix I.2). We plot FID of all sampler settings in Figure with guidance weight w 1.5. We find pushforward sam-
8

---

Inductive Moment Matching
id/cos
id/FM
sEDMicos
sEDM/FM
cFM
FID-50k
CIFAR-10
3.77
3.45
2.39
2.10
2.53
m(s. =1
40.19
ImageNet-266x250
256x256
46.44
47.32
27.33
28.67
27.01
F10-00
8 9
+ELBO weight
96.43
Table FID results with different flow schedules and network parameterization.
+00
33.44
+1/(a3 + oZ)
27.43
plers with uniform schedule to work the best on ImageNet- 256x256 and use this as our default setting for multi-step generation, Additionally, concede that pushforward com- bined with restart samplers achieve superior results. We leave such experiments to future works.
Toining Steps
2003
Table 4. Ablation of weight- ing function w(s,)) () on 156geN
Figure 10. ImageNet-256x222 FID progression with different t, gap with M = 4.
conerann dec In
coroiaant dec in
corsiantt dec Ie
consraans A
7.4. Scaling Behavior
FID-su
concent in. in 1//
convontt 1//
6
Similar to diffusion models, IMM scale with training and inference compute as well as model size on ImageNet- 256x256. We plot in Figure 7 FID vs. training and in- ference compute in GFLOPs and we find strong correlation between compute used and performance. We further vi- sualize samples in Figure 8 with increasing model size, i.e. DiT-S, DÝT-B, DiT-L., DiT-XL,, and increasing inference steps, i.e. 1, 2, 4, 8 steps. The sample quality increases along both axes as larger transformers with more inference steps capture more complex distributions. This also explains that more compute can sometimes yield different visual content from the same initial noise as shown in the visual results.
504
1008
Figure 11. FID progression on different types mapping function r(s, f) for CIFAR-10 (left) and 2556 (right).
the form of (7max - T/minn/ 1/2² for an appropriately chosen k. We show in Figure 10 that the performance is relatively stable across k E {11, 12, 13) but experiences instability for k = 14. This suggests that, for given particle number, there exists a largest k for stable optimization.
Weighting function. Table we first ablate the weight- ing factors three groups: (1) the VDM ELBO factors ( दr Ar), ), (2) weighting or (i.e. when a 1), and (3) weighting 1/(a, + of). We find it necessary to use on jointly with ELBO weighting because it converts u-pred net- work to a e-pred parameterization (see Appendix C.9), con- sistent with diffusion ELBO-objective. Factor 1/(a/ + of) upweighting middle time-steps further boosts performance, helpful practice also known for FM training (Esser et al., 2024). We leave additional study of the exponent a to Ap- pendix 1.5 and find that 2 emphasizes optimizing the loss when is small while = 1 more equally distributes weights to larger t. As result, = 2 achieves higher quality multi-step generation than a = 1.
7.5. Ablation Studies
All ablation studies are done with DDPM++ architecture for CIFAR-I0 and DiT-B for ImageNet-256x225 FID compar- isons use 2-step samplers by default.
Flow schedules and parameterization. We investigate all combinations of network parameterization and flow sched- ules: Simple-EDM cosine (sEDM/ cos), Simple-EDM + OT-FM (FED Euler-FM + OT-FM (eFM), Iden- tity cosine /lld/cos) Identity + OT-FM (id/FM). Iden- tity parameterization consistently fall behind other types of parameterization, which all show similar performance across datasets (see Table 3). We see that on smaller scale (CIFAR-10), FMM works the best but larger scale (mageNNet2266x556 eFM works the best, indicating that OT-FM schedule and Euler paramaterization may be more scalable than other choices.
8. Conclusion
We present Inductive Moment Matching, framework that learns few-step generative model from scratch. It trains by first leveraging self-consistent interpolants to interpolate between data and prior and then matching all moments of its own distribution interpolated to be closer to that of data. Our method guarantees convergence in distribution and gen- eralizes many prior works. Our method also achieves state- of-the-art performance across benchmarks while achieving orders of magnitude faster inference. We hope it provides new perspective on training few-step models from scratch and inspire new generation of generative models.
Mapping function r(s, l). Our choices for ablation are (1) constant decrement in Th. (2) constant decrement in (3)
constant decrement in A = (at /a²). ). (4) constant incre- ment in 1/7,. (see Appendix C.6). For fair comparison, we choose the decrement gap so that the minimum t - r(s.t) is sJ 10-3 and use the same network parameterization. FID progression in Figure 11 show that (1) consistently outper- forms other choices. We additionally ablate the mapping gap using M = in (1). The constant decrement is in

---

Inductive Moment Matching
Impact Statement
Chen, N., Zhang, Y., Zen, H., Weiss, J., Norouzi, M., and Chan, W. Wavegrad: Estimating gradients for waveform generation. arXiv preprint arXiv:2009.000 00000, 2020.
This paper advances research diffusion models and gen- erative AL, which enable new creative possibilities and de- mocratize content creation but also raise important consider- ations. Potential benefits include expanding artistic expres- sion, assisting content creators, and generating synthetic data for research. However, we acknowledge challenges around potential misuse for deepfakes, copyright concerns, and impacts on creative industries. While our work aims to advance technical capabilities, we encourage ongoing discussion about responsible development and deployment of these technologies.
Dhariwal, P. and Nichol, A. Diffusion models beat gans on image synthesis, Advances in neural information processing systems, 34:8780-8794, 2021.
Esser, P., Rombach, R., and Ommer, B. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVE conference on computer vision and pattern recogni- tion, pp. 12873-12883, 2021.
Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F,, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference on Machine Learning, 2024.
Acknowledgement
Frans, K.. Hafner, D., Levine, and Abbeel, P. One step diffusion via shortcut models, arXiv preprini arXiv:2410. 2024.
We thank Wanqiao Xu, Bokui Shen, Connor Lin, and Sam- rath Sinha for additional technical discussions and helpful suggestions.
Geng, Z., Pokle, A., Luo, W., Lin, J., and Kolter, J. Z. Consistency models made easy. arXiv preprint arXiv:2406. 14548, 2024.
References
Goodfellow, L. Pouget-Abadie, J., Mirza, M., Xu, B., Warde- Farley, D., Ozair, S.,, Courville, A., and Bengio, Y. Generative adversarial networks. Communications of the ACM, 63(11): 139-144, 2020.
Albergo, M. S. and Vanden- Vanden-Eijnden, E, Building normalizing flows with stochastic interpolants. arXiv preprint arXiv:2209.1 1557,,, 2022.
Gretton, A., Borgwardt, K. M.. Rasch, M. J., Schölkopf. B., and Smola, A. A kernel two-sample test. The Journal of Machine Learning Research, 13(1):223-773 2012.
Albergo, M. S., Boffi. N. M., and Vanden- Eijuden E. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.00 08797, 2023.
Heek, J., Hoogeboom, E., and Salimans, T. Multistep consistency models. arXiv preprint arXiv:2403.06600 2024.
Auffray, Y. and Barbillon, P. Conditionally positive definite ker- nels: theoretical contribution, application to interpolation and approximation. PhD thesis, INRIA, 2009.
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochre- iter, S. Gans trained by two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Bao, F., Nie, S., Xue, K., Cao, Y., Li, C., Su, H., and Zhu, J. All are worth words: A vit backbone for diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 22669 22679, 2023.
Ho, J. and Salimans. T. Classifier-fic free diffusion guidance. arXiv preprint arXiv:2207.12552 12598, 2022.
Berthelot, D.,, Autef, A., Lin, J., Yap, D. A., Zhai, S., Hu, S., Zheng, D., Talbott, W., and Gu, E. Tract: Denoising diffuusion models with transitive closure time-distlllation. arXiv preprint arXiv:2303.0424 2023.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. Advances in neural information processing systems. 33:6840-66 (000-6851222 2020.
Ho, J.. Chan, W., Saharia, C., Whang, J., Gao, R., Gritsenko, A., Kingma. D. P., Poole, B., Norouzi, M., Fleet, D. J., et al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02.0 0230000,, 2022a.
Blattmann, A., Dockhor, T., Kulal, S., Mendelevitch, D., Kilian, M., Lorenz, D., Levi, Y.. English, Z. Voleti. V.. Letts, A., et al. Stable video diffiusion: Scaling latent video diffusion models to large datasets. arXiv preprint arXiv:2311.15127) 2023.
Ho, J., Saharia, C., Chan, W., Fleet. D. J., Norouzi. M., and Salimans. T. Cascaded diffusion models for high fidelity image generation. Journal of Machine Learning Research, 23(47): 1-33, 2022b.
Brock, A. Large scale gan training for high fidelity natural image synthesis. arXiv preprint 1rX9v:1809. 1096, 2018
Chang, H., Zhang, H., Jiang, L., Liu, C., and Freeman, W. T. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11315-11325, 2022.
Hoogeboom, E.,, Heek. J., and Salimans. T. simple diffusion: End- to-end diffuusion for high resolution images. In International Conference on Machine Learning, pp. 13213-13232. PMIR, 2023.
Chen, J., Yu, J., Ge, C., Yao, L., Xie, E., Wo, Y., Wang, Z., Kwok, J. Luo, P., Lu, H., et al. Pixart-o: Fast training of diffusion transformer for photorealistic text-to-image synthesis. arXiv preprint arXiv:2310.0042. 2023.
Kang, M., Zhu, J.-Y., Zhang, R., Park, J., Shechtman, E., Paris, S., and Park, T. Scaling up gans for text-to-image synthesis. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 10124-10134, 2023.
10

---

Inductive Moment Matching
Karras, T., Laine, S., Aittala, M., Hellsten, J., Lehtinen, J., and Aila, T. Analyzing and improving the image quality stylegan. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 8110-8119, 2020.
Luhman, and Luhman, T. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv: 10138228, 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, 2024a.
Karras, T., Aittala, M., Aila, T., and Laine, S. Elucidating the design space of diffusion- generative models. Advances in neural information processing systems, 35:26565 2657-,66777222 2022.
Luo, W. Huang, Z., Geng, Z., Kolter, J. Z., and Qi. G.-j. One-step diffussio distillation through score implicit matching. arXiv preprint ar Xrv: 16794, 2024b.
Karras. T., Aittala, M.. Lehtinen. J.. Hellsten, J. Aila, T., and
Laine, S. Analyzing and improving the training dynamies of
diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 24174- 24184, 2024.
Ma. N., Goldstein, M., Albergo, M. S., Boffi, N. M., Vanden- Eijnden, E., and Xie, S. Sit: Exploring flow and diffusion- based generative models with scalable interpolant transformers. arXiv preprint arXiv:2401. 00.00400000 2024.
Kim, D., Lai, C.-H., Liao, W.-.., Murata, N., Takida, Y., Uesaka, T., He, Y., Mitsufuji. Y., and Ermon, S. Consistency trajectory models: Learning probability flow ode trajectory of diffusion. arXiv preprint arXiv:2310.02277 2023.
Meng, C., Rombach, R., Gao, R., Kingma, D., Ermon, S., Ho, J., and Salimans, T. On distillation of guided diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 14297-14306, 2023.
Kingma, D. and Gao, R. Understanding diffusion objectives as the elbo with simple data augmentation. Advances in Neural Information Processing Systems, 36, 2024.
Müller, A. Integral probability metries and their generating classes of functions. Advances in applied probability, 29(2):429 443, 1997.
Kingma, D., Salimans, T., Poole, B., and Ho, J. Variational dif- fusion models. Advances in neural information processing systems, 34:21696-21700 2021.
Nichol, A. Q. and Dhariwal, P. Improved denoising diffusion probabilistic models. In International conference on machine learning, pp. 8162-8171, PMLR, 2021.
Kong, Z.. Ping, W.. Huang. J.. Zhao, K., and Catanzaro, B. Dif- fwave: versatile diffiusion model for audio synthesis, arXiv preprint arXiv:2009.0976 2020.
OpenAL. Video generation models world simulators. ht tttps: / /ooeeaii... . com//ora// 2024.
Li C.-L,, Chang, W.-C., Cheng, Y. Yang, Y., and Póczos, B. Mmd gan: Towards deeper understanding of moment matching network. Advances in neural information processing systems, 30, 2017.
Peehles, W. and Xie, S. Scalable diffusion models with transform- ers. In Proceedings of the IEEF/CVF International Conference on Computer Vision. pp. 1195-4205- 2023.
Podell, D.,, English, Z., Lacey, K., Blattmann, A., Dockhorn, T., Müller, J., Penna, J., and Rombach, R. Sdxl: Improving latent diffusion models for high-resolution image synthesis. arXiv preprint arXivv2307.01933
Li, T., Tian, Y. Li, H., Deng. M.. and He, K. Autoregressive image generation without vector quantization. arXiv preprint 2rXiv:2466.18888 2024.
2023.
Li, Y., Swersky, K., and Zemel, R. Generative moment matching networks. In International conference on machine learning, pp. 1718-1727. PMLR, 2015.
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 conference on computer vision 2022.
and pattern recognition, pp.
Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Lee, M. Flow matching for generative modeling. arXiv preprint arXiv:2210.2277..22 2022.
Saharia, C., Chan, W., Saxena, S., Li, L., Whang, J., Denton, E. L., Ghasemipour, K., Gontijo Lopes, R., Karagol Ayan, B., Sali- mans, T., et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in neural infor- mation processing systems, 35:36479 979444994999994 2022.
Liu, H. Chen, Z., Yuan, Y., Mei, X., Liu, X., Mandic, D., Wang, W., and Plumbley, M. D. Audioldm: Text-to-audio generation with latent diffusion models, arXiv preprint arXiv:2301 120.1 2023.
Salimans, T. and Ho, J. Progressive distillation for fast sampling of diffusion models. arXiv preprini arXiv:2202.005// 005122 2022.
Liu, X., Gong, C., and Liu, Q. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint Sii:22209000003. 2022.
Salimans, T., Mensink, T., Heek, J., and Hoogehoom, F. Multistep distillation of diffusion models via moment matching. arXiv preprint 2xXiv22006.04100 2024.
Lu. C. and Song, Y. Simplifying, stabilizing and scal- ing continuous-time consistency models. arXiv preprint arXiv:2410.1108 2024.
Sauer, A., Lorenz, D., Blattmann, A., and Rombach, R. Adversarial diffusion distillation. In European Conference on Computer Vision, pp. 87-103. Springer, 2025.
Lu, C., Zhou, Y. Bao, F. Chen, J., Li, C., and Zhu, J. Dom-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5787, 2022.
-bh--Dickttein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequilibrium thermody- namies. In International conference on machine learning, pp. 2256-2265. PMIRR 2015.

---

Inductive Moment Matching
Song, J., Meng, C., and Ermon, S. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020a.
Song. Y. and Dhariwal, Improved techniques training consis- tency models. arXiv preprint arXiv:2310.44189 2023.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochas- tic differential equations. arXiv preprint arXiv: 1.14466, 2020b.
Song, Y., Dhariwal, P., Chen, M., and Sutskever, I. Consistency models. arXiv preprint arXiv:2333301469) 2023.
Steinwart, 1. and Christmann, A. Support vector machines. Springer Science & Business Media, 2008.
Tee, J. J., Zhang, K., Yoon, H. S., Gowda, D. N., Kim, C., and Yoo, C. Physics informed distillation for diffusion models. arXiv preprint arXiv:2411.00
2024.
Tian, K., Jiang, Y., Yuan, 7., Peng. B., and Wang, L. Visual autoregressive modeling: Scalable image generation via scale prediction. arXiv preprint ar Xiv: 02905, 2024a.
Tian, Y.. Z. Chen, H., Hu, J. Xu, C., and Wang, U-dits: Downsample tokens in u-shaped diffusion transformers. arXiv preprint arXiv:2405.022
2024b.
Xiao, Z., Kreis, K., and Vahdat, A. Tackling the generative learn- ing trilemma with denoising diffiusion gans. arXiv preprint 222:2//22.7704 2021.
Xu, Y., Deng, M., Cheng, X., Tian, Y., Liu, Z., and Jaakkola, T. Restart sampling for improving generative processes. Advances in Neural Information Processing Systems, 36:76806-76838, 2023.
Yin, T., Gharbi, M., Zhang, R,, Shechtman, E.,, Durand, F., Free- man, W. T., and Park, T. One-step diffusion with distribution matching distillation. In Proceedings of the IEEE/CVF Confer- ence on Computer Vision and Puttern Recognition, pp. 6613- 6623, 2024.
Zhang, R., Isola, P., Efros, A. A., Shechtman, E., and Wang. O. The unreasonable effectiveness of deep features perceptual metric. Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 586-595, 2018
Zheng, H., Nie, W., Vahdat, A., Azizzadcnesheli, K., and Anandku- mar, A. Fast sampling of diffusion models via operator learning. In International conference on machine learning, pp. 42390- 42402. PMLR, 2023.
Zhou, M., Zheng, H., Wang, Z., Yin, M., and Huang, H. Score identity distillation: Exponentially fast distillation of pretrained diffusion models for seesttep generation. In Forty-first Interna- tional Conference on Machine Learning, 2024.
12

---

Inductive Moment Matching
A. Background: Properties of Stochastic Interpolants
note some relevant properties of stochastic interpolants for our exposition.
Boundary satisfaction. For an interpolant distribution qe (xx|x, e) defined in Albergo al. (2023), and the marginal q: (X;)
defined in Eq. (2), we can check that 91 (x1) p((1) and (x(x)) q (xx)) that X1 and Xo
/ʃ qi (xi x, (x)
qi (x1)
(15)
/| 6(x1 - =))(())(()dxdl
(16)
// ((x))p(x3)d )dx
(17)
p(x1)
(18)
/f 9o (xn x, ())(x))(()dxde
(xn)
(19)
// 6(xo x)g(x)p(e)dxxll
(20)
/|
((xo)p(e)dx
(21)
g(xo)
(22)
Joint distribution. The joint distribution of x and x, is written as
(x(x. = qi (x(x)xxee
(23)
Indepedence of joint at
ai (xx,x)) f (x((xixx,))) ()(x))((e)
(24)
-f 6(xl <)(()p(()dd
(25)
(())((x))
(26)
((x)p(()
(27)
in which case X1 = E.
B. Theorems and Derivations
B.1. Divergence Minimizer
Lemma 3. Assuming marginal-preservinn interpolant and metric D(:s)), minimizer 0* of Eq. (7) exists, i.e. 1.
ge (x(xx;)) and the minimum
Proof. We directly substitute qe (x(x:) into the objective to check. First,
// 2a/o (x, x (x.
(x,)
(28)
//
(x) (xx
(29)
(a)
ao (Xx)
(30)
13

---

Inductive Moment Matching
where (a) is due definition of marginal preservation. So the objective becomes
Eggt (x)
=E,,e [w]ss, ())(((((((()).
(31)
In general, the minimizer yr (x exists. However, this does not show that the minimizer is unique. In fact, the minimizer is not unique in general because a deterministic minimizer can also exist under certain assumptions on the interpolant (see Appendix B.6).
Failure Case without Marginal Preservation. We additionally show that marginal-preservation property of the interpolant (x)e(Xa |x, xx) is important for the naive objective in Eq. (7) to attain 0 loss (Lemma 3). Consider the failure case below where the constructed interpolant a generalized interpolant but not necessarily preeginin.. Then we show that there exists t such that P (x,) can never reach qa (xx) regardless of
Proposition 2 (Example Failure Case). Let gxx) 8(x), p(e) ô(E 1), and suppose interpolant Lsje(x, xx) (L 1)xx fx, and Yolt : 2 f1(t 1), then (x,) P2 (x.)) 0 for all < regardless of the learned distribution V/ (x(x:) given any metric D(,,)).
Proof. This example first implies the learning target
(x.) = // 9, (xxx,)x,))g(())((e))
(32)
= /f 8(x,, -(11 - s)x sc) 1)dxde 5(x, -8)
(33)
(34)
is a delta distribution. However, we show that if we select any t < 1 and 0 < s < t, P/ can never be delta distribution.
Px)c(xs)= = /ʃ Veje(xalxx, x (x)
(35)
/f 1/N((1-9)))
- -)x
5²
=2 )1 - t)dx,dx
8²
7x 1 -
(1
p
(36)
=fN(1 8 ixx 22 2a
$2
52
Dp' (x t)dx
(37)
Now, we show the model distribution has non-zero variance under these choices of and 8. Expectations are over p/ (x.)
or conditional interpolant (njo (x,)xx for all equations below.
Var(x,) = E [
E [x,]]
(38)
E
[x3] x, xx] p (x(xx)xxxxxxxx) (xxxxxxx IE [x.]²
(39)
// EE [x* x xx] 1 x )
[x,]²
(40)
= -// Var(x, x xi
|x, x,
p (x)) - 2E [x,] + E [x,]
(41)
/f
Var X
x,
Xx, (x.
1 (x
(42)
-Iʃ E [x Xx xxx p
x
(xx EE [xx]]²
=// Var(x, /x. xx) [x, x,,
E
[x,
[x, + E [x,
[x,
pējc (x/xx, (x, )xxdxx
(43)
(0)
14

---

Inductive Moment Matching
where (a) can be simplified
2 > 0 [x, x1] - E (xx1)
Var(xx|XxxX) + [x,
(44)
because (ar(x, x, x²) > 0 for all s < <<t due to its non-zero Gaussian noise. Therefore, Var(x,) > implying (x,) can never be delta function regardless of model Pa (x(x)). valid metric D) over probability p. (x,) and
qo (x,) implies
Pxa) - 0 + (xs) (x()))
which means
(x,) (x = P2 >
B.2. Boundary Satisfaction of Model Distribution
The operator output P (x.) satisfies boundary condition.
Lemma (Boundary Condition). For all E [0, 1] and all 0, the following boundary condition holds.
Us (x())) = (x.))
(45)
Proof.
// ao (x, x.) polo (x/xx, (x,5x))9,,1 (x.)dx.d
(x,)
8(x.-
I Pajo (xx(x (x,,
-x((x) / (x(xx))
= (xxx)
B.3. Definition of Well-Conditioned r(s,t) For simplicity, the mapping function r(s,t) is well-conditioned if
r(s,)) 1) (0,0,6,,-000
(46)
where A(t) 2 > 0 is a positive function such that r(s, t) is increasing for t 2 co(s) where co(s) is the largest t that mapped to s. Formally, co(s) = sup(t r(s,t) = s}. For 2 co(s), the inverse w.r.t t exists, i.e. (s,)) and rs, t)) = t. All practical implementations follow this general form, and are detailed Appendix C.6.
B.4. Main Theorem
Theorem 1. Assuming r(s, t) is well-conditioned, the interpolant is marginal-preserviee and 6*, a minimizer of Eq. (8) for each n with infinite data and network capacity, for all E [0, 1], s E t0,
lim MMD2 (1.(s) = 0. +00
(9)
Proof. We prove by induction on sequence number n. First, r(s, t) is well-conditioned by following the definition in Eq. (46). Furthermore, for notational convenience, we let r, 1 (s. = r-
(s,..., be n nested application
of r (s,,) on the second argument, Additionally, ro ((,, t) = f.
15

---

Inductive Moment Matching
Base case: 1. Given any 0, r(s, u) for all 8 a S co(s), implying
MMD² (x.) Peju (x,) MMD (x,), Paju (x) (x,)
(8)
(47)
for u S co(s) where (a) is implied by Lemma and (6) is implied by Lemma 3.
Inductive assumption: 7 - 1. assume Palu (Xs) = (x(x)) for all s S u S ry 2
8
We inspect the target distribution (s)r(s. (xx) in Eq. (8) if optimized
- - (8, cn(s)) On this interval, can apply r(s,-) :) to the inequality and get r(s, 8) I r(s, u) (s..oi) r co(s)) since r(s, is increasing. And by inductive assumption Psir(s.u) (x+) (x,) for r(, S ( cs) this implies minimizing
8
6.
E E..u (ss M)MMD (p) -:- (p )
(x,) P/.. (xx))
o < Susrr (m, is equivalent to minimizing
Egau w[s (x.
(48)
for s S uuSu surr (s, co())) Lemma 3 implies that its minimum achieves Paju (xx) = 9s (x,).
Lastly, taking - oo implies lim T (8, co()) = and thus the induction covers the entire [s, 1] interval given each Therefore, lim MMD(g, (x.)), (x.) for all 0 S 11515
B.5. Seff-Consistency Implies Marginal Preservation
Without assuming marginal preservation, it is important to define the marginal distribution x, under generalized interpolants qs (x, x, x1) as
(x)) - /f (x, /x, x xx)):(x(xxx (x, dx
(49)
and we show that with self-consistent interpolants, this distribution is invariant of t, i.e. (x)((xs) = (x(x)).
Lemma If the interpolant qatt (x, |x. x:) is self-consistent, the marginal distribution (o/r(X,)) as defined Eq. (49)
satisfies (x.)) (x)e(xa) for all & [s, 1].
Proof. For [s,
/ʃ xx)
(x,)
(50)
// (x)((x/Xxxxx) (xxx(x/xx,x)))) (( / 9:((x x ())( ()) dedx,dx
(51)
(x(),)
// (x xx)( / qe ce(xeexx x, edexxxd
(52)
/ʃ ((x)p(e) (f Xx) xx, e)dx; dedx
(53)
// ((x)p(E) (ʃ
Uote(xslx xx. dedx
(54)
(a) // 4x.1 (x, ()((x)p(())ddxx
(55)
/ʃ as (x, |x, c){(x)p((eddexx
(56)
(x,))
(57)
where (a) uses definition of self-consistent interpolants and (b) uses definition our generalized interpolant.
16

---

Inductive Moment Matching
We show in Appendix C.I that DDIM is example self-consistent interpolant. Furthermore, DDPM posteior (Ho et al., 2020; Kingma et al., 2021) is also self-consistent (see Lemma 6).
B.6. Existence of Deterministic Minimizer
We present the formal statement for the deterministic minimizer.
Proposition 3. for all E [0, 1), [0, c), Yafe = 0, [x,( xx) is invertible w.r.t and there exists C1 < 0o such that ex, Cillx ell, then there exists a function haje RR" + RD such that
(xxx) /ʃ 2ojo (0 xx)0(x - haje( (xx))))))(xxeeel))
(58)
(y, Proof. Let [ (x) be the inverse of oje w.r.t. x such that T (I,je(X, xx), = x and [i (y, xt), xx),x) x²) = y for all x, x P2 Since there exists O oo such that | (x, e)|
e|| for all [0, 1], the PF-ODE of the original interpolant (x,(x,e) e) Ix(x,e) E) exists for all & [0, 1] (Albergo al., 2023). Then, for all [0, 1], [0, f). let
(x))(Xe)) = =x: + I Ex,e
Ou [x xx du,
(59)
which pushes forward the measure (xx) to 9a (x.). We define:
1/ I (h) xt).
(60)
Then, since Ys = 0, Gole (x,,,xx,x,,xx = 6(x, L (x)) where x ~ 6(x - ha (xx)) Therefore,
Iajee xx) = I [ T/ (haje xx), xx) = haje (xee
(61)
whose marginal follows (x,) due to it being the result of PF-ODE trajectories starting from ge (xxx...
Concretely, DDIM interpolant satisfies all of the deterministic assumption, the regularity condition, and the invertibility assumption because it is a linear function of x and xs. Therefore, any diffusion or FM schedule with DDIM interpolant will enjoy a deterministic minimizer P (x|xx).
C. Analysis of Simplified Parameterization
C.1. DDIM Interpolant
We check that DDIM interpolant is self-consistent. By definition, (x)x(x,/xxxx) = 5(x, DDIM(X;, t, s)). We check
that for all < r Srst.
/
x, (x, =
6(xx- DDIM((,,,,,,))))555 DDIM(x,,
=6(x, - DDDIM(DDMMIxxxXA)......
where
DDMM(DDIM(x,,, = 0, (o. -xx)] a,x)
= a,x + (o./0///)00////)))) xx
= =0,x -a,x+ + (o5/00//) (xx)
DDIM(XX,X,1,,))
Therefore, 8(x, DDMM(DDDM(XXXX...
8)) 5(x, DDIM(X, x,t, s)). So DDIM is self-consistent.
It also implies a Gaussian forward process qi (x+[x) = .N(oxx, ofo l) as in diffusion models. By definition,
@t (x, (x(xx/))
q: (x e))(E)dee
t, 1) = a,x + o,e, which implies gxxxx(x) so that x, is deterministic transform given x and e, i.e., x; = DDIM(E,, x,
N(axx, T).
17

---

Inductive Moment Matching
C.2. Reusing x for x,
We propose that instead of sampling x, via forward flow a, + o,e we reuse x, such that x, DDIM(Xx x, r, t) to reduce variance. In fact, for any self-consistent interpolant, one can reuse x, via X, Or)o(X,, Xx, x,) and x, will follow qr (x,) marginally. We check
(x(x,) (a) (r|((x,)) lI C- (x) (x)x(( (x) ) /I arie(x))x, / ae (xi 1x, (x(x))()))x)
where (a) due to Lemma 5. We can see that sampling x, X2 first then x, ~ (x,,( xx) respects the marginal distribution gr (x,).
C.3. Simplified Objective
We derive our simplified objective. Given MMD defined in Eq. (1), we write our objective as
(0 = Es, w t)
Ex [a[ (xx), )] Ex. [k (xx), )]
(62)
2 H4
(a) Bytt (s, 1)
xx [k: [ (x)): 1: 1(F.Tr (x+);
(63)
H
Eu,t
w1s t)
Ex [k [M, (x)). k(f (xx), Ex,
- KCF2 (x.): {))]
(64)
- East w[s 1) ExIxx, xx,x (k: (k( (x)))))()))))))))))))))))))))( )), (MU,,
(x.),
(65)
(M(F,, (x))) )) A.f"", (xx). - (K( (24), (xx))
= Ex x x4 Saatt ()) [MC,, (xx) F2 (x4) + M(S (x+), f (x+)) - -k(f, (x:), f (xx)) k(f (x)), f
(66)
where (-,-) is in RKHS, (a) is due to the correlation between x, and x by re-using Xx.
C.4. Empirical Estimation
As proposed in Gretton et al. (2012), MMD is typically estimated with V-statistics by instantiating a matrix of size M x M such that a batch of B x samples, (x(")) B is separated into groups of M (assume B is divisible by M) particles B/M.M where each group share a (s", rf, (+) sample. The Monte Carlo estimate becomes
(1((5)) ((6k))))//A-4
B/M
.M
M
(0M = B/M
Σ
w(s
((.))) + k( (a() S2 (a MC (0 (10))) S2 (4
(67)
M
- (2 (1.)) S2 (c ((
Computational efficiency. First we note that regardless of M, we require only model forward passes one with and one without stop gradient, since the model takes in all instances together within the batch and produce outputs for the entire batch. For the calculation of our loss, although the need for M particles may imply inefficient computation, the cost of this matrix computation is negligible in practice compared to the complexity of model forward pass. Suppose forward pass for single instance is O(K), ), then the total computation for computation loss for batch of B instances is O(BK) + (BB)). Deep neural networks often has K > M, so O(BK) dominates the computation.
C.5. Simplified Parameterization
We derive f2,( for each parameterization, which now generally follows the form
f2. Cakp t)xt + com (s, Go t)xe,, Cavise (s), Cuoise (t))
18

---

Inductive Moment Matching
Identity. This is simply DDIM with z-prediction network.
Cnoise (s), Croiso )
(x,) (a,
a,
Xx adv oz
Jo
0f
oz
0t.
Simple-EDM.
]] -ix
f2,((1)) = (a. aa))
0g
a2
JdGo
X;
Cnoise (s), Cnoise (t)
07
2
+ of
Va + oj
Gayaz + to
0,04 + 0g00
0.0t 0,at odGo
X
Croise (8), Croise (1)
0i + X
a + to
dav +
When noise schedule is cosine, (x.) cos() 8)) sin() - 8) oa)o
Caoise (s); Croiiee (t)
oo
And similar to Lu & Song (2024), we can show that predicting X, DDIM(X, x, s,t) with L2 loss is equivalent to u-prediction with cosine schedule.
f2.((1)) 1((:
(10.-02
-e
Lo (xxx) )(
0t
0t
0,0t + 0,0t
0,01 - 0z0l
ogGo
X odV at + oz
Cnoise (s), Cnoise (t)
((0,
a,
a1
+ 0x X:
at + 0
Va + of
01
ts
o.Go
X1
Cnoise (s), Cnoise (t)
V (_vai of 0gg0; - 0g00t
(-
( at
x+
_x
aga; + Jg00
at + oo X1
JaV at o
oz
Guget
where
Va + oz
(as 0. ax)x x + -x, X₁
0g00 + 1.0t
Gurrgtt
2 + of
0lg 0g0g
Vai + +
(o, 0g 0 +
0.a0 got a,000 - Izod
(axx + (0xx 0,6)
0.0t 0,09
0, o7)
Vai + +o
((a,o, - a,,)) + of) + t - a,aio, x 000 a,a00s
0,0; 0,04
(. + of)
a + +
Va + of
(or,o, o,,0 + of) (0,0: (a,0: -
(as0, 0201)an 0z(0ī + oî)
0,00 0,00
of t
0,6 - 0,X
Vo + oi
This reduces to u-target if cosine schedule used, and deviates from v-target if FM schedule used instead.
19

---

Inductive Moment Matching
Euler-FM. We assume OT-FM schedule.
]]
(x,) ((1 -5) -) - (1--))) 1))
Xt tottaGGa
X₁ oaV an oz
Cnoise (s), Croise (t)
+ -xx
Cnoise (+), Croise (1) ) + 7x,
=(1 -)
7)
X loaGu
JdV Of oi
X₁ (t s)oaGo
Cnoise (8), Cnoieei (t)
od 0dVoz
0z
+ o
This results in Fuler ODE from X₂ to Xs. We also show that the network output reduces to e-prediction if matched with x, = -DDIM(X,, x, s,t). To see this,
w(s,t) t) (x.) ((11-)- - -1 t))x
8
8 xx)](
1
t
= =w(s, 1)
a.t (x.) (11 --)x [
+ xxx
-
l|
t
|
= -s
X 11 - G
Cuoise (+), Croissc
- -xx
+ s
oav at + oo
12
(s.t)
GgG
X,
Cnoiso (s), Croise (t)
(- 5) (
- xx+ ( 1)x1
OdV oav a oo
(t -
= w(s t)
G.Go
Caoiso (s), Cnoiso (t)
(t s) (11
8 8 x + ( 1)xx)
OJV at + t
Guzgi
where
G urgee
(1 x + ( -1)((1 (201=-)x+10 +)x+ + tc))
(t
(11 (1-2)) 8 ) + ( - 1+ 1+()x + (s-1)6)
(t 8)
- ((-=s)x + (s t)c)
X
which is u-target under OT-FM schedule. This parameterization naturally allows zero-SNR sampling and satisfies boundary condition at s = 0, similar to Simple-EDM above. This is not true for Identity parametrization using Gg as it satisfies boundary condition only at s > 0.
C.6. Mapping Function (s),t)
We discuss below the concrete choices for r(s, f). We use a constant decrement > 0 in different spaces.
Constant decrement in rt) = 77 = 01/az. This is the choice that we find to work better than other choices in practice. First, let its inverse be 17 1-).
r(s, t) max (s +))
We choose e = (7mxx - 7/2 for some k. We generally choose Mmax xe 160 and 7mmin 0. We find (10...., 15) works well enough depending on datasets.
20

---

Inductive Moment Matching
Constant decrement in
r(s,t) max (8, E)
We choose (T- -)/2
Constant decrement in X(1) log-SNR, = 2log(ar/ot).
Let its inverse be (-), then
r(8, max (8, (A(1) - +)))
We choose e (Amax /2 This choice comes close to the first choice, but we refrain from this because r(s, t) becomes close to both when 0 and instead of just f. 1. This gives more chances for training instability than the first choice.
Constant increment in 1/n(t).
10''(1NN (1/1(()) )
r(s,t) max s
We choose (1/m(t) 1/n((t mmx))))))
C.7. Time Distribution t)
In all cases we choose p(t) = (e. T) and p(s|t) = ((e, for some 2 0 and T < 1. The decision for time distribution is coupled with r(s,t). We list the constraints on p(s, t) for each r(s, t) choice below.
Constant decrement in vy(t). We need to choose < because, for example, assuming OT-FM schedule, 7) = =t/(1 t), one can observe that constant decrement in 4) when around We need to define T < 1 such that r(s, T) is not too close T for s reasonably far away. With Mama Rs 160, we choose T 0.994 for OT-FM and T 0.996 for VP-diffusion.
results in r(s, t) that is too close to due to 7/e" exploding gradient
Constant decrement in t. No constraints needed. T 1, 0.
Constant decrement in A One can similarly observe exploding gradient causing r(s, 1) to be too close to at both
1 R 0 and / x 1, so we can choose e > 0, e.g. 0.001, addition to choosing = 0.994 for OT-FM and T = 0.996 for VP-diffusion.
Constant increment in 1/n-- This experience exploding gradient for t 0, we require e 0, e.g. 0.005. And T 1.
C.8. Kernel Function
For our Laplace kernel k(x,, y) = exp(-i(s,,) max([|x vll2;c)/DD we 10-8, Looking its gradient w.r.t. x,
be a reasonably small constant, e.g.
m(x.t) =(o.0))l=~vll//22
| 2l12 e
V -(..) max([
0,
otherwise
one can notice that the gradient self-normalized to be unit vector, which is helpful in practice. In comparison, the gradient of RBF kernel of the form k(æ, y) exp( (xp(-1$ee - 1)([ v11 /D
ū i t) ((x.t) -- /D -y)
V ](..0)]==~111///
(68)
D
whose magnitude can vary a lot depending on how far x is from
For iss 1), we find it helpful write out the La loss between the arguments. For simplicity, we denote Go (x,, 8,
21

---

Inductive Moment Matching
G#
N1
Caoise (s); Canise (1)
ouva2 o0
(0ss t) (x)) fhr (x2)112
ts) cosi (s, - courls, (s, G)x(, - (akip(s,))xx + cour(s, r) s,
(ss t) (s, Go(
Cour (3, (csxip(s, r)x, + Cour (s, ( s,r) Cskip(s, sxxp(s;t)xx))
We simply set (8,,6) 6) C/ (8, 1) for the overall weighting be 1. This allows invariance of magnitude of kernels w.r.t. 6.
C.9. Weighting Function wiss
To review VDM (Kingma al., 2021), the negative ELBO loss for diffusion model
(O) 2 Bx,,,t [(-at 2e) (cco(xe,)) - <1l²
(69)
where 6g is the noise-prediction network and A log-SNR, The weighted-ELan BLBB loss proposed in Kingma & Gao (2024)
introduces an additional weighting function w(t) monotonically increasing in (monotonically decreasing in (og,
understood as a form of data augmentation. Specifically, they use sigmoid as the function such that the weighted ELBO
written
(w-ELBO( = Ex,c, o oob-- Xx) (- |cco(xe,)- <1p)
1
(70)
where (-) is sigmoid function.
The ag is tailored towards the Simple-EDM and Euler-FM parameterization as have shown in Appendix C.5 that the networks o2Go amounts to u-prediction in cosine and OT-FM schedules. Notice that ELBO diffusion loss matches instead of v. Inspecting the gradient of Laplace kernel, we have (again, for simplicity we let Gg( = X, Cnoise (8), Cuoise (t)))
Gol tuy/a?+o?
0 M(stt) (()(x))) ff (x.)))
00
w(s,t) - (x.) ||/ fa.(xi) - 8 (x,))
30 So, (xi)
(xx)
D
[(,((x)) 1 (x,)
112
ū(s,t) -...0)) (x))) || Go G
②
0g fo, (x,)
D
Go(Xi, Grunget
for some constant G Grarget- We can see that gradient à (x.) is guided by vector Gø(x,, s,t) Gurgett Assuming 36
Gx(X,, 8, t) is u-prediction, as is the case for Simple-EDM parameterization with cosine schedule and Euler-FM parame- terization with OT-FM schedule, we can reparameterize to e-prediction with eg as the new parameterization. We omit arguments to network for simplicity.
We show below that for both cases Eø Eurget =g(GG Guurge) for some constants Eurget and Grarges For Simple-EDM, we know z-prediction from u-prediction parameterization (Salimans & Ho, 2022), x a,X; o2 G.O. and we also know
22

---

Inductive Moment Matching
r-prediction from e-prediction, X9 (x, 0. We have
0:60 OtXt G,Go
(71)
X 0168 aix: a.o.Go
(72)
(1 a x + GG = 0:59
(73)
oix; + a.j, Go -FEE
(74)
60 oi + G,
(75)
6o Caarget oix: + o;Go - (oix; + ar Gungel)
(76)
Eg Eurget aG Gangel)
(77)
For Euler-FM, we know z-prediction from u-prediction parameterization, X = X, -tG and we also know z-prediction from e-prediction, Xe = (x: 1Eo)/(1
1). We have
x tcD
1-1 = -x -x (Go
(78)
X tEg (1 t)x, - t(1 G)Go
(79)
xx, + f(1 G) teg
(80)
6o x + (1 G)
(81)
Co X + Go
(82)
Eg Etargot xt + agGo - (xr + ar Gurgot)
(83)
60 Gurget = =r(GG Granget)
(84)
In both cases, (Go(x,, 8, Grarget) can be rewritten to (En(X,, 8, t) Etarget) by multiplying a factor dez, and the guidance vector now matches that of the ELBO-diffusion loss. Therefore, we are motivated to incorporate a into w(s t) as proposed.
The exponent for of takes value of either or We explain the reason for each decision here. When 1, we guide the gradient 8 f (xt), with score difference (Eo(Xe, 8, Enarge)) To motivate 2, first note that the weighted gradient
30
a ((s,t) t) 20 f2,(())) |cour (s, t)l 00 (x,) = = 00
0 Go (Xi, 8,t)
and as shown above that e-prediction is parameterized as
co(x,, 8, = GrGo (x+, s,t).
Multiplying an additional ar to ao Go s, therefore implicitly reparameterizes our model into an e-prediction model. The case of therefore implicitly reparameterizes our model into e-prediction model guided by the score difference (ea (Xe, s,t) t) Emrger). Empirically, of additionally downweights loss for larger compared to a,, allowing the model to train on smaller sime-steps more effectively.
Lastly, the division of of + of is inspired by the increased weighting for middle time-steps (Esser et al., 2024) for Flow Matching training. This is purely an empirical decision.
D. Training Algorithm
We present the training algorithm in Algorithm 3.
E. Classifier-Free Guidance
We refer readers to Appendix C.5 for analysis of each parameterization. Most notably, the network Go in both (1) Simple- EDM with cosine diffusion schedule and (2) Fuler-FM with OT-FM schedule are equivalent to u-prediction parameterization
23

---

Inductive Moment Matching
Algorithm 3 IMM Training
Input: model f data distribution g(x) and label distribution 4(c/x) (if label is used), prior distribution (o, o time distribution p(t) and p(s]t), DDIM interpolator ( x, 8, t) and its flow coefficients Ct,, os. mapping function r(s,t), kernel function ie(-, -), weighting function wa(s, t), batch size B, particle number M, label dropout probability p
Output: learned model fo Initialize n 0 A
while model not converged do
Sample a batch of data, label, and prior, and split into B/M groups, ((xb))) clf5), ( B/M,M
For each group, sample {(s), +=)== RMM
and r¹ r(st, 14) for each This results in a tuple ((s r (()) m (i.f) <- (DI xlt,) 11 axar(td) + (t,e v(i,j)
Feet
(i.5)
(1 < DDIM (53) alth) ("),)) V(i,j)
(Optional) Randomly drop each label e(f-5) to be null token ø with probability
On-1 optimizer step by minimizing LIMM(0.) using model fou (see Eq. (67)) (optionally inputting clt-)) into
network)
end while
in diffusion (Salimans & Ho, 2022) and FM (Lipman et al., 2022). When conditioned on label during sampling, it is customary to use classifier-frre guidance to reweight this u-prediction network via
G (Can (f)xe, Cavise (s), Cauise (t),c) = =wGo (Cim (t)xe, Xt. Cuoise( (s), Cuoise (t), + (w- (Cin)(((())))) Croise Caoise (t), 2))
(85)
with guidance weight w that the classifier-free guided fb, f,, (xx) is
fa (Xt: = Cakip(s, t)x, + Cour GV (Gi((t) (f)x,, Croise (8), Guoise (t),
(86)
F. Sampling Algorithms
Pushforward sampling. See Algorithm 4. We assume a series of N time steps (4) with T = IN > IN-1 > ...>
t2 > f > to = for the maximum time T and minimum time e. Denote aa as data standard deviation.
Restart sampling. See Algorithm 5. Different from pushforward sampling, N time steps {ta), do not need to be > 2 2 e (assuming T e). Different from
strictly decreasing for all time steps, e.g. T IN 2 1N-1 pushforward sampling, restart sampling first denoise clean sample before resampling noise to be added to this clean sample. Then clean sample is predicted again. The process is iterated for N steps.
G. Connection with Prior Works
G.I. Consistency Models
Consistency models explicitly match PF-ODE trajectories using a network gx t) that directly outputs a sample given any X2 ~ (x,... The network explicitly uses EDM parameterization satisfy boundary condition ge 0) Xg and trains via loss Ex, x [x,,, t) 9e- (x,, I|? where x, is deterministic function of x, from ODE solver.
We show that CM loss is a special case of our simplified IMM objective.
Lemma 1. When X1 = x, x, X,. k(o, v) -1|x v1l², and > 0 is a small constant, Eq. (12) reduces to CM loss Ex,xx ((t))lgg(xxtt)) go- (x,, r)1 for some valid mapping rt) <
Proof. Since X, xt, X, xx,, we have (x.) (x,( and fx,, f2 (x), So (x)) (2,(x))) for (x,)) = 0 by definition. Since k(x,v) = - yll it is easy to see Eq. (12) reduces to
(x.()) (xx
E Ex,,
(87)
24

---

Inductive Moment Matching
Algorithm 4 Pushforward Sampling
Input: model f" time steps {4)) prior distribution (00, of (optional) guidance weight w
Output: Xto
Sample x NN(
for i N₂ ,....:1 do
(Optional) <- 1 if N = I/ can optionally discard unconditional branch for N
Xt: fh 1.4. (Xt,) or fi:
a.t.. (X.,)
end for
Algorithm 5 Restart Sampling
Input: model fe, time steps (4) prior distribution (N oār), DDIM interpolant coefficients a, and o, (optional)
guidance weight w
Output: Xto
Sample xN ~ ~N
for i N,....1 do
(Optional) w < 1 if N =1
// can optionally discard unconditional branch for N =
x + flo.1. (xt,) or Flo.e.. (x::)
if i * 1 then
?~ N0, r
Xt. Ott, & + 0t: ,
|/ more generally Xt, ~ ge(xx; ,88 E)
else
Xta =x
end if
end for
where w(s,t) f) is weighting function. If is small positive constant, we further have (x,( Rs go (xxe))) where we drop s as input. If tx,,,,t) itself satisfies boundary condition at s O, we can directly take s = O in which case (x.) = (x,, f). And under these assumptions, our loss becomes
Ex,x. (w(() (x9(x,)) - (x.,r))1))
(88)
which is simply a CM loss using lg distance.
However, one can notice that from a moment mateehnn perspective, this loss significantly deviates from a proper divergence between distributions, and problematic in two aspects. First, it assumes aangleparticll estimate, which now ignores the entropy repulsion term in MMD that arises only during multi-particle estimation. This can contribute mode collapse and training instability of CM. Second, the choice of energy kernel is not proper positive definite kernel required by MMD. At best, it only matches the firs/ moment (its Taylor expansion cannot cover all moments as in RBF kernels), which is insufficient for matching two complex distributions! We should use kernels that match higher moments in practice. In fact, we show in the following Lemma that the pseudo-huber loss proposed in Song & Dhariwal (2023) matches higher moments as a kernel.
Lemma 2. Negative pseudo-hule huber loss (z(2,, y) c V/ vll² + e² for e > a conditionally positive definite kernel that matches all moments of and y where weights on higher moments depend c.
Proof. We first check that negative pscudo-huber loss c- v || yll + e² is a conditionally positive definite kernel (Auffray Barbillon, 2009). By definition, k(r, u) is conditionally positive definite for Ti, F E TRD and 61, Go E B with Σ C: =
(cyk(1,, 2
(89)
25

---

Inductive Moment Matching
We know that negative Lg distance -||æ y|
conditionally positive definite. We prove this below for completion. Due
triangle inequality, -|| vll 2 -|1=||| ||yll Then
Gej (-112; - x51l)
(- ||al||| - |,r Il)
(90)
im
cs |لرااا||
Σ Cj ==1
Ecllaa))
(91)
i=1
(a)
(92)
where (a) is due Ei= G O. Now since c
V1211 + +²² 2 -||=|| for all e 0, we have
Σ Σ city (c (c-v |1=
(c-vIla- - 23l1 +
Ge (- 11 - x,,1)> 2
(93)
i= ==1
So negative pseudo-huber loss is a valid conditionally positive definite kernel.
Next, analyze pseudo-huber loss's effect on higher-order moments by directly Taylor expanding V |||| + c²
0
VIE c - c=
||]|
||
+
||
5
|| +
0(1==11)
(94)
2c
8c
16c
128c
1
5
|
y1
8c3 ||
v1
+
||
v11
|| - -11l +0 - v/l
(95)
2c
16c5
128c
(96)
where we substitute = x - y. Each higher order || - yl1 for > 2 expands to a polynomial containing up to k-th moments, ie.. {z, 22........... .... 2² }, fy. x² ....... y to the weightings in front of each term.
thus the implicit feature map contains all higher moments where c contributes
Furthermore, we extend our finite difference (between r(s, t) and t) IMM objective to the differential limit by taking r(s, t) + & in Appendix H. This results in new objective that similarly subsumes continuous time CM (Song et al., 2023; Lu & Song, 2024) as single-particle special case.
G.2. Diffusion GAN and Adversarial Consistency Distillation
Diffusion GAN (Xiao al., 2021) parameterizes its generative distribution as
(x) xx)
[.]r(xx/Go(x, 2), xx (z))dl
where Gg is neural network, p(z) is standard Gaussian distribution, and (x/o(xx x,, is the DDPM posterior
(x x, xt) (Wo,,,
(97)
0,0
oi 0² a o? x
PQ ago!
+ a.
a 0²
2 -
a of
Note that DDPM posterior a stochastic interpolant, and more importantly, it is self-consistent, which we show in the Lemma below.
Lemma 6. For all 0 < < $ 1, DDPM posterior distribution t to s as defined in Eq. (97) is a sels-consistent Gaussian interpolant between x and Xt-
26

---

Inductive Moment Matching
Proof. Let x, ~ (r/e(x,, x, x1) and X3 2o/r (x,|X, x,), we show that x, follows Is e (x,, |x, xt).
O007 X ar(1
2
2 o² a oi c1
x,
x
0rl
(98)
0,0%
0 oz
0,00 s X, + X, a20,
o o
a o?
a,
0x
1
(99)
a o
0 o
where E1, C2 ~ NN(0 I) are i.i.d. Gaussian noise. Directly expanding
0,0;
2
a ot a
aj o²
a, o
x,
0²07
Xx + t,,(1
x
+
+o//
o₂ ?
E1
+ 0,
a o
x +
0z
E2
(100)
0,07
0
07 o
a 0²
a,0?
3 0² a of
a
ag97
2
X1
+
7
a o
?
0,
a oi
x
or1 1
E1
+ 0,
1
0 o
(101)
0
0i -
0,
1 a o²
a 0² a 00
7 Xt
0s (1
x
El + 0x
62
(102)
0,00
022 %
0r
0f d
(a) 0:03 ag
of o²
of 0² a of c3
X + a1
x
+
02
(103)
a oi
where (a) is due to the fact that sum of two independent Gaussian variables with variance a² and 62 is also Gaussian with variance a² + b², and cs ~ I) is another independent Gaussian noise. We show the calculation of the variance:
a o (1
at o. + )
a o² a
a
oe
af oia?
a g
+
0²-
a o2
a of
o g
of
a o²
0g o0
This shows x, follows (x)e(X,) x, xt) and completes the proof.
This shows another possible design of the interpolant that can be used, and diffusion GAN's formulation generally complies with our design of the generative distribution, except that it learns this conditional distribution of x given x, directly while we learn marginal distribution. When they directly learn the conditional distribution by matching P (x(x,) with qi (x(x,), ), the model is forced learn ge (x(x;) and there only exists one minimizer. However, in our case, the model can learn multiple different solutions because match the marginals instead.
GAN loss and MMD loss. We also want to draw attention to similarity between GAN loss used in Xiao et al. (2021); Sauer et al. (2025) and MMD loss. MMD is an integral probability metric over set of functions F in the following form
Dx(P, Q) = sup E Ex- - By-Q(s)fy)) fEr
where a supremum is taken on this set of functions. This naturally gives rise to adversarial optimization algorithm if F is defined as the set of neural networks. However, MMD bypasses this by selecting F as the RKHS where the optimal can be analytically found. This eliminates the adversarial objective and gives stable minimization objective in practice. However, this is not to say that RKHS is the best function set. With the right optimizers and training scheme, the adversarial objective may achieve better empirical performance, but this also makes the algorithm difficult to scale to large datasets.
G.3. Generative Moment Matching Network
It is trivial to check that GMMN is a special parameterization. We fix = 1, and due to boundary condition, r(s, = 0 (x,) = (xx) is the data distribution. Additionally, (xx, is simple pushforward of prior implies training target Peir
p(E) through network (e( where drop dependency on and s since they are constant.
27

---

Inductive Moment Matching
H. Differential Inductive Moment Matching
Similar to the continuous-time CMs presented (Lu & Song, 2024), our MMD objective can be taken to the differential limit. Consider the simplifed loss and parameterization in Eq. (12), we use the RBF kernel as our kernel of choice for simplicity.
Theorem 2 (Differential Inductive Moment Matching). Let f2,( be twice continuously differentiable function with bounded first and second derivatives, let k(-:-) be RBF kernel with unit bandwidth, x, ~ g(x), X2 ~ qi (x(x((x) x ~ (x(x(1x),, x, (x,, x, 1,r) r) and x,_ (D xx,1,, the following objective
CIMM-c((0,
lim p-2 (t r)2 Ex X 2 x k (ff (x), f., (x))) + 1o (x.), (x) (x,)
(104)
k (s,, (x)) St. (x.)) k (x),,
can be analytically derived as
xx -11 sixlilll (Uria
T (x))((x))
(105)
dt
af2,(xx)) (f2 f2 f2.((xt)) (F2,(xx)) (x)( T (x()
dt
dt
Proof. Firstly, the limit can be exchanged with the expectation due to dominated convergence theorem where the integrand consists of kernel functions which can be assumed to be upper bounded by 1, e.g. RBF kernels are upper bounded by 1, and thus integrable. It then suffices to check the limit of the integrand. Before that, let us review the first and second-order Taylor expansion of e =|==~=1? We let a, b R constants to be expanded around. We note down the Taylor expansion to second-order below for notational convenience.
e- <-=l==b|) a:
e -111-2 1a- 1a-001////2 ( T) -
-0)-a)+ a)+
7 - Te --n---11/// ((a-b)(a b)T - 1) (x a)
1
2
-lvloool12/2 at y b:
c-lbbal11 -e-110-al11200 - T) b) + 2 6) -1lb-al1 ((b a)(b - -a)T - 1) (y - b)
1
- 10-1==-11// at a,y=a = b:
e -laa-b11//2 1 ||0 /2 (a -) b) (x - a) - -1bbala 1// - a) (y -b)
2 (y b) - 110 / /2 ((b a)(b a)T 1) (v b)
1
(y - b) -100-a 1/2 ((b a)(b a) 1) (V - b)
+ 2
+ (x a)Te T -lb-all (1- (a - b)(a b) (y - b)
Putting together, the above results imply
e - 1>ll~v1///2 2 ||-1 e e 1o-bb) -1yy-a112/2
= (x a) -1b--al1? 2 (I - (a b) T) ) (y - b)
since it easy to check that the remaining terms cancel.
28

---

Inductive Moment Matching
Substituting S2,(( (xe), a Sa, (x,): y (x),((()) Sh,r (x,), furthermore have
lim P-r
(x 0) = lim p-t
(f (x,)(
(106)
(t-r) dl (xe) +(ll - -r13)
lim
(107)
r-t r
dt
(x, di
(108)
Similarly,
(x,)) f2, (x))(x()
lim
b) lim r
(109)
dt
Therefore, CIMM (0, can be derived
xx 1 (x)))) pd T
(110)
dl
(x) - f2((x1)) (x),((x)) SS, (x,( ") (x()
(111)
dt
(1 ( ( T (x)((x))
( dl.
Ex,x x
(112)
di
df (x)))
T =df(a((())
(2,( (F2,(x)) -
dl
dl
H.1. Pseudo-Objective
Due to the sto-ggradeent operation, we can similarly find pseudo-objective whose gradient matches the gradient of LIMM-OO (0, t) in the limit of + t.
Theorem Let (x)) be a twice continuously differentiable function with bounded first and second derivatives, k(.,-) be RBF kernel with unit bandwidth, x,x' q(x). X2 90(x; x), x4 (x(xx(xx) x,, xD, and
x DDIM(x,, x ,t,r), the gradient of the following pseudo-objective
VoLnMM.o lim (t rt)
1
VoEx, x-
lo (x)), Fe. (x))) + & { (x,), fa (xx)) (x+)) -
(113)
- ((),,(())) (x,,
can used to optimize and can analytically derived as
1 $2,7
(x()
12.7 (x.)
E Ex.x/
e
(114)
df (xx) dt
dfof (xx)
fo, (xi) - -/, (x()
s (xi) fh, (xx))
Volt.
dt
d df (xx)
df (xx)
fo, (xx) fh, (xx)) 11 Voste(x())
[ro, (xx)
foa (xx)
dt
29

---

Inductive Moment Matching
Proof. Similar the derivation of LoMM-oo (0, 1), let -2,,((x1) f2, (x,), - (x)), S2. (x,). have
lim (t (t-77 r) Volx - a) T f--1bb-al ||b-a|l;22 (a- - -b (y - b)
lim rot
(t e - 1101-a1//// 8/2
(y (y-b)) b) - -(a b)_ (y - b) (a b) Vox+
(a=
(x ax- ((a-b) b) (x a) a)) (a b) Voy
where
df (xx) di
lim 1 r-t (1 r) (y ()) - aff, (x() dl
lim r- r) (x (x)a)=
a)
Note that ae (x:) dr is now parameterized by B- instead of f because the gradient is already taken w.r.t. f², (x,) outside of the brackets, so (z a) and (y b) merely require evaluation at current with no gradient information, which 0 satisfies. The objective can be derived as
V VeLiMM (0
b 11 (x,) 1x)--11 12. (xoll( (
Ex
(d (d (xx) dt
[Fa (xx) - faa. (xx))
dso. (xx)
[Fa, (xx) fa, (xx)] T
Vott for (xx) - F2. (xx) 7) Vott Vostaalx)))
")
dt
d (fs (xx) dt
afor (xx)
[2a, (xi) - faz (xx)
dt
H.2. Connection with Continuous-Tinne CMs
Observing Eq. (105) and Eq. (110), we can see that when x/, and x, X being small positive constant, then (x() = 22 (x:), and exp -- f2, fl, (xx) = 1, and (x) RJ go t) where since is fixed we discard the dependency on 8 as input. Then, Eq. (105) reduces
d (x,r (xt)
Ex,
(115)
dt
which is the same differential consistency loss (Song et al., 2023; Geng al., 2024). And Eq. (110) reduces to
d (x+)
Ex,
Vof%,(xi)
(116)
dt
which is the pseudo-objective for continuous-time time CMs (Song et al., 2023; Lu & Song, 2024) (minus weighting function of choice).
I. Experiment Settings
L.I. Training& & Parameterization Settings
We summarize our best runs in Table Specifically, for ImageNet-256x256. we adopt a latent space paradigm for computational efficiency. For its autoencoder, we follow EDM2 (Karras et al., 2024) and pre-encode all images from
30

---

Inductive Moment Matching
CIFAR-10
15igeNet-256x256
Parameterization Setting
Architocture
DDPM++
DIT-S
DIT-B
DIT-L
DIT-XL
DIT-XL
GFlops
21.28
6.06
23.01
80.71
118.64
118.64
Purams (M)
55
33
130
458
675
675
Crose( (f)
1000
1000
1000
1000t
1000
1000t
2dd time conditioning
8 OT-FM Simple-EDM
Flow Trajectory
OT-FM Buler-FM
OT-FM Buler-FM
OT-FM Buler-FM
OT-FM Euler-FM
OT-FM
(x,.
Buler-FM
0.5
0.5
0.5
0.5
0.5
0.5
Training iter
400K
1.2M
1.2M
1.2M
1.2M
1.2M
Training Setting
Dropout
0.2 RAdum 10
Optimizer
AdamW 10-
Adam 10
Adam 10-8
Adam
Adam 10
Optimizere
10-
P1
0.9
0.9
0.9
0.9
0.9
0.9
82
0.999
0.999
0.999
0.999
0.999
0.999
Leurning Rate
0.0001
0.0001
0.0001
0.0001
0.0001 0 4096
Weight Decuy
0 4096
Batch Size
4096
4096
4096
4096
M
4 Laplace
4
4 Laplace
Kerne
Laplace
Laplace
Laplace
Laplace
r(s,t)
max(s,
(m
max(s, 1
(yp
Co
21
212
Minimum i gap p(t)
10-4
24(0.006, 0.994)
u(0, 0.994)
14(0, 0.994)
u400, 0.994)
t((0, 0.994)
((000.994)
1
5
4
EMA Rate
0.9999
0.9999
0.9999
0.9999
0.9999
0.9999
Lahel Dropout
0.1
0.1
0.1
0.1 Falsc Falsc
0.1 False Falsc
z-flip
True
False False
False False
False Falsc
EDM augment (Karras et al., 2022)
True
Inference Setting
Sampler Type
Pushforward
Pushforward
Pushforward
Pushforward Uniform
Pushforward Uniform
Pushforward
Number Steps
2
Schedule Type
1.4
Uniform
Uniform
Uniform
FID-50K (w 0)
1.98
FID-50K (10 = 1.0, i.c. no guidance)
42.28
26.02 9.69
9.33
7.25
6.69
FID-50K (10 = 1.5)
20.36
2.80
2.13
1.99
Table Experimental settings for different architectures and datasets.
ImageNet into latents without flipping, and calculate the channel-wise mean and std for normalization. We use Stable Diffusion VAE* and rescale the latents by channel mean [0.86488, -0.27787343, 0.21616915, 0.3738409) and channel std (4.85503674, 5.31922414, 3. 3.93725398, 3.9870003]. After this normalization transformation, we further multiply the latents by 0.5 so that the latents roughly have std 0.5. For DiT architecture of different sizes, we use the same hyperparameters for all experiments.
using mapping function r(s, t) by constant decrement in Te, keep Tmax RJ 160. This implies that for time distribution of the form ((e, set T 0.996 for cosine diffusion and T 0.994 for OT-FM. For we set it differently for pixel-space and latent-space model. For pixel-space on CIFAR-10, we follow Nichol & Dhariwal (2021) and set to a small positive constant because pixel quantization makes smaller noise imperceptible. We find 0.006 to work well. For latent-space on 1magcNet-256x256, we have no such intuition as in pixel-space. We simply set in this case.
Choices for T' and By default assuming
Exceptions occur when ablate other choices of r(s, l), e.g. constant decrement in Xe in which we 0.001 to prevent r(s, t) for being too close to when t is small.
Injecting time The design for additionally injecting s can be categorized into 2 types injecting directly and injecting stride size (1 - s). In both cases, architectural designs exactly follow the time injection of We simply extract positional time embedding of s (or - s) fed through 2-layer MLP (same as for t) before adding this new embedding to the embedding of after MLP. The summed embedding is then fed through all the Transformer blocks as standard DiT architecture.
31

---

Inductive Moment Matching
1-step
2-step
4-step
8-step
1-step
2-step
4-step
8-step
10-step
16-step
32-step
TF32
7.97
4.01
2.61
2.13
7.97
2.61
FID w
1,90
FPI6 w/
8.73
4.54
3.03
2.38
0 2
8.28
4.08
2.60
2.01
FPI6 w/ =
8.05
3.99
2.51
1.99
Table FID 155geN be- yond 8 steps presented in the main text. At 16 steps. IMM already outperforms VAR with 2B parameters (1.92 FID) and its performance saturates beyond that, with 32 steps performing marginally bet- ter.
Table 7. Ablation of exponent in the weighting function on ImageNet- 256x256. We see that multi-step generation while lagging slightly behind in 1-step generation.
Table 8. Ablation lower precision training maageNet-256x225 For lower precision training, we employ both minimum gap A 10-4 and (1 r) conditioning.
2 excels
Improved CT baseline. For 1mageNet-256xx566, we implement iCT baseline by using our improved parameterization with Simple-EDM and OT-FM schedule. We the proposed pseudo-huber loss for training but find training often collapses using the same r(s, I) schedule ours. We carefully tune the gap to achieve reasonable performance without collapse and present our results in Table 2.
1.2. Inference Settings
Inference schedules. For all one-step inference, directly start from ~N00 ( at time T to time through pushforward sampling. For all 2-step methods, we set the intermediate timestep 1 such that nt, = 1.4: this choice is purely empirical which we find to work well. For 2 4 steps we explore two types of time schedules: (1) uniform decrement in t with 17o < 7h < VJN where
N T
t,
(e (T-T)
(117)
and (2) EDM (Karras et al., 2022) time schedule. EDM schedule specifies 1jo < m
7/N where
1
N 1 and P
Tmmxx) ))
T
Mmmax + (Tmin
(118)
We slightly modify the schedule so that 77u = Mmmi is the endpoint instead of ?/ Mmin and fJo originally proposed, since our y/ can be set without numerical issue.
We also specify the time schedule type used for our best runs in Table 5 and their results.
I.3. Scaling Settings
Model GFLOPs. We reuse numbers from DiT (Peebles & Xie, 2023) for each model architecture.
Training compute. Following Peebles & Xie (2023), we use the formula model GFLOPs batch size training steps 4 for training compute where, different from DiT, we have constant 4 because for each iteration we have 2 forward pass and I backward pass, which estimated as twice the forward compute.
Inference compute. We calculate inference compute via model GFLOPs number of steps.
I.4. Scaling Beyond Steps
We investigate when the performance saturates for ImageNet-256x256 in Table 6. We see continued improvement beyond steps and at 16 steps our method already outperforms VAR with 2B parameters (1.92 FID).
I.5. Ablaiion on exponent a
We compare the performance between a = 1 and a = 2 on full DIT-XL architecture in Table 7, which shows how a affects results of different sampling steps. We observe that a = causes slightly higher 1-step sampling FID but outperforms a = in the multi-step regime.
32

---

Inductive Moment Matching
1.6. Caveats for Lower-Precision Training
For all experiments, we follow the original works (Song al., 2020b; Peebles & Xie, 2023) and use the default TF32 precision for training and evaluation. When switching to lower precision such FP16, we find that mapping function, i.e. constant decrement in ?hs cause indistinguishable time embedding after some MLP layers when is large. To mitigate this issue, we simply impose a minimum gap A between any t and for example, A 10-4. Our resulting mapping function becomes
A, r(s,t) = max min (t 1, (m(()) (o)(()))))) o1) -c)
Optionally, we can also increase distinguishability between nearby time-steps inside the network by injecting (r- () instead of as our second time condition. We use this as default for FP16 training. With these simple changes, we observe minimal impact on generation performance.
Lastly, if training from scratch with lower precision, recommend FP16 instead of BF16 because of higher precision that is needed to distinguish between nearby and r.
We show results in Table 8, For FP16, a = 1 causes slight performance degradation because of the small gap issue at large This is effectively resolved by a = 2 which downweights losses at large t by focusing on smaller t instead. At lower precision, while not necessary, a an effective solution achieve good performance that matches or even surpasses that of TF32.
J. A Aditiinna Visualization
We present additional visualization results in the following page.
33

---

Inductive Moment Matching
2
Figure 12. Uncurated samples on CIFAR-10, unconditional, steps.
34

---

Inductive Moment Matching
Figure 13. Uncurated samples on 2mageNe using D/T-XI//2 architecture. Guidance w 1.5, 8 steps.
35

---

Inductive Moment Matching
Figure 14, Uncurated samples on 2mageNe using D/T-XIL/2 architecture. Guidance w 1.5, 8 steps.
36


Reviews:
[{"review_id": "SaEjhve5Li", "reviewer": "Reviewer_SxrL", "summary": "This paper introduces a novel generative model enabling the few-step generation of high-quality photorealistic images. The approach builds on prior work in consistency trajectory models [1] and flow map matching [2], where a generator learns to move a noisy image from one timestep to another. However, unlike [1,2], which focus on pointwise matching along the Probability Flow ODE (PF-ODE), the proposed method instead matches the marginal distribution.\n\nThe training follows a bootstrapping approach similar to consistency models, while distribution matching is achieved via Maximum Mean Discrepancy (MMD), eliminating the need for auxiliary networks required in prior methods like DMD [3] and MMD [4].\n\nBeyond theoretical contributions, the paper introduces practical improvements for stable training and demonstrates strong performance on ImageNet with models trained from scratch.\n\n\n[1] Kim, Dongjun, et al. \"Consistency trajectory models: Learning probability flow ode trajectory of diffusion.\" arXiv preprint arXiv:2310.02279 (2023).\n\n[2] Boffi, Nicholas M., Michael S. Albergo, and Eric Vanden-Eijnden. \"Flow Map Matching.\" arXiv preprint arXiv:2406.07507 (2024).\n\n[3] Yin, Tianwei, et al. \"One-step diffusion with distribution matching distillation.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2024.\n\n[4] Salimans, Tim, et al. \"Multistep distillation of diffusion models via moment matching.\" Advances in Neural Information Processing Systems 37 (2024): 36046-36070.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Overall Update:\n\\\nWe thank all reviewers for their insights and helpful suggestions. We would like to announce several important updates.\n\n- To distinguish our method from other distillation-based post-training techniques, we change our title and model name from **“Moment Matching Self-Distillation”** to **“Inductive Moment Matching”**. \n- Better experimental results that outperforms the ones reported in submission.\n\n|     | CIFAR-10 FID |\n| -------- | :-------: |\n| 1-step  |  3.20    |\n| 2-step |  **1.98**     |\n \n\n|     | ImageNet-256x256 FID |\n| -------- | :-------: |\n| 1-step (w=1.25)  | 7.77   |\n| 2-step (w=1.25)   |  5.33     |\n| 4-step (w=1.25)   |  3.66     |\n| 8-step (w=1.25)   |  2.77     |\n| -------- | ------- |\n| 1-step (w=1.5)  | 8.05   |\n| 2-step (w=1.5)   |  3.99     |\n| 4-step (w=1.5)   |  2.51     |\n| 8-step (w=1.5)   |  **1.99**     |\n\n- Scaling beyond 8 steps for ImageNet-256x256\n\n|     | ImageNet-256x256 FID |\n| -------- | :-------: |\n| 10-step (w=1.5)  | 1.98   | \n| 16-step (w=1.5)   |  **1.90**     |\n| 32-step (w=1.5)   |   **1.89**  |\n\nThe results continue to improve beyond 8 steps. We see that at 16 steps it achieves 1.90 FID and already outperforms the 2B parameter VAR baseline (1.92 FID). We see saturation beyond 16 steps, with marginal improvement at 32 steps (1.89 FID).\n \n\\\n&nbsp;\n------------------------------------------\n## For reviewer SxrL:\n\nThank you for your comments and suggestions. We would like to address your concerns below.\n\n> I would like to see additional results on T2I or T2V applications. These experiments should be relatively easy to do if initialized from a pretrained diffusion model and could greatly enhance the case for broader adoption.\n\nWe test our algorithm on a text-to-image model trained on Datacomp [*] at 512x512 resolution. Samples using 8 steps can be found at [this link](https://drive.google.com/file/d/1e_2d0S1g4TStIkNtWE8UqGOO-BZZc9r0/view?usp=drive_link). While this is a preliminary result, we can see that our algorithm can easily transfer to T2I settings.\n\n\n\n> Q1: Would increasing the number of steps beyond 8 consistently improve performance? Is this improvement dependent on the training strategy, such as the choice of r(s,t)?\n\nYes the result continue to improve beyond 8 steps. See results above in “Overall Update” that 16 steps attain 1.90 FID and performance tends to saturate beyond that, with 32 steps attaining 1.89 FID. Notably, this outperforms the VAR variant with 2B parameters. We additionally show a comparison between 8 step generation and 16 step generation in [this link](https://drive.google.com/file/d/1T8h-k7IW4b4srOMZbBppuwFT-LRk1CLP/view?usp=drive_link). We notice that extending to 16 steps only give minor shifts in visual content in most cases. This demonstrates that 8 steps already give near optimal solutions -- our method scales *efficiently* with sampling compute. With low probability, there can be major shift in content although their low-frequency component look similar (i.e. they look alike from afar) (see row 2 col 3). Additionally, the relative improvement in FID should be similar across different $r(s,t)$, although their convergence rate can significantly differ. This is evident during training (see Figure 6) where constant decrement in $t$ noticeably lags behind constant decrement in $\\eta_t$.  \n\n \n\n\n\n\n> Q2: In MMD [1], a DDIM-style sampler is also used during inference. Could this be applied to the proposed method? Currently, it appears that only the consistency sampler is utilized.\n\nWe actually investigate both DDIM sampler (i.e. pushforward sampler) and consistency sampler (i.e. restart sampler) (see Sec 4.3). Pushforward sampler is equivalent to DDIM (while additionally injecting $s$ as conditioning) because $f_{s,t}^\\theta(x_t)$ is defined as a DDIM step from $t$ to $s$ with $g_\\theta(x_t, s, t)$ as the $x$-prediction network (see line 193-194). In fact, we also show that DDIM sampler is better than consistency sampler in Sec 7.3 and Figure 8. \n\n\\\n&nbsp;\n\n[*] Gadre, Samir Yitzhak, et al. \"Datacomp: In search of the next generation of multimodal datasets.\" Advances in Neural Information Processing Systems 36 (2023): 27092-27112.", "claims_and_evidence": "Yes, they are well-supported by theoretical proofs and strong empirical results.", "methods_and_evaluation_criteria": "This paper is well-motivated, and the proposed method is intuitively sound. It follows a recent trend in multi-step generation, where the generator is trained to produce samples resembling those encountered at intermediate noisy distributions rather than directly targeting the clean distribution. This approach aligns with methods such as Piecewise Rectified Flow [1], the consistency trajectory model [2], and MMD [3], among others. It is exciting to see a novel method emerge along the same lines but driven by a different distribution matching objective.\n\nFurthermore, the use of Maximum Mean Discrepancy (MMD) as a training signal is well-justified. It eliminates the need for auxiliary network training, simplifying the approach while enabling stable training and strong performance.\n\nBeyond these methodological innovations, the overall results are solid, with rigorous benchmarking, systematic comparisons, and convincing ablation studies.\n\n[1] Yan, Hanshu, et al. \"Perflow: Piecewise rectified flow as universal plug-and-play accelerator.\" arXiv preprint arXiv:2405.07510 (2024).\n\n[2] Kim, Dongjun, et al. \"Consistency trajectory models: Learning probability flow ode trajectory of diffusion.\" arXiv preprint arXiv:2310.02279 (2023).\n\n[3] Salimans, Tim, et al. \"Multistep distillation of diffusion models via moment matching.\" Advances in Neural Information Processing Systems 37 (2024): 36046-36070.", "theoretical_claims": "I didn't check the full proofs.", "experimental_designs_or_analyses": "Yes, the design is sound.", "relation_to_broader_scientific_literature": "It has broader implications for the generative modeling community, enabling end-to-end training of models that support both few-step and many-step inference. By eliminating the need for an additional distillation step during deployment, this approach streamlines the training-to-inference pipeline.", "essential_references_not_discussed": "The discussion on related works are adequate.", "other_strengths_and_weaknesses": "I would like to see additional results on T2I or T2V applications. These experiments should be relatively easy to do if initialized from a pretrained diffusion model and could greatly enhance the case for broader adoption.", "other_comments_or_suggestions": "n/a", "questions_for_authors": "Q1: Would increasing the number of steps beyond 8 consistently improve performance? Is this improvement dependent on the training strategy, such as the choice of r(s,t)?\n\nQ2: In MMD [1], a DDIM-style sampler is also used during inference. Could this be applied to the proposed method? Currently, it appears that only the consistency sampler is utilized.\n\n[1] Salimans, Tim, et al. \"Multistep distillation of diffusion models via moment matching.\" Advances in Neural Information Processing Systems 37 (2024): 36046-36070.", "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": "SXEBEVhyxo", "reviewer": "Reviewer_oBFV", "summary": "The paper introduces Moment Matching Self-Distillation (MMSD), a novel framework for training few-step generative models from scratch. MMSD offers a single-stage training procedure that avoids the need for pre-training or optimizing two networks. It leverages self-consistent interpolants to match the moments of its distribution to that of the data, ensuring distribution-level convergence.\n\n## update after rebuttal\n\nI thank the authors for their detailed response, which clarified my concerns about the use of classifier-free guidance in distillation models and the relationship between their method and Consistency Training. Given the authors' clarifications and commitments to improving notation and readability, I am happy to maintain my Accept recommendation.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the insights and suggestions. We would like to first refer the reviewer to the **Overall Update** section  in our rebuttal for reviewer SxrL for important updates to the paper. And we would like to address the concerns below.\n \n\n> The rationality of classifier-free guidance for distillation models is questionable, particularly for one-step models. The experiment also demonstrates that cfg = 1.5 yields inferior results compared to cfg = 1.25 when step = 1.\n\nIt is true that theoretically CFG in the diffusion/Flow Matching context does not work well in 1-step context because we are no longer modeling the velocity field at any $t$. In addition, the transition kernel defined by this 1-step CFG should be subject to an acceptance rate as in Metropolis adjusted MCMC (see [3] for details) without which the generated distribution can deviate from data distribution. In practice, since we cannot evaluate the likelihood, we always accept the samples. In general, however, we empirically find that adding some CFG values indeed helps with quality. We find that without CFG, 1-step sampler yields 12.21 FID on ImageNet-256x256, which is significantly higher than the results using CFG (7.97 and 7.12). We hypothesize that with 1 step, the linear combination of conditional and unconditional branches can still effectively correct visual details that are inaccurately modeled using either branch. Our conclusion is that CFG is helpful for 1-step sampler, but CFG values that achieve superior results in multi-step regime may not necessarily transfer to 1 step because 1-step CFG is not as well theoretically motivated as its multi-step counterparts. We call for more studies on this phenomenon in future works.\n\n> The writing and organization of this paper need further refinement to improve readability. Some proofs, such as the proof of Theorem 1, which I believe uses induction to extend from an infinitesimal step to a long-range step, are somewhat obscured by overly complex notation.\n\nWe will streamline our notation in proofs in our revision.\n\n\n> Questions regarding CT as a special case.\n\nWe stress an important point: minimizer of Consistency Distillation (CD) loss is the PF-ODE of the pretrained diffusion but minimizer of Consistency Training (CT) loss is NOT the PF-ODE. Minimizer of CT coincides with PF-ODE only when data distribution is delta distribution, as assumed by many proofs in the original CM paper [1] (Appendix B.1 Remark 4) and iCT [2] (Sec 3.2).\n\nTo see CT does not learn PF-ODE, recall CM loss $$\\mathbb{E}\\_{x_t,x,t}[w(t) || g_\\theta(x_t,t) - g_{\\theta^-}(x_r, r)  ||^2]$$, where $x_r$ is ODE solution from $x_t$. Assume $w(t)=1$, the minimizer of this loss is $$ g_{\\theta^*}(x_t,t) = \\mathbb{E}\\_{x|x_t}[g_{\\theta^-}(x_r, r)  ] $$. For CD, $x_r$ is ODE solution using pretrained score so $x_r$ does not depend on $x$, i.e. the conditional expectation can be dropped and $$ \\mathbb{E}\\_{x|x_t}[g_{\\theta^-}(x_r, r)  ] = g_{\\theta^-}(x_r, r)$$. However, for CT, $x_r$ depends on $x$, and the conditional expectation is irreducible. If we assume $g_{\\theta^-}(x_r, r) $ is a PF-ODE result, this new minimizer at $t$ deviates from any single PF-ODE due to the expectation. The original CM paper [1] and iCT [2] have a delta-distribution assumption that similarly allow elimination of the conditional expectation, which downplays the aforementioned problem in the general case.\n\nWe therefore call for understanding its loss at a distribution level, and find that it is a special case from this moment-matching perspective.\n\n> Our method does not guarantee learning of PF-ODE.\n\nWe do not guarantee our method learns PF-ODE. However, our method can converge to a different and equally valid solution whose *distribution* matches the data distribution. Consider a toy 2D Gaussian distribution $\\mathcal{N}(0,I)$ as data, and the same Gaussian $\\mathcal{N}(0,I)$ as prior. The PF-ODE is an identity function mapping any point $x$ to itself. However, consider another function $$f_{s,t}(x_t) = \\text{rotate}(x_t, 2\\pi*(t-s))$$ where $\\text{rotate}(x, \\phi)$ is a rotation operation around the origin by angle $\\phi$. Since Gaussian is rotation-invariant, distribution of $f_{s,t}(x_t)$ stays  $\\mathcal{N}(0,I)$. This is another valid solution under distribution matching objective which our method can also possibly learn.\n\n\n\\\n&nbsp;\n\n\n[1] Song, Yang, et al. \"Consistency models.\" (2023).\n\n[2] Song, Yang, and Prafulla Dhariwal. \"Improved techniques for training consistency models.\" arXiv preprint arXiv:2310.14189 (2023).\n\n[3] Du, Yilun, et al. \"Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and mcmc.\" International conference on machine learning. PMLR, 2023.", "claims_and_evidence": "The claims made in the submission are supported by clear and convincing evidence. The authors demonstrate the effectiveness of MMSD through extensive experiments on CIFAR10 and ImageNet 256×256.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are well-suited for the problem and application at hand. The use of benchmark datasets such as CIFAR-10 and ImageNet 256×256 is appropriate, as these are widely recognized and challenging benchmarks for evaluating generative models. The evaluation metric, Fréchet Inception Distance (FID), is a standard and reliable measure for assessing the quality and diversity of generated images.", "theoretical_claims": "I have gone through the theoretical proofs in the paper to a reasonable extent, and while I did not perform a 100% detailed verification, the proofs appear to be correct and well-justified.", "experimental_designs_or_analyses": "The experimental setups are well-structured and comprehensive. The authors evaluate MMSD on standard image benchmarks, including CIFAR-10, and ImageNet 256×256, which are widely used and respected in the field of generative modeling.", "relation_to_broader_scientific_literature": "The key contributions of the paper are related to diffusion/flow matching and distribution matching training/distillation, such as MMD-GAN, consistency model, and adversarial training.", "essential_references_not_discussed": "After reviewing the paper, I did not find any particularly critical or essential references that were missing from the discussion.", "other_strengths_and_weaknesses": "Strength: \n\n1. A key strength of this paper is its novel and efficient approach to training few-step generative models from scratch.\n2. By leveraging moment matching and self-consistent interpolants, MMSD guarantees distribution-level convergence and maintains stability across various settings, making it both robust and practical.\n3. The theoretical formulations are rigorous.\n\nWeakness:\n\n1. The rationality of classifier-free guidance for distillation models is questionable, particularly for one-step models. The experiment also demonstrates that cfg = 1.5 yields inferior results compared to cfg = 1.25 when step = 1.\n2. The writing and organization of this paper need further refinement to improve readability. Some proofs, such as the proof of Theorem 1, which I believe uses induction to extend from an infinitesimal step to a long-range step, are somewhat obscured by overly complex notation.", "other_comments_or_suggestions": "I do not have any additional comments or suggestions for the paper.", "questions_for_authors": "I don't think consistency models can be entirely considered a special case of this paper's framework. While consistency models can be understood from a distribution matching perspective, they are theoretically guaranteed to be trajectory-based distillation. Based on my understanding, this paper's method doesn't theoretically guarantee that the learned deterministic mapping is consistent with the PF-ODE mapping, since the framework is at the distribution level. And re-using $x_t$ for $x_r$ is aligned with consistency training only if r is very close to t. I'm not sure if this understanding is correct.", "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": "K8rRyRaxB9", "reviewer": "Reviewer_m3Qo", "summary": "# Update\n\nI gave the score of 4. My complaints are minor, and the authors have addressed them in the rebuttal. I'm comformtable with the paper being published in the form close what it is right now. So, I decided to not change my evaluation.\n\n# Old Summary\n\nThis paper presents moment matching self distillation (MMSD), an algorithm for training a generative model from scratch that is based on the stochastic interpolant framework.\n\nGiven a data distribution $q(x)$ and a noise distribution $p(\\epsilon)$, the stochastic framework constructs a stochastic process $\\\\{ q_t(x_t) : 0 \\leq t \\leq 1\\\\}$ such that $q_0(x_0) = q(x_0)$ and $q_1(x_1) = p(x_1)$. The generative model trained by MMSD is of the form $f^{\\theta}\\(x_t, s, t)$ where $\\theta$ denotes the model parameter. The specification is that, for any $0 \\leq s \\leq t \\leq 1$, if $x_t \\sim q_t$, then $f^{\\theta}(x_t, s, t) \\sim q_s$. A well-trained model can thus be used to generate a data sample from a noise sample in 1 evaluation of $f^\\theta$ or any other number of function evaluations.\n\nThe proposed algorithm is as follows. In each training iteration: \n1. We sample two times $s$ and $t$ such that $0 \\leq s < t \\leq 1$. We also compute an intermediate time $r = r(s,t)$ such that $s < r < t$.\n2. We then sample $N$ pairs of data items and noise vectors $(x_1, \\epsilon_1)$, $(x_2, \\epsilon_2)$, $\\dotsc$, $(x_M, \\epsilon_M)$ where $x_i \\sim q$ and $p_i \\sim p$.\n3. For each $1 \\leq i \\leq N$, we compute $x_{t,i} = \\alpha_t x_i + \\sigma_t \\epsilon_i$ and $x_{r,i} = \\alpha_r x_i + \\sigma_r \\epsilon_i$ where $\\alpha_t$ and $\\sigma_t$ are the noise schedule functions used in the formulation of $q_t$.\n4. Compute two sets of samples $\\\\{ y_{s,t,i}: y_{s,t,i} = f^{\\theta}(x_{t,i}, s, t) \\\\}$ and $\\\\{ y_{s,r,i} : y_{s,r,i} = f^{\\theta^-}(x_{r,i}, s, t) \\\\}$ where $\\theta^-$ is the gradient-stop version of $\\theta$.\n5. Compute the loss $\\mathcal{L}$ as the distance between the two sets of samples $\\\\{ y_{s,t,i} \\\\}$ and $\\\\{ y_{s,r,i} \\\\}$ using the minimum mean discrepancy (MMD).\n6. Update $\\theta$ with the gradient of $\\mathcal{L}$.\n\nThe algorithm is strikingly simialr to a previous algorithm, consistency training (CT) [1]. The main difference MMSD uses MMD to compute the loss, but CT simply computes the loss as $\\sum_i \\mathcal{d}(y_{s,t,i}, y_{s,r,i})$ where $d$ is a distance metric. Another difference is that, in CT, the time $s$ is always zero, but it is variable in MMSD.\n\nThe paper provides theoretical basis for the MMSD algorithm and shows that CT is one of its special case. It demonstrates that MMSD is effective by competitive FID scores acheived by MMSD-trained models on the CIFAR-10 and ImageNet 256x256 datasets. It also claims that MMSD training in more stable than some other algorithms, and that MMSD does not require hypertuning.\n\n*Citation*\n* [1] Song et al. \"Consistency Models.\" ICML 2023.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the insights and suggestions. We would like to first refer the reviewer to the **Overall Update** section in our rebuttal for reviewer SxrL for important updates to the paper. And we would like to address the concerns below.\n\n> Claims on (4) training stability and (5) not requiring extensive tuning.\n\nWe stress our abstract only claims that our method, similar to diffusion models, remains stable in the sense that common parameter choices can lead to a stable training process and decent quality, without the model becoming degenerate. We do not claim that they all obtain “good scores”, which may have been misinterpreted by the reviewer. In fact, even diffusion models require specific designs in the architecture to reach optimal scores – for example, the DiT paper found that AdaLN gives best performance with timestep conditioning, and in the EDM paper it is shown that variance preserving is much less sensitive to hyperparameter choices than variance exploding (see config A and config B in the paper).  For best performance in IMM, similar to diffusion and Flow Matching, we should expect some degree of tuning and different hyperparam choices. The stability is substantiated in Figure 4,5,6, and 7 which show that, as claimed in abstract, our method converges across different time embedding, time conditioning functions, $M$, $r(s,t)$, and $w(s,t)$ params. We introduced variability in these choices precisely to show our method’s stability across them.\n\nFor $M$, we emphasize that our convergence stability is not sensitive to $M$ as long as $M\\ge 4$, which is a reasonable range for practical purposes. No deep learning methods can train stably for *all* parameter choices (e.g. even for diffusion too large of lr or incorrect weighting for some $t$ can lead to degeneracy). We can only safely guarantee the stability exists for a large range of choices, which is sufficient for practical success. Similarly for $r$, as long as $r$ and $t$ has a reasonable finite gap (i.e. $k\\le 12$ in our case), we can achieve good performance.\n\n\nWe also argue $w(s,t)$ did not require extensive tuning and is simply carried over from VDM and Flow Matching. As motivated in Appendix C.9, since no weighting can be derived from MMD itself, we resort to the closest diffusion counterpart as in VDM [1], from which the terms $\\frac{1}{2}\\sigma(b-\\lambda_t) (-\\frac{d}{dt}\\lambda_t)$ are carried over. The term $\\alpha_t$ is also theoretically motivated to follow VDM gradients. These terms may look complex at first to suggest extensive tuning but they come as a group and do not require ablation of each individual subterms. The only additional  $\\alpha_t^2 + \\sigma_t^2$ term is simply motivated by [2] for Flow Matching schedule that upweights middle timesteps. The same effect can be achieved by sampling more middle timesteps instead.\n \n\n\nWe also use EDM conditioning as a simple unifying notation for different parameterization choices (e.g. DDIM and Euler sampler can both be unified under this one notation, see Appendix C.5). We do not reuse the complex conditioning introduced in EDM. In fact, we find the most simple Euler parameterization $f_{s,t}(x_t) = x_t + (s-t)G_\\theta(x_t,s,t)$ to work the best for ImageNet-256x256 (see Appendix C.5 and Table 4) and do not suggest deviating from this standard choice.\n\n\n\n\n> Discussion of Tee et al.'s physics-informed distillation (PID)\n\nPID is related as a distillation technique that explicitly matches the network’s own velocity field with that of diffusion. Different from other distillation methods, PID inputs noise $z$ and outputs $x_t$ whose derivative is matched with pretrained diffusion velocity field. The skip connection shares similarity with CM but requires $c_\\text{out}(T)=1$ instead of $c_\\text{out}(T)=0$, and different from distribution-matching distillation, it does not need to jointly train two networks.\n\nIn addition to PID, we will discuss the two additional works in our revision.\n\n> Shouldn't this be \"where each groups shares the same $t$ and $s$?\n\nYes. This group should share the same $t$ and $s$. \n\n\n> Wall-clock time training with different group sizes M but the same batch size B.\n\nWith $B=4096$ with DiT-XL architecture, we experimented across $M=2,4,8,16$ and find that all choices have per-step wallclock time of 0.53-0.55 seconds (one step here means a full optimization step accounting for both forward and backward passes). This is consistent with our analysis in Appendix C.4 that forward/backward pass is the computational bottleneck and time for computing the $M\\times M$ matrix is negligible. \n\n\\\n&nbsp;\n\n\n[1] Kingma, Diederik, and Ruiqi Gao. \"Understanding diffusion objectives as the elbo with simple data augmentation.\" Advances in Neural Information Processing Systems 36 (2023): 65484-65516.\n\n[2] Esser, Patrick, et al. \"Scaling rectified flow transformers for high-resolution image synthesis.\" Forty-first international conference on machine learning. 2024.", "claims_and_evidence": "The paper makes several claims.\n\n1. MMSD is theoretically sound.\n2. MMSD generalizes CT.\n3. MMSD results in competitive and fast generative models.\n4. MMSD training is stable.\n5. The abstract implies that MMSD does not require extensive tuning.\n\nI do not have particular problems with (1), (2), and (3). The FID scores on CIFAR-10 and ImageNet 256x256 are good given that the models were trained from scratch.\n\nHowever, (4) and (5) are not so evident. MMSD introduces a quite number of hyperparameters: the kernel function to use for MMD, the number of groups M inside a minibatch, how to compute $r(s,t)$, and the weighting function $w(s,t)$. It also inherits conditioning parmaeters from the EDM and EDM2 papers. These hyperparameters have to be picked carefully in order for the model to achieve good scores. (Some hyperparameters such as the weighting function $w(s,t)$, are derived from previous works, and I believe this is a form of extensive tuning.)  Moreover, the paper indicates that hyperparameters have impact on training stability. For example, picking $M \\leq 2$ or picking $r$ to be too close to $t$ lead to training to become unstable.", "methods_and_evaluation_criteria": "The paper uses CIFAR-10 and ImageNet 256x256, which are standard datasets for benchmarking generative models. It also uses the standard metric, the FID. There are no issues with these choices. The experiments are quite straightfoward as well: training models with MMSD and compare their FID scores to prior works. The paper includes scores upto NFE = 2 for CIFAR-10 and NFE = 8 for ImageNet 256x256 (along with guidance scale used), which makes comparison easy.", "theoretical_claims": "I skimmed through most proofs, but did not check all the calculations. The proofs were generally easy to follow and seem to be sound.", "experimental_designs_or_analyses": "The experiments are simple: training models with MMSD on benchmark datasets using well-known architectures and compare the resulting FID scores to previous works. I find no particular problems.", "relation_to_broader_scientific_literature": "Training a fast generative model in a stable and reproducible manner from scratch is an important problem. This paper contributes a novel and relatively simple algorithm that works well on widely used benchmarks, and I think this is an important contribution.", "essential_references_not_discussed": "The authors cite Tee et al.'s physics-informed distillation (PID) paper but did not discuss its approach. PID tries to make the network under training so that its velocity field conforms to ODE that defines the sampling trajectory of a diffusion model. This approach is quite different from the approach taken by the paper under review or other papers, and so a discussion can be illuminating. There are other several papers that take this approach that the authors might consider citing.\n\n*Citation*\n* [1] Boffi et al. Flow map matching. 2024.\n* [2]  Yang et al. Consistency Flow Matching:Defining Straight Flows with Velocity Consistency. 2024.", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "In Section 7.2, \"we haver a fixed batch size $B$ in which the samples are grouped into $B/M$ group of $M$ where each group shares the same $t$. Shouldn't this be \"where each groups shares the same $s$ and $t$?\n\nHow are $s$ and $t$ sampled during training? Am I missing something?", "questions_for_authors": "I would like to know how much wall-clock time training with MMSD is takes with different group sizes M but the same batch size B. Including this information in the appendix would be quite helpful for other practioners.", "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: The paper introduces Moment Matching Self-Distillation, renamed Inductive Moment Matching (IMM) during the rebuttal, as a new approach for pretraining unconditional or label-conditional deep generative models from scratch. IMM has a built-in capability to generate images in one or a few steps, while continuing to improve as the number of inference steps increases. The method generalizes the consistency model as a special case and illustrates how classical distribution matching techniques such as Maximum Mean Discrepancy (MMD) can be reimagined for pretraining high-performing generative models.

While the reviewers are all in favor of accepting the paper, they have identified several instances of potential overclaiming that should be moderated to more fairly represent prior work. For example, the manuscript criticizes existing methods as being unstable and requiring extensive tuning, while simultaneously adopting their well-studied architectures and hyperparameters. In particular, the abstract makes a broad assertion: “distilling them into few-step models often leads to instability and extensive tuning.” Such statements should be revised or substantiated with concrete evidence from the paper. Otherwise, there is a risk of misattributing isolated shortcomings of one method (e.g., Method A, a member of the distillation family) as general flaws of the entire methodological family that Method A belongs to. This kind of unsupported generalization could propagate misinterpretations in future literature, and hence should be avoided.

Upon reviewing the paper, the AC also notes that the claims would be better supported by including comparisons with consistency models and score distillation methods on ImageNet 512×512. Such comparisons would help position the performance of IMM more clearly.

The AC fully acknowledges the significance of the paper in demonstrating how to pretrain generative models from scratch with the built-in ability to generate high-quality samples in a few steps, and recognizes the strong performance achieved under this setting. However, it is important to clarify that state-of-the-art results for one or few-step generations for unconditional/label-conditional generations are still currently held by methods using diffusion pretraining combined with score distillation.

Another important point that requires clarification is the number of function evaluations (NFEs). For IMM, if the AC's understanding is correct, the NFEs should be calculated as **twice** the number of inference steps, due to the use of classifier-free guidance (CFG) during generation. In contrast, many distillation-based methods do **not** apply CFG directly to the generator, but rather to the teacher and/or an auxiliary network. Therefore, reporting only the number of inference steps—without accounting for the doubled computational cost from CFG—can lead to misleading comparisons across methods. If the AC’s assessment is accurate, the paper should explicitly report NFEs to support fair and transparent performance evaluations.

Finally, the AC would like to highlight that Reviewer m3Qo has provided an excellent summary of the actual implementation of the proposed algorithm—one that the AC finds clearer and more accessible than the current version in the paper. The authors are encouraged to incorporate these details to improve clarity and make the implementation more understandable for readers.
Decision: Accept (oral)