VCT: Training Consistency Models with Variational Noise Coupling

We thank the reviewer for the insightful comments.

W1. On Eq. (8), Typos, and Its Tightness

R1. The correct form of Eq. (8) should be:

N\\sum_{i=0}^N||f_\\theta(\\psi_{t_{i+1}}(x_0; x_1),t_{i+1})- f_{\\theta^-}(\\psi_{t_i}(x_0; x_1),t_{i})||^2.$$ This follows from the CS inequality. We establish its connection to the continuous-time CM as $ N\\rightarrow\infty $ and analyze their optimality. Applying a Taylor expansion under the assumption that $ \\Delta t := t_{i+1} - t_i = \frac{1}{N} $ (which can be relaxed) and the above inequality, we obtain: $$- \\log p_{\\theta}(x_0) \\leq \\frac{1}{2\\sigma^2}\\, E_{q_{\\phi}(x_1| x_0)} \Big\Vert x_0 - f_\\theta(x_1, 1)\Big\Vert^2 +\text{KL}(q_{\phi}(z | x_0)||p(z)) + C$$

\leq \frac{1}{2\sigma^2} N\sum_{i=0}^N \Big\Vert \frac{d}{dt} f_\theta(\psi_t, t) \Big\vert_{t=t_i} \Big\Vert^2 (1/N)^2 + \text{KL}(q_{\phi}(z | x_0) || p(z)) + C$$

$$$$

for a constant $C$.

Taking $N\\rightarrow \\infty$ is reasonable even in practical scenarios, as CT and ECT propose designing the $N$ scheduler in a coarse-to-fine manner. Additionally, we observe that at the optimal values of $\\theta$ and $\\phi$, the reconstruction losses (i.e., the first term in the upper bounds) specifically recover the consistency function. Consequently, the bound becomes tight at the optimum. We will incorporate this discussion in the camera-ready version.

W2. On Gradient Clipping.

R2. To evaluate the impact of gradient clipping, we applied this technique to our baseline model, iCT-VE, on CIFAR10. In the following table, we report results for both baseline and out VC method with and without gradient clipping (GC):

From the table it is clear that the main improvement comes from the learned coupling. Interestingly, in this case using GC for the baseline resulted in a performance improvement, even though generally GC is not mentioned in other CM literature, and we originally added it to our method to prevent early training instabilities for the learned noise distribution. However, we agree that applying GC to the baselines ensures a fairer comparison. We will include this discussion and report the results for all baselines with GC.

W3. On ImageNet.

R3. We also run the ImageNet experiments with 100k iterations. For ECM-LI-VC we increased $\beta$ to $\beta=100$ as $\beta=90$ diverged during training. Note that for runs with 100k iterations, we sometimes encountered divergences for our models with small $\beta$ and sometimes also for the baselines, while it seems to be solved when training for 200k iterations. The 1-step/2-step FID results are as follows:

For the settings with 100k iterations, our method performs similarly or slightly worse than the baseline. We believe this is due to the fact that the encoder for our model requires more iterations to learn the coupling, as demonstrated by the improved results for 200k iterations. We agree with the reviewer that including these results in the paper is important. We will add them, along with the corresponding results with OT coupling, to the camera-ready version.

W4. On $\beta$-Selection.

R4. In our experiments, $\beta$ was tuned with a coarse grid search with different values with a gap $10$. For iCT on CIFAR10, we initially tested the values $\beta=[10, 20, 30, 40]$, of which $\beta=30$ gave the best performance, then tested also for $\beta=[25, 35]$ which did not improve the performance. Similarly, for ECM, we tested for $\beta=[10, 20, 30, 40]$, and after achieving the best performance for $\beta=10$, we tested $\beta=[5, 15]$ which did not improve the performance. The tuning was done with the VE kernel and the best values were used also for the LI kernel. We used the best values of $\beta$ also in FashinMNIST and FFHQ without additional tuning. For ImageNet we observed on early runs that $\beta$ needed to be much bigger, so we initially tuned for $\beta=[30, 60, 90, 120]$. After achieving the best results with the VE kernel for $\beta=90$, we further tuned for $\beta=[70, 80, 90, 100, 110]$ for both VE and LI kernels, and found $\beta=100$ to be the best for VE and $\beta=90$ for LI. We will add a discussion about this process in the camera ready version, as we agree with the reviewer that guidelines on how to choose $\beta$ are important for the practitioners.

We are grateful to the reviewer for the careful and thoughtful review. Below we address some of the points raised by the reviewer, especially the correctness of Eq. (8), how our results compare with the baselines, and the novelty of the method.

W1. About Eq. (8).

R1. We thank the reviewer for pointing out the mistake in Eq. (8). The correct form of Eq. (8) should be:

N\\sum_{i=0}^N||f_\\theta(\\psi_{t_{i+1}}(x_0; x_1),t_{i+1})- f_{\\theta^-}(\\psi_{t_i}(x_0; x_1),t_{i})||^2.$$ This follows from the Cauchy–Schwarz inequality. We notice that this modification does not affect our discussion. In particular, we demonstrate the connection between our training objective (an upper bound of the negative log-likelihood as in VAE) and the continuous-time Consistency Model as $N\\rightarrow\\infty$ in **R1.** to the **Reviewer zkoD**. We will incorporate this correction and discussion in the camera-ready version. ### **W2. On Related Baselines and Their Comparisons.** **R2.** We agree with the reviewer that having additional baselines can be beneficial for the paper. Upon camera-ready version, we plan to add a more comprehensive table, similarly to Table 1 from TCM [3], and including results from other relevant consistency model works such as SCT from [5], sCT from [4], and TCM. ### **W3. On FID comparison.** **R3.** Regarding our FID results, on CIFAR-10 the aforementioned methods achieve 1-step/2-steps FID of 2.92/2.02 (SCT), 2.85/2.06 (sCT), and 2.46/2.05 (TCM), which are comparable to ours especially in the 2-steps regime. It is also important to consider that those improvements are orthogonal to ours and could in principle be combined. On Imagenet $64\times 64$, our results are generally worse than the ones reported by the aforementioned methods, but it is important to take into account that we used minimal settings in terms of network size and training budget due to computational constraints. Regarding the iCT baseline, there is no official open source implementation available, and not being able to reproduce the exact results seems to be a common problem found in other papers too (see for example [1,2]). ### **W4. On Novelty.** **R4.** We agree that the concept of coupling and its advantages have been explored before in other works. However, in the context of Flow Matching, it has been shown that coupling generally results in straighter trajectories which result in improved generation with less function evaluation, which does not entail that coupling would necessarily result in improved performance in CMs. While, as pointed by the reviewer, minibatch OT-coupling was already explored in CMs, our learned coupling shows improved scalability compared to minibatch OT with respect to data dimensionality and batch size, as we can see from our experiments on $64\times64$ images. ### **References** [1] Issenhuth, T., Santos, L. D., Franceschi, J.-Y., and Rakotomamonjy, A. Improving consistency models with generator-induced coupling. arXiv preprint arXiv:2406.09570, 2024. [2] Lee, J., Park, J., Yoon, J., and Lee, J. Stabilizing the training of consistency models with score guidance. In ICML 2024 Workshop on Structured Probabilistic Inference & Generative Modeling, 2024a. [3] Lee, S., Xu, Y., Geffner, T., Fanti, G., Kreis, K., Vahdat, A., and Nie, W. Truncated consistency models. arXiv preprint arXiv:2410.14895, 2024b. [4] Lu, C. and Song, Y. Simplifying, stabilizing and scaling continuous-time consistency models. arXiv preprint arXiv:2410.11081, 2024. [5] Wang, F.-Y., Geng, Z., and Li, H. Stable consistency tuning: Understanding and improving consistency models. arXiv preprint arXiv:2410.18958, 2024.

We thank reviewer for the comments and feedback. We address here some of the points and concerns raised.

W1. Figure 3 Does Not Have Enough Support.

R1. We believe that the initial disadvantage of our method compared to the baseline is due to the fact that the encoder is still in early training, which results in an initial increase in variance. As training proceeds and the encoder’s quality improves, the variance reduces accordingly and the FID performance of our method surpasses the one of the baseline. While at first the improvement can seem only marginal, we would like to highlight that for models that already perform relatively well on a given dataset, small FID improvements can correspond to a significant image quality enhancement.

W2. On Comparison with CD.

R2. We appreciate the reviewer’s suggestion that including additional baselines could strengthen the paper. Our focus is on training-from-scratch methods that do not assume access to a pre-trained teacher model. Thus, we primarily compare our approach with the CT counterpart. Nevertheless, we recognize the importance of a comprehensive evaluation and will include the FID results in a table for clarity in the camera-ready version.

W3. $\beta$-Weighting.

R3. We thank the reviewer for spotting the mistake. The correct formula for $\lambda_{kl}$ is $\beta \lambda_{ct}(t_N)$ when using the adaptive loss, while simply $\beta$ when using the weighting like in EDM. We will update the paper accordingly in the camera ready version.