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.

Claims on (4) training stability and (5) not requiring extensive tuning.

We 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.

For $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.

We 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.

We 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.

Discussion of Tee et al.'s physics-informed distillation (PID)

PID 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.

In addition to PID, we will discuss the two additional works in our revision.

Shouldn't this be "where each groups shares the same $t$ and $s$?

Yes. This group should share the same $t$ and $s$.

Wall-clock time training with different group sizes M but the same batch size B.

With $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.

 

[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.

[2] Esser, Patrick, et al. "Scaling rectified flow transformers for high-resolution image synthesis." Forty-first international conference on machine learning. 2024.

Overall Update:

We thank all reviewers for their insights and helpful suggestions. We would like to announce several important updates.

• 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”.
• Better experimental results that outperforms the ones reported in submission.

• Scaling beyond 8 steps for ImageNet-256x256

The 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).

 

For reviewer SxrL:

Thank you for your comments and suggestions. We would like to address your concerns below.

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.

We 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. While this is a preliminary result, we can see that our algorithm can easily transfer to T2I settings.

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)?

Yes 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. 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$.

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.

We 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.

 

[*] 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.

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.

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.

It 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.

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.

We will streamline our notation in proofs in our revision.

Questions regarding CT as a special case.

We 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).

To 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.

We therefore call for understanding its loss at a distribution level, and find that it is a special case from this moment-matching perspective.

Our method does not guarantee learning of PF-ODE.

We 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.

 

[1] Song, Yang, et al. "Consistency models." (2023).

[2] Song, Yang, and Prafulla Dhariwal. "Improved techniques for training consistency models." arXiv preprint arXiv:2310.14189 (2023).

[3] Du, Yilun, et al. "Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and mcmc." International conference on machine learning. PMLR, 2023.