Paper: Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion (ICLR-2024)

Reviewer comment (Reproducibility):
The paper itself presents a novel concept, but the methodology employed appears to be excessively complex. The distillation process involves three networks, and various techniques, such as $\lambda_{GAN}$ for ensuring stable training, are necessary to attain the results reported in Figure 15. Nevertheless, it is worth noting that this complexity also contributes to the uniqueness and innovation of the paper.

Author response:
Thank you for your valuable feedback. In response, we highlight CM's inherent speed-quality imbalance and then concentrate on elucidating the rationale behind CTM's design. The below arguments are integrated in our paper revision on Section 3.
- **[Why Not CM?]** CM, consisting of 2 model components (teacher/student) and 1 distillation loss (local consistency), lacks a balance between speed and quality (the red curve of Figure 9-(a) of original version). Our Theorem 2 (of first draft) proves that this unbalance is inevitable due to the overlapping between jump (from $t$ to $0$) to jump (from $t'$ to $0$), and will last as a fundamental issue as long as we train the jump only to zero-time. Therefore, CTM has been introduced to resolve this fundamental problem of jump overlapping.
- **[Why DSM Loss?]** A neural network, well trained for small jumps, will create high-quality sample for large-NFE. However, as the network gradient is proportional to jump step-size (Lemma 1 in the original manuscript), the neural jump for small step is challenging to train. This small-step jump precision can be boosted by applying DSM loss, serving as a regularizer for infinitesimally small jump in Figure 14-(b) of the original manuscript.
- **[Why Soft Consistency Distillation Loss?]** CTM (soft consistency) constructs distillation target $x_{target}$ using teacher with multiple random steps of ODE solving, while CM (local consistency) relies on teacher for only 1 step. This difference in teacher utilization leads to a relatively weaker signal for CM's distillation compared to CTM, as shown in Figure 13 (of the initial submission). Additionally, we observed that training with local consistency is prone to unstable to train and shows occasional divergence, whereas such instability has never seen in our experiments with soft consistency.
- **[Why GAN Loss?]** Now the question is whether we could also improve the quality of small-NFE samples. This challenge of long-jump training can be addressed by introducing GAN training, albeit with the addition of a new component, the discriminator. Specifically, the incorporation of GAN training enhances the precision of high-frequency (fine-details), as demonstrated in the comparison between Figure 11-(b) and 11-(c) of the first version. This finding aligns with the results from VQGAN [1]. For a detailed analysis, please refer to **Q3**.
- **[Student Beats Teacher]** It is known that DSM loss is closely linked to KL divergence [2,3,4], and GAN loss serves as a proxy for $f$-divergence [5] or IPMs [6]. Consequently, our overall loss function can be interpreted as an approximation of $$L_{distill}(teacher,student)+D_{KL}(data,student)+D_{f}(data,student).$$ Indeed, training with our loss function, we firstly observe that the well-established phenomenon in classification distillation, "student can be as good as teacher," is transferable to generation tasks. Here, $D_{KL}(data,student)+D_{f}(data,student)$ plays a role similar to the cross entropy, a statistical divergence between student and the oracle classifier, in the classification distillation. 

[1] Esser, Patrick, Robin Rombach, and Bjorn Ommer. "Taming transformers for high-resolution image synthesis." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021.

[2] Song, Yang, et al. "Maximum likelihood training of score-based diffusion models." Advances in Neural Information Processing Systems 34 (2021): 1415-1428.

[3] Kim, Dongjun, et al. "Soft Truncation: A Universal Training Technique of Score-based Diffusion Model for High Precision Score Estimation." International Conference on Machine Learning. PMLR, 2022.

[4] Kingma, Diederik P., and Ruiqi Gao. "Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation." Thirty-seventh Conference on Neural Information Processing Systems. 2023.

[5] Nowozin, Sebastian, Botond Cseke, and Ryota Tomioka. "f-gan: Training generative neural samplers using variational divergence minimization." Advances in neural information processing systems 29 (2016).

[6] Arjovsky, Martin, Soumith Chintala, and Léon Bottou. "Wasserstein generative adversarial networks." International conference on machine learning. PMLR, 2017.

Reviewer comment (Novelty):
One of the reasons for the popularity of diffusion models lies in their scalability with respect to data complexity. In Figure 15, it is evident that the distillation of diffusion models heavily depends on the underlying GAN training structure. This prompts the question of whether it might be more straightforward to solely employ GAN training for a single NFE generation. Essentially, this would involve passing through the computationally expensive pretrained Diffusion Model, followed by the resource-intensive distillation process (as CTM requires simulating trajectories), ultimately converging to the realization that the fundamental component driving performance improvement is the GAN training objective. Given the current state of the art where GANs can achieve superior results, it raises the question of why we necessitate a complex distillation methodology, notwithstanding the novelty of such an approach.

Author response:
We highly appreciate for the critical review and believe answering this feedback is important. As the reviewer points out, modern GAN models are highly compatible with diffusion models. In response to the reviewer's concerns, we have considered three factors: **Flexibility**, **Training**, and **Performance**.
- **[Flexibility]** Recent diffusion and distillation models offer a high degree of flexibility, allowing for easy fine-tuning with a minimal number of parameters or, in some cases, can be utilized without any additional training for tailored datasets. An illustrative example is presented in [7], where a combination of LoRA [8] parameters from distillation models and fine-tuned diffusion models enables the generation of samples with specific styles, eliminating the need for extra training of distillation models. In contrast, the development of such application-oriented techniques to GAN models poses greater challenges as they need to retrain GAN models for their specific usage. Additionally, CTM facilitates the likelihood computation and supports score-based inference techniques, which are not applicable to GANs. Given the significant divergence in potential use cases, we advocate for training the model with a combined adversarial loss and distillation loss, as opposed to relying solely on the adversarial loss.
- **[Training]** Training GAN models alone remains challenging. Although state-of-the-art GAN models work well with their proposed hyperparameters and network architectures, it requires expensive hyperparameter search if we retrain or fine-tune the generator with another dataset. To mitigate the instability, most modern GAN algorithms use progressive training across resolutions, whereas CTM training is end-to-end, meaning CTM training is resolution-agnostic. Additionally, in our experiments, CTM (even with adversarial loss) shows training stability and benign performance regardless of the hyperparameter choice because the existence of distillation loss and DSM loss significantly stabilize the training. Please refer to the paragraph **[Insight from VQGAN]** in the response of **Q3**.

|Model|NFE|FID$\downarrow$|IS$\uparrow$|Recall$\uparrow$|
|-|-|-|-|-|
||||||
|*Validation data*||1.41|64.10|0.67|
||||||
|**DM and GAN**|
|ADM|250|2.07|-|0.63|
|EDM|79|2.44|48.88|**0.67**|
|BigGAN-deep|1|4.06|-|0.48|
|StyleGAN-XL|1|2.09|**82.35**|0.52|
|**Distillation Models (NFE 1)**|
|PD|1|15.39|-|0.62|
|BOOT|1|16.3|-|0.36|
|CD|1|6.20|40.08|0.63|
|CTM (*ours*)|1|1.92|70.38|0.57|
|**Distillation Models (NFE 2)**|
|PD|2|8.95|-|0.65|
|CD|2|4.70|-|0.64|
|CTM (*ours*)|2|**1.73**|64.29|0.57|

- **[Performance]** We would like to compare samples from CTM and GAN in various metrics with the above updated table on ImageNet $64\times64$:
    - **[Inception Score]** StyleGAN-XL generates samples showing excessive confidence, as assessed by the metric IS, in comparison to validation data. This means that StyleGAN-XL tends to generate samples with a higher likelihood of being classified for a specific class, even surpassing the probabilities of real-world validation data. In contrast, CTM demonstrates an intermediate IS, positioning itself between the diffusion models and the GAN models. Notably, it closely aligns with the IS of the validation data.
    - **[Recall]** CTM outperforms GAN models in sample diversity. In fact, while CTM exhibits an intermediate level of recall, the diversity of its actual samples is comparable to that of EDM or CM, as visualized in Figure 16. Furthermore, Table 1 illustrates that CTM achieved state-of-the-art likelihood ($2.43$ bits/dim) and surpassed the teacher model (EDM) ($2.56$ bits/dim) that alleviates potential concerns regarding mode collapse in CTM.

[7] Luo, Simian, et al. "LCM-LoRA: A Universal Stable-Diffusion Acceleration Module." arXiv preprint arXiv:2311.05556 (2023).

[8] Hu, Edward J., et al. "LoRA: Low-Rank Adaptation of Large Language Models." International Conference on Learning Representations. 2021.

Reviewer comment (Theory):
Some concepts need more elaboration. For example, why the adversarial training can lead to better performance than the teacher model? Why Taylor's expansion is the main reason for the discretization error?

Author response:
We think your valuable question directly touches the essential aspects of our contribution. Let us consider the **Property of L2 Distance** and the **Insight from VQGAN**.

- **[Property of L2 Distance]** Of the possible distances between $x_{target}$ and $x_{est}$ in the distillation loss, the Euclidean L2 distance is the most popular one. The L2 distance is decomposed into two components: the low-frequency part and the high-frequency part. The high-frequency part is relatively ignorable compared to the scale of the low-frequency part [9,10]. Therefore, analogous to the autoencoder, if we train the model with the L2 distance in pixel space, the generated samples become blurry as it only captures the low-frequency part in training.
    This imperfect generation is largely mitigated but still remains when we use a perceptual loss, i.e., the feature distance. The feature distance calculates the L2 distance between two features $f_{target}=FE(x_{target})$ and $f_{est}=FE(x_{est})$, where $FE$ is a pre-trained feature extractor. Similar to the pixel distance, the high-frequency part of feature is incorrect after optimizing with the perceptual distillation loss. We tested CTM with different feature extractors [11,12,13], resulting in consistently inferior images across all distances for NFE 1 samples. Notably, the choice of the feature extractor did not yield discernible quality differences in the generated images in both human eyes and evaluation metrics.
- **[Insight from VQGAN]** The adversarial loss can be understood as a distance on a (feature) space that is optimized to maximize the distance between the real data and the fake data. Consequently, if the high-frequency components are not adequately trained, the discriminator captures and optimizes these frequencies, enabling the successful training of the entire frequency spectrum through the adversarial approach. Notably, VQGAN [1] incorporates a combination of adversarial loss and reconstruction loss, yielding two significant advantages: 1) it enhances the stability of adversarial training; 2) it substantially improves the quality of reconstruction. In particular, the reconstruction FID of VQGAN approaches near zero. As the primary modeling goal of distillation models is the accurate reconstruction of teacher samples, we posit that the use of adversarial loss is pivotal in enhancing distillation performance.

[9] Durall, Ricard, Margret Keuper, and Janis Keuper. "Watch your up-convolution: Cnn based generative deep neural networks are failing to reproduce spectral distributions." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2020.

[10] Kim, Nahyun, et al. "Unsupervised Image Denoising with Frequency Domain Knowledge." The 32nd British Machine Vision Conference. British Machine Vision Conference, 2021.

[11] Zhang, Richard, et al. "The unreasonable effectiveness of deep features as a perceptual metric." Proceedings of the IEEE conference on computer vision and pattern recognition. 2018.

[12] Tan, Mingxing, and Quoc Le. "Efficientnet: Rethinking model scaling for convolutional neural networks." International conference on machine learning. PMLR, 2019.

[13] Touvron, Hugo, et al. "Training data-efficient image transformers \& distillation through attention." International conference on machine learning. PMLR, 2021.

Reviewer comment (Experiments):
For a fair comparison with extremely small NFE, the comparison should be with Consistency Models instead of EDM.

Author response:
We appreciate the reviewer's insightful comments. The original submission predominantly compares CTM with EDM and StyleGAN-XL in Section 5.1, with limited emphasis on CM. We concur with the reviewer's suggestion to enhance the discussion on the CTM and CM comparison in small-NFE settings. Accordingly, we revised Section 5.1 to provide a more detailed exploration of the differences between CTM and CM. Additionally, we emphasized the numerical results from the ablation studies in Figures 13-15 of the original submission to offer a comprehensive comparison between CTM and CM in our revision.

Reviewer comment (Novelty):
It would be great if the authors could provide further motivation for using such complicated training objective functions with pretrained model over just training GANs.

Author response:
In addition to the response to **Q1. [Methodology Excessively Complex]** and **Q2. [Why Not Solely Adversarial Loss?]**, we would like to mention that diffusion models are easy for controllable generation due to its explicit form of sample generation process, whereas GAN models should be considered case-by-case for controllable generation. Moreover, functionalities in diffusion models like classifier-free guidance, cost-quality trade-off, and exact likelihood evaluations, are infeasible in GAN models. Nonetheless, we consider the GAN loss pivotal for enhancing sample quality in generative models. Hence, rather than relying solely on adversarial loss, we incorporate GAN loss alongside the reconstruction-based loss, i.e., the distillation loss.

Reviewer comment (Experiments):
It would be helpful if the author could provide the wall-clock of the training procedure. It seems like the training will be even more expensive than the diffusion model or GAN.

Author response:
|                |Epochs (official report)       |GPU Days (calculated in identical machine)                         |
|----------------|-------------------------------|-----------------------------|
| | | |
|EDM|1951            |896            |
|CM          |959            |964            |
|CTM          |48| 112 |

We compare the training time of ImageNet $64\times64$ in the above Table. We calculate the number of epochs based on the official papers and re-calculate the GPU days based on our NVIDIA Ampere machine with 40Gb memory for fair comparison. According to the Table, CTM reaches to the convergence in nearly 1/20 of the iterations required by CM. In terms of wall-clock training time, CTM completes training in approximately 1/10 of the total time needed by CM. We include the above training time in our revision at the first paragraph of Section 5.1.

Reviewer comment (Experiments):
More comparisons with CMs are needed.

Author response:
We appreciate the reviewer's insightful comments. The original submission predominantly compares CTM with EDM and StyleGAN-XL in Section 5.1, with limited emphasis on CM. We concur with the reviewer's suggestion to enhance the discussion on the CTM and CM comparison in small-NFE settings. Accordingly, we revised Section 5.1 to provide a more detailed exploration of the differences between CTM and CM. Additionally, we emphasized the numerical results from the ablation studies in Figures 13-15 of the original submission to offer a comprehensive comparison between CTM and CM in our revision.

Reviewer comment (Reproducibility):
CTM loss brings a lot more additional training cost compared to other distillation methods. It takes at least 4 model forward passes of CTM, multiple steps of the ODE solver at each training step, the cost of computing DSM loss, and also the cost of training the discriminators. Can you provide a detailed description of how many model forward passes are used per training step? Can you also provide the total training time in Table 3 and compare it to the cost of training the teacher model?

Author response:
**Response to Q2.**
|                |Epochs (official report)       |GPU Days (calculated in identical machine)     | #Forward Passes Per Iteration                    |
|-|----------------|-------------------------------|-----------------------------|
| | | | |
|EDM|1951            |896            | 1 |
|CM          |959            |964            | 5 |
|CTM          |48| 112 | 9+n |

We compare the training time of ImageNet $64\times64$ in the above Table. As the reviewer expected, CTM requires more model evaluations per iteration due to the multiple steps of the ODE solver. More precisely, CM computes $5$ forward passes with 1 for teacher (wrapped with `no_grad` in PyTorch), $2$ for student, and the other $2$ for LPIPS. CTM takes 9+n forward passes with n for the number of ODE steps (teacher wrapped with `no_grad` in PyTorch), $2$ for $s\rightarrow 0$ prediction $G_{sg(\theta)}(\cdot,s,0)$, $2$ for student ($G_{\theta}(\cdot,t,s)$ and $G_{sg(\theta)}(\cdot,u,s)$), $2$ for LPIPS, $1$ for DSM loss, and the other $2$ for discriminator. The number of ODE steps $n$ varies each iteration, depending on random $u$. For ImageNet experiment, $n$ is $10$ in average.

We emphasize that CTM achieves significantly faster convergence in terms of wall-clock time compared to EDM and CM, despite requiring more function evaluations. To ensure a fair comparison, we calculate the number of epochs based on the official papers and recalculate the GPU days using our NVIDIA A100 machine with 40GB memory. As indicated in the above table, CTM reaches convergence in nearly $1/10$th of the total training time required by the teacher model (EDM) and the baseline (CM). We included the above argument in the first paragraph of our revised Section 5.1.

Reviewer comment (Experiments):
Training CTM seems to require many tricks like the warm-up scheme for GAN loss and a particular time sampling scheme as described in appendix D.1. But there is little ablation study about them. Can you elaborate more on how you design these training strategies and how important they are?

Author response:
**Response to Q3.**

We appreciate for the constructive comments. Below, we have considered five factors separately: **Warm-up**, **Time Sampling**, **GAN Details**, **EMA**, and **Adaptive Weights**. In our revision, we included below discussions in the revised manuscript more explicitly.

- **[Warm-up]** We apply warm-up for GAN training mainly to stabilize CTM training at the initial phase. The insights of VQGAN is that it is better to incorporate the GAN loss on top of the reconstruction loss once the autoencoder captures the image semantic after sufficient iterations. In our experiments, we found that applying this warm-up to generator is more important than discriminator, and we did not warm-up the discriminator in our training.
- **[Time Sampling]** An well-designed time sampling is critical for score training. When the forward-time SDE is represented by $dx_{t}=\sqrt{2t}dw_{t}$, the large-time score function closely resembles a contraction function towards the origin, exhibiting nearly proportional to $-x_{t}$. Consequently, in EDM, the learning target (oracle score) is predominantly linear for large-time, rendering frequent sampling on large-time unnecessary. In contrast, CTM learns complicated trajectory across the diffusion time $s<t$, demanding that the neural network deviates significantly from linearity when processing large-time inputs. In the process of learning this non-linear complexity for $s<t$, thus, the neural network easily forgets the linearity for $s=t$ on large-time. Without sampling more on large-time for better score matching, the FID significantly deteriorates with large-NFE sampling.
- **[GAN Details]** We mostly follow the practice of StyleGAN-XL in our GAN training. In particular, except the fake data construction, our GAN training mirrors that of StyleGAN-XL, including upscaling to $224\times224$, using cross-channel mixing, and processing with multiple discriminators for different feature maps. While these training details could be substituted with other well-crafted components of state-of-the-art GAN models, we defer further exploration of alternative configurations to future study.
- **[EMA]** When constructing $sg(\theta)$, we adopt the practice from CM by incorporating Exponential Moving Average (EMA) to enhance training stability. While CM used the EMA scale ($\mu$) of 0.0 on CIFAR-10 and 0.95 on ImageNet $64\times64$, we experiment with higher values of $\mu$, specifically 0.999 or 0.9999. This adjustment significantly contributes to stabilizing GAN training, and notably, in our experiments on CTM, we have observed minimal FID fluctuation.
- **[Adaptive Weights]** In addition to warm-up and EMA, we apply the adaptive weighting for $\lambda_{DSM}$ and $\lambda_{GAN}$ to scale each component (the distillation loss, the DSM loss, and the GAN loss) at identical level. The adaptive weighting computes $\lambda_{DSM}=stopgrad(\frac{\Vert \nabla_{\theta_{L}}L_{CTM}(\theta;\phi) \Vert}{\Vert \nabla_{\theta_{L}}L_{DSM}(\theta) \Vert})$ and $\lambda_{GAN}=stopgrad(\frac{\Vert \nabla_{\theta_{L}}L_{CTM}(\theta;\phi) \Vert}{\Vert \nabla_{\theta_{L}}L_{GAN}(\theta,\eta) \Vert})$, where $\theta_{L}$ is the last layer of the UNet output block. This adaptive weight prevents one component to dominate other factors. For instance, without the adaptive weight, after sufficient iterations, the GAN loss suddenly breaks the stability and ruins CTM training.

Reviewer comment (Reproducibility):
Why is the retrained CM much worse than the official report?

Author response:
**Response to Q4.**

The authors of CM provided distinct codebases for CIFAR-10 (implemented in Jax) and ImageNet (implemented in PyTorch). Our retrained CM is grounded in the official PyTorch code (designed for ImageNet), with all hyperparameters calibrated to align with the released Jax code for regenerating CM on CIFAR-10.

While the reasons for the performance disparity between PyTorch and Jax remain unclear, we advocate for a unified codebase to propel the advancement of the distillation community. In alignment with this vision, we have opted to develop our code using PyTorch and will subsequently release a comprehensive, unified codebase that spans across various datasets.

Reviewer comment (Evaluation):
The official report of DFNO is 3.78.

Author response:
**Response to Q5.**

We sincerely appreciate for detailed review and the reviewer's devotion to the community. We modified the performances of DFNO and Rectified Flow in our revision.

Reviewer comment (Presentation):
The FID of Rectified Flow 4.85 is from 2-Rectified Flow with distillation. It could be more concrete about which Rectified Flow model is compared in the table.

Author response:
**Response to Q5.**

We sincerely appreciate for detailed review and the reviewer's devotion to the community. We modified the performances of DFNO and Rectified Flow in our revision.

Reviewer comment (Presentation):
The presentation of the paper could be improved a little. Figure 1 is not very helpful I think -- I would advocate making Figure 2 the main Figure of the paper. There are some typos that could be fixed with more careful reading. For example, there is a parenthesis missing in the citation at the top of page 3.

Author response:
We acknowledge the reviewer's comment regarding the significance of figures in our manuscript. In response to this feedback, we propose relocating Figure 1 to the Appendix and featuring Figure 2 as the primary thumbnail. Also, we have rechecked the manuscript thoroughly and revised typos including the citation of page 3 in our revision. We appreciate for the reviewer of the careful review.

Here, we wish to provide a brief contextual explanation for Figure 1 (of the original manuscript), and (if encouraged) reflect this discussion in the manuscript revision. One of the key motivations driving this research is the creation of a novel family of multi-purpose models that extend beyond traditional score-based models (requiring a minimum of $10$ NFE for sampling) and distillation models (whose performance wasn't consistently enhanced with large NFE). In larger-scale learning scenarios, such as dealing with diverse data modalities, training and saving both score-based and distillation models separately may not be cost-effective. Additionally, practitioners face the dilemma of choosing between a score-based model and a distillation model for their application, unable to fully leverage the merits of both. However, CTM, accommodating training from scratch, distillation, and score-based approaches, offers a unified and versatile model. This single model seamlessly integrates all the advantages of separate models, catering to diverse applications and providing a more economical and practical solution.

Reviewer comment (Theory):
In the definition of the $g$ function (Lemma 1), I think the integral bounds $s, t$ have to change order (or a minus sign needs to be added in front of the integral).

Author response:
We thank the reviewer for the detailed feedback. In response to the reviewer's comment, we carefully recheck our derivations with particular emphasis on the time direction and integration order. Based on our careful examination, to the best of our knowledge, we confirm that there are no mathematical flaws in our original manuscript. For the benefit of clarity and transparency, we present a comprehensive account of the derivation below.

The reverse-time SDE is 
$$dx_{T-\tau}=g^{2}(T-\tau)\nabla\log{p_{T-\tau}(x_{T-\tau})}d t+g(T-\tau)dw_{\tau},
$$ where the forward-time SDE is given by $$dx_{\tau}=g(\tau)dw_{\tau}.$$
The reverse-time PF-ODE thus becomes $$dx_{T-\tau}=\frac{1}{2}g^{2}(T-\tau)\nabla\log{p_{T-\tau}(x_{T-\tau})}d \tau.$$
Therefore, by integrating from $T-t$ to $T-s$ ($s<t$), we obtain $$x_{s}-x_{t}=\int_{T-t}^{T-s}\frac{1}{2}g^{2}(T-\tau)\nabla\log{p_{T-\tau}(x_{T-\tau})}d \tau.$$
By change-of-variable with $u=T-\tau$, the equation is derived to be $$x_{s}=x_{t}-\int_{t}^{s}\frac{1}{2}g^{2}(u)\nabla\log{p_{u}(x_{u})}d u.$$
With $g(u)=\sqrt{2u}$ and the Tweedie's formula $\nabla\log{p_{u}(x_{u})}=\frac{E[x\vert x_{u}]-x_{u}}{u^{2}}$, we derive Eq. (2) in our paper: $$x_{s}=G(x_{t},t,s)=x_{t}+\int_{t}^{s}\frac{x_{u}-E[x\vert x_{u}]}{u}d u.$$
Now, we derive the following equations: 
$$
G(x_{t},t,s)=\frac{s}{t}x_t+ \Big(1-\frac{s}{t}\Big) x_t+ \int_{t}^{s} \frac{x_{u}-E[x\vert x_{u}]}{u} d u 
=\frac{s}{t}x_t+ \Big(1-\frac{s}{t}\Big) \bigg[ x_t + \frac{1}{(1-\frac{s}{t})} \int_{t}^{s} \frac{x_{u}-E[x\vert x_{u}]}{u} d u \bigg] $$
$$=\frac{s}{t}x_{t}+\Big(1-\frac{s}{t}\Big) \bigg[x_{t}+\frac{t}{t-s}\int_{t}^{s}\frac{x_{u}-E[x\vert x_{u}]}{u}d u \bigg]
=\frac{s}{t}x_{t}+\Big(1-\frac{s}{t}\Big)g(x_{t},t,s),
$$ 
where $g(x_{t},t,s):=x_{t}+\frac{t}{t-s}\int_{t}^{s}\frac{x_{u}-E[x\vert x_{u}]}{u}d u$.

Reviewer comment (Experiments):
Following this point, I think being able to approximate the conditional expectation with the same network is not really important. For this method to work, we need to start from a pre-trained score model anyway. Additionally, the proposed method for learning this doesn't actually work because of numerical instabilities, and hence the authors resort to adding an additional loss term.

Author response:
Thank you for the insightful review. We agree with the reviewer's observation that once an any-to-any jump $G(\mathbf{x}_{t},t,s)$ is approximated, the necessity for infinitesimal jumps is eliminated. Below, given that we agree with the reviewer, we would like to provide the detailed motivation of our model design and discuss that the infinitesimal jump is gained as a side effect of our model design. We have revised Section 3.1 in line of below response of reviewer's question.


- **[Constrained to Unconstrained]** Directly approximating $G$ using a neural network is challenging due to the constrained optimization problem posed by the initial condition  $G_{\theta}(x_{t},t,t)=x_{t}$: $$\min_{\theta}L(\theta)\quad\text{subject to}\quad G_{\theta}(x_{t},t,t)=x_{t}.$$ By setting $c_{skip}(t,s)=\frac{s}{t}$ in $G_{\theta}(x_{t},t,s)=c_{skip}(t,s)x_{t}+c_{out}(t,s)NN_{\theta}(x_{t},t,s)$, the initial condition $G_{\theta}(x_{t},t,t)=x_{t}$ is satisfied for free, transforming the constrained problem into an unconstrained optimization task $\min_{\theta}L(\theta)$. 
- **[$c_{out}$ for Stable Training]** In our design, $G_{\theta}=c_{skip}(t,s)x_{t}+c_{out}(t,s)NN_{\theta}$, $c_{out}(t,s)$ acts as a time-weighting factor for $NN_{\theta}$ and influences the loss gradient. We experimented with $c_{out}(t,s)\equiv 1$, indicating equal contribution of various time, but this resulted in extremely unstable training. Therefore, inspired by the Euler solver (Appendix C.1), we adopted $c_{out}(t,s)=1-\frac{s}{t}$ as a mathematically oriented weighting function, which contributed to stabilizing the training. 
- **[Denoiser As Bonus]** The conditional expectation (i.e., denoiser function), as a natural consequence of our design, is now approximated by $NN_{\theta}$ as $s\rightarrow t$. This inherent capability allows CTM to naturally generalize to both diffusion models ($s=t$) and distillation models ($s<t$). After examining various candidates for $c_{skip}$ and $c_{out}$, we found that the Euler-solver motivated design of $c_{skip}(t,s)=\frac{s}{t}$ and $c_{out}(t,s)=1-\frac{s}{t}$ demonstrated to be the most stable, to the best of our knowledge. Please refer to Appendix C.1 for a detailed analysis in this regard.

Reviewer comment (Writing):
Section 3.2 is a little bit confusing. The notation added is too much and makes it hard to follow what's going on. I think (?) that the main point is that another network is used for the ODE solving and the student network learns these predictions. If that's the case, this needs to be a little bit more clear.

Author response:
We agree that the notation could be difficult to understand. Hence, in the revision of Section 3.2, we aim to streamline equations, focusing on providing more comprehensive yet simplified explanations of formulas and notations. We appreciate the reviewer to point out this.

Reviewer comment (Evaluation):
Why does this method give better results than Consistency Models when s=0?

Author response:
The reviewer's question pertains to understanding why CTM achieves a more accurate estimation of $G(\mathbf{x}_{t},t,0)$ compared to CM. Building upon the insights provided in **Q9. [Why Student Beats Teacher?]** that discusses supplementary components such as DSM/GAN losses, this thread further emphasizes the differences in distillation losses between CTM and CM. We also reflect the below arguments in the first paragraph of Section 5.2 and the entire new section of Appendix C.3.

Our focus centers on 1-step generation $G_{\theta}(x_{T},T,0)$, where we examine the distillation mechanisms employed in CM (using *local consistency*) and CTM (using *soft consistency*), respectively. To align with the implementation, we adopt a discrete time setup.

- **[Local Consistency (Implicit Information from Teacher)]**  Let us assume that at some training iteration the maximum time $T$ is sampled as a random time $t$. Then CM matches the long jump provided by $G_{\theta}(x_{T},T,0)$ and $G_{sg(\theta)}(x_{T \rightarrow T-\Delta t}^{\phi},T-\Delta t,0\big)$. Here, $x_{T \rightarrow T-\Delta t}^{\phi}$ denotes a 1-step denoised sample of the *pre-trained teacher* (with parameters $\phi$), obtained using an ODE solver, starting from $x_T$ at time $T$ and terminating at $T-\Delta t$. Hence, the student model $G_{\theta}(x_{T},T,0)$ distills only the teacher information within the interval $[T-\Delta t,T]$ and may lack precision for the trajectory within $[0,T-\Delta t]$. The transfer of teacher information for the interval $[0,T-\Delta t]$ may occur in another iteration with a random time $t\le T-\Delta t$. In this case, the student model $G_{\theta}(x_{t},t,0)$ distills teacher information solely within the interval $[t-\Delta t,t]$, where the network is trained with $x_{t}$ as the input. 

    However, for 1-step generation, $G_{\theta}(x_{T},T,0)$ still lacks perfect knowledge of the teacher's information within $[t-\Delta t,t]$. This is because, when the student network input is $x_{T}$, the teacher's information for the interval $[t-\Delta t,t]$ with $t\le T-\Delta t$ has not been explicitly provided, as the student was trained with the input $x_t$ to distill information within $[t-\Delta t,t]$. Given the non-overlapping intervals with distilled information from local consistency, the student neural network must extrapolate and attempt to connect the scattered teacher information. Consequently, this implicit signal provided by teacher results in slow convergence and inferior performance. We reference Figure 13 in the manuscript, wherein it illustrates the inferior performance of local consistency (blue curve).
- **[Soft Consistency (Explicit Information from Teacher)]** At the opposite end of local consistency is *global consistency*, where the teacher prediction is constructed using an ODE solver to cover the entire interval $[0, T]$ (or $[0, t]$ for a random time $t$). In this case, the student model can explicitly extract information from the teacher. However, this approach is resource-intensive (3$\times$ slower than local consistency on CIFAR-10) during training due to the ODE solving calls on the entire interval at each iteration.

    In contrast, CTM constructs the teacher prediction with our innovative technique, *soft matching* through the use of an ODE solver, spanning from $T$ to a random $u$. Importantly, $u$ is not limited to $T-\Delta t$ but can take any value in the range of $[0, T]$. As a result, the teacher information has the opportunity to be distilled and transmitted over a broader range of $[u, T]$. More precisely, if a random $u$ is sampled with $u\le t-\Delta t$, the range $[u,T]$ containes the interval $[t-\Delta t,t]$, and the student network directly distills teacher information of the interval $[t-\Delta t,t]$. As $u$ is arbitrary, the student of CTM with input $\mathbf{x}_{T}$ will ultimately receive explicit information from the teacher for any intermediate timesteps. This renders CTM superior to CM for 1-step generation, as evidenced in Figure 13 (green curve), while maintaining training efficiency (2$\times$ faster than global consistency on CIFAR-10).