Paper: Alleviating Exposure Bias in Diffusion Models through Sampling with Shifted Time Steps (ICLR-2024)

Reviewer comment (Theory):
Due to the conditions imposed by the theoretical derivation, almost 10% sampling time can not be simply ignored.

Author response:
Regarding the additional +9.5% sampling time when sampling with DDPM on CIFAR-10:

1. In practice, we believe it is acceptable, for example, when we are sampling with 20 steps, it only requires less than 10ms additional sampling time.

2. We would not consider the conditions imposed by the theoretical derivation as the cost. 
In fact, our method requires nearly no additional sampling time when working with a bigger backbone ADM[1] as shown in Figure 6 and Table 2.  It makes our method stand out in nearly a “free lunch” situation, where we make significant improvements in performance without sacrificing efficiency.
The additional computation we introduce to the sampling method only involves (1) computation of the sample variance (2) calculating $\mathop{\arg\min}(\cdot)$ within a small window size of timesteps. These two computations can be quickly achieved in modern GPUs.

Reviewer comment (Experiments):
As mentioned 'seamlessly integrated to existing sampling algorithms', how about the performance combined with the DPM-solver [1], DEIS [2].

Author response:
We report the results of our method combined with DPM-solver-2 and DEIS-order-2 on CIFAR-10 in the table below. Limited to resources and rebuttal time, we report the results using 5, 10 and 20 sampling steps. (The results can also be found in Table 9 in Appendix I of the revised manuscript.) Our model can still improve the performance of both DPM-solver and DEIS samplers. However, in comparison to DDIM, our method again provides substantial improvements although the improvements for these samplers are a bit smaller than for the other tested samplers.This could be attributed to the fact that both the DPM-solver and DEIS utilize the particular structure of the semi-linear ODE, which already largely reduces the error of estimating $x_t$. 

 | Sampling Method | 5 steps | 10 steps | 20 steps|
|---|---|---|---|
|DPM-solver-2 | 32.30 | 10.92 | 4.30 |
|*TS*-DPM-solver-2| **31.02 (+3.96%)** | **9.82 (+10.07%)** | **4.11 (+4.42%)** | 
|DEIS-order-2-*t*AB | 24.64 | 5.88 | 4.13 | 
|*TS*-DEIS-order-2-*t*AB | **22.57 (+8.40%)** | **5.41 (+8.00%)** | **3.62 (+12.30%)** |

Reviewer comment (Theory):
In my humble opinion, the analysis of the exposure bias problem is empirically not theoretically as mentioned in the 'contribution'.

Author response:
Thank you for acknowledging our sound empirical findings on the exposure bias problem. Regarding the theoretical aspect of the problem, we would like to point out that: 

1. We derive the analytic form of the optimal variance $\theta_s$ for the optimal time step $t_s$ corresponding to any $t-1$ in Theorem 3.1. 

2. We show that the cutoff mechanism is a necessity to reach the derivation of Equation 21 (in Appendix J). It provides the theoretical foundation for the cutoff values. 

3. The purpose of the selection window of timesteps is already well embedded in our theorem. 

4. We further add the derivation of the analytic estimation for window size in Appendix J.2 as you suggested.
We believe our theoretical contribution to the problem is just as important as our empirical findings.

Reviewer comment (Novelty):
There is no obvious advantage over the recent training-free sampling method.

Author response:
Limited to resources, we cannot test all the existing sampling methods. In Table 1, we already show the significant improvements obtained when applying our method combined with S-PNDM and F-PNDM [1], arguably two of the most performant sampling methods at this moment. Also as requested, we conduct experiments with  DPM-solver [2] and DEIS [3] as shown below (Q1). These experiments show that our method could be  seamlessly integrated to existing sampling algorithms. Thus, it is suitable to see our method as a way to further augment other sampling methods. 

Furthermore, our contributions go beyond just a sampling method. We show that the exposure bias problem of diffusion models could be alleviated without retraining the model, which could inspire future research on designing new samplers. We consider the theoretical and empirical findings regarding the exposure bias problem that we present in the paper  to be pioneering to the field.


[1] Liu et al. Pseudo numerical methods for diffusion models on manifolds. ICLR22 

[2] Lu et al. DPM-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. NeurIPS22

[3] Zhang and Chen. Fast sampling of diffusion models with exponential integrator. ICLR23

Reviewer comment (Experiments):
Can this method combine with other training-free sampling methods, such as DPM-sovler, DEIS?

Author response:
We report the results of our method combined with DPM-solver-2 and DEIS-order-2 on CIFAR-10 in the table below. Limited to resources and rebuttal time, we report the results using 5, 10 and 20 sampling steps. (The results can also be found in Table 9 in Appendix I of the revised manuscript.) Our model can still improve the performance of both DPM-solver and DEIS samplers. However, in comparison to DDIM, our method again provides substantial improvements although the improvements for these samplers are a bit smaller than for the other tested samplers.This could be attributed to the fact that both the DPM-solver and DEIS utilize the particular structure of the semi-linear ODE, which already largely reduces the error of estimating $x_t$. 

 | Sampling Method | 5 steps | 10 steps | 20 steps|
|---|---|---|---|
|DPM-solver-2 | 32.30 | 10.92 | 4.30 |
|*TS*-DPM-solver-2| **31.02 (+3.96%)** | **9.82 (+10.07%)** | **4.11 (+4.42%)** | 
|DEIS-order-2-*t*AB | 24.64 | 5.88 | 4.13 | 
|*TS*-DEIS-order-2-*t*AB | **22.57 (+8.40%)** | **5.41 (+8.00%)** | **3.62 (+12.30%)** |

Reviewer comment (Theory):
Is there exits a analytical solution for the window size?

Author response:
Thank you for the question. We provide the derivation for the analytical estimation of the window size in Appendix J.2.

Reviewer comment (Experiments):
The selection of window sizes and cutoff values is primarily based on limited experience.

Author response:
The characteristics of our empirical findings indicate that the selection procedure of window sizes and cutoff values are  robust and can be extended to different settings.
The empirical findings are (1) very intuitive. For example, more sampling steps require smaller window sizes. It is well motivated by the fact the violation of Gaussian assumption is less a problem for more sampling steps. And the cutoff value is closely related to the changes of sample variance along different timesteps as presented in Figure 2. (2) The empirical findings are showing very similar patterns across datasets, models and sampling methods. As also reflected by the results obtained in various settings, given the same number of sampling steps, we find that the optimal values for window sizes and cutoff values always fall into the same range. For a more analytic solution of window size, please refer to the answer to the Q3 of reviewer zG28 or the Appendix J.2.

Reviewer comment (Experiments):
How can we design a more effective strategy for selecting window sizes and cutoff values when dealing with random datasets and image sizes?

Author response:
It is intriguing to study the design of a more systematic and effective selection strategy for window sizes and cutoff values. As stated before, we find the optimal values for window sizes and cutoff values show remarkably similar patterns across datasets and models. 
Specifically, as stated in Section 5.4, (1) one should opt for large window sizes (e.g. 40, 80) for fewer steps(e.g. 5,10 steps) and very small window sizes (e.g. 2) for more steps (e.g.50,100 steps). This holds true for all our experiments. (2) The optimal cutoff values are always within the range of [200,400]. Under different experimental settings, only less than 5% of the cases in our experiments the optimal cutoff value falls outside the range of [200, 400]. For a more analytic solution of window size, please refer to the answer to the Q3 of reviewer zG28 or the Appendix J.2.

Reviewer comment (Writing):
The writing needs improvement as it is currently a bit ambiguous at certain places. (See 1. and 2. in questions below for further details). For instance, plotting details for Figure 3 are unclear from the appendix. Similarly, it is unclear to me how $var(x_{t-1})$ is computed in line 9 of Algorithm 3. The overall clarity of this paper will greatly improve if these paragraphs are rewritten by adding additional details. Similarly, mathematical expressions should be added at multiple places along with text for improved clarity. For instance, in Figure 2, writing $var(x_t)$ is more informative than simply writing $x_t$ for the label of y-axis. In Figure 3, the mathematical form of error can be included instead of labelling y-axis as error.

Author response:
> **[W2] suggestions on added details in paper**

Thank you for the suggestions. We’ve updated the manuscript as follows: (1) changed y-axis to $var(x_t)$ in Figure 2; (2) added the details for plotting Figure 3 in Appendix B, including a visualized example presented in Figure 8; (3) added details in computing $var(x_{t-1})$ of Algorithm 3 in Appendix B. Please also refer to the comments we made below for Q1 & Q2.

Reviewer comment (Theory):
Certain parts of derivation of proof of Theorem 3.1 need further explanation. The proof assumes that for an image P, and pixels $p_i, p_j \in P$, $p_i \perp p_j$ if $i \neq j$ which is usually not true in practice as neighboring pixels in image usually have high correlation. It also assumes that each pixel in image $P$ follows distribution $\mathcal{N}(\mu_i, \sigma)$, but later it claims that $\mu_i$ is the mean of the distribution of $\hat{x}_{t-1}$ which looks incorrect as the mean of the latter distribution is a vector/tensor (as it is an image) while $\mu_i$ is a scalar as it is mean of pixel values.

Author response:
> **[W3] explanation on independence assumption in derivation**

Thanks for raising this insightful question. Indeed, the assumption of independence is not true for natural  images.  However, considering the image generation process of the diffusion model, we start with Gaussian distribution and gradually remove the noise to generate images. Initially, the assumption of independence is valid, since we start with pure Gaussian noise. Yet, as time approaches zero, the pixels are more and more dependent on each other, making the assumption of independence invalid. Thus our derivation is applicable when $t$ is large. This also explains the introduction of  the cutoff mechanism, where we avoid applying the time shift when $t$ is smaller than a certain predefined value. 

We could provide a more general derivation that takes into account the correlation between pixels; however, we would still assume that this correlation can be ignored when $t$ is large.

To construct a more general derivation, one only needs to make the following modifications to our original derivation. The independence condition on $p_i$ and $p_j$ is only used in Equation 16. When assuming a general Gaussian distribution $\mathcal{N}(\mu, \Sigma)$ on the image, we need to replace the $\mathbb{E}(p_i)\mathbb{E}(p_j)$ term by  $\mathbb{E}(p_i p_j)$. Equivalently we can keep the first equation unchanged while adding an extra $\sum_{i \neq j}\mathbb{E}(p_i p_j) - \mathbb{E}(p_i)\mathbb{E}(p_j)$, which gives the covariance between $p_i$ and $p_j$, denoted by $\Sigma_{\{i,j\}}$ Therefore, Equations 16 to 19 have an additional $\sum_{i \neq j}\frac{\Sigma_{\{i,j\}}}{d(d-1)}$ at the end. Equation 20 has an additional $\sum_{i \neq j}\frac{\Sigma_{\{i,j\}}}{d-1}$ at the end and Equation 21 has an additional $-\sum_{i \neq j}\frac{\Sigma_{\{i,j\}}}{d-1}$  on the right hand side.

Thanks for bringing up the ambiguity of $\mu$.  $\mu_i$ is the mean of each pixel of the $x_{t-1}$. Thus it is also a scalar.  We have corrected this in the revision.

Reviewer comment (Writing):
The explanation of details for plotting Figure 3 is unclear from the description given in Appendix B. Perhaps writing mathematical equations might make the idea more concise and clear. Also, for the purpose of rebuttal, could the authors include this expression here, or alternately provide an explanation of what the figure indicates? It is unclear to me why the error computation for this figure is split into two different stages. Any intuition/reasoning behind choosing the methodology of computing errors in Figure 3 is appreciated.

Author response:
> **[Q1] details and intuition for plotting Figure 3**

Thanks for raising this important question. Figure 3 shows the mean square error (MSE) between the prediction and ground truth at each step in the backward process. Given an image denoted as $x_0$, by applying Equation 2 ($x_t=\sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon_t$)  we could obtain a sequence of $x_t, t=1,2,\cdots, T-1$.  Taking each $x_t$ and the paired time step $t$ to run the backward process, we obtain a sequence of predicted $\hat{x}_0^t, t=1,2,\cdots, T-1$. Ideally, we would expect all these predicted $\hat{x}_0^{t}, t=1,2,\cdots, T-1$ to be exactly equal to the ground truth $x_0$, as they are generated using the given $x_0$. However, this is not the case in practice. In our experiments, we found that only when $t<t_s$ (around 650 steps in our experiment using DDIM) we could obtain the original $x_0$ by running the backward process with paired $(x_t,t)$. For $t>t_s$, the image created using $(x_t, t)$ differs from the original $x_0$. This observation also reveals that the image generation process of diffusion models basically contains two stages. In the first stage, the model moves the Gaussian distribution towards the image distribution and no modes are presented at this stage, which means we can not know which images will be generated. In the second stage, the prediction shows clear patterns and modes are presented. We can predict which images will be generated following the backward process. This observation led us to divide the error computation into two stages. The full  explanation, including the equations, are added to Appendix B with Figure 8 showing the visualized procedure.

Reviewer comment (Evaluation):
Other metric should also be included, such as precision and recall. Is there any reason why text-to-image performances are not added? It would be better if authors could include such results as the task is one of most important tasks of diffusion model.

Author response:
We followed previous works [1,2] by reporting results in FID on the widely used benchmark datasets such as CIFAR-10, CelebA and LSUN. But, yes, reporting results on other metrics and the results in text-to-image generation tasks would definitely make our contribution stronger. Thus, we report these results as follow:

| Sampling Method | 10 steps | 20 steps|
|---|---|---|
|DDIM | 27.80 | 25.47 |
|*TS*-DDIM| **26.32 (+5.32%)**|**24.80 (+2.63%)**|

The results are also presented in Table 10 in Appendix I of the revised manuscript.

We report precision (P) and recall (R) with ADM as backbone on CIFAR10 using DDIM and our *TS*-DDIM. The results are also included in Table 7 in Appendix I of the revised manuscript.

| Sampling Method | 5 steps, P | 5 steps, R | 10 steps, P | 10 steps, R| 20 steps, P | 20 steps, R | 50 steps, P | 50 steps, R |
|---|---|---|---|---|---|---|---|---|
|DDIM | 0.59 | 0.47 | 0.62 | 0.52 | 0.64 | 0.57 |0.66 | 0.60 |
|*TS*-DDIM| 0.57 | 0.46 | 0.62 | 0.55 | 0.64 | 0.60 | 0.65 | 0.62 |

[1] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021.

[2] Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. In International Conference on Learning Representations, 2022.

Reviewer comment (Evaluation):
Could you also visualize the selecte timestep trajectory? It also would better to have more analysis on the selected timesteps and around which timesteps are mostly important.

Author response:
We visualize the selected timestep trajectory for experiments on CIFAR-10 with *TS*-DDIM using DDPM backbone and 10 sampling steps. Please refer to Figure 12 in Appendix G of the revised manuscript. The red dashed line is the original timesteps in DDIM (uniform selection). The blue line shows the trajectories using *TS*-DDIM. For better visibility, we also zoom in for the range of [300,600]. In the range of (300, 600), more time shifts happen, and time steps can shift to both larger or smaller time steps within the window. Figure 12 demonstrates that most of the time shifting happens in the intermediate range of the backward process. This phenomenon might be due to the fact that at the initial stage (when $t$ is large), the samples are predominantly influenced by the initialized normal distribution, leading to minimal variance change. Conversely, in the later stage (when $t$ is small), the samples are predominantly composed of image data, containing most of the image information. This makes it easier for the model to make accurate predictions.