Paper: On Error Propagation of Diffusion Models (ICLR-2024)

Reviewer comment (Theory):
Modular Error Definition: The modular error is defined as the expected KL divergence from $p_{\theta}$ to $q$. However, the KL divergence is not symmetric. Could the authors please explain why the modular error is defined in such a way? Why not using the reversed KL divergence as in the original loss function?

Author response:
Thanks for raising this insightful point. We did consider the asymmetry of KL divergence. The key is that **the error propagation happens to the learnable backward process $p_{\theta}$** (instead of the predefined forward process $q$). Therefore, it makes more sense to define the module error as

$E_{x_t \sim p_{\theta}(x_t)}[ D_{KL} ( p_{\theta}(x_{t-1} | x_t) \mid\mid q(x_{t-1} | x_t) ] = \int p_{\theta}(x_t) ( \int p_{\theta}(x_{t-1} | x_t) \ln (\cdot) dx_{t-1}) dx_t = E_{x_t \sim p_{\theta}(x_t)}[ E_{x_t \sim p_{\theta}( x_{t-1} | x_t)}[ \ln(\cdot) ]  ] $,

such that **expectation integrals are operated on the distributions $p_{\theta}(x_t), p_{\theta}(x_{t-1} | x_t)$ of the backward process** to average the distribution gap $\ln(\cdot)$. For your mentioned reversed case, the expectation operations will be mistakenly applied to the distributions $q(x_t), q(x_{t-1} | x_t)$ of the forward process.

Reviewer comment (Theory):
In the proof of Proposition A.1 in Appendix A, equation (19) holds only in the limit as $T\to\infty$. The current proof falls short of demonstrating Proposition A.1 for finite values of $T$.

Author response:
Thanks for pointing out this. We address your concern by **introducing a weaker precondition, which still makes our proposition hold and is also supported by current works [1,2]**.

In our proof to Proposition A.1, the use of that precondition $T \rightarrow \infty$ is to let $q(x_T | x_0) \rightarrow N(x_T; 0, I)$. Considering the fact that 

$q(x_t | x_0) = N(x_t; \sqrt{\bar{\alpha}_t} x_0, \sqrt{1 - \bar{\alpha}_t} \epsilon)$, 

we can achieve the goal with a much weaker assumption: $\lim_{t \rightarrow T} \bar{\alpha_t} \rightarrow 0$. Notably, **this new assumption is a standard configuration in current diffusion models** (e.g., DDPM [1] and SGM [2]), which lets $x_T$ contain no information about $x_0$.

Reviewer comment (Theory):
Several typographical errors in equation (24) in Appendix B raise concerns about the validity of the proof. Specifically, the $p_{\theta}$ term in both the numerator and denominator should be consistent and denoted as $p_{\theta}(x_{t-1}|x_{t})$ in order for equation (29) to hold. Additionally, the first term in the second line of (24) appears to be missing a logarithmic notation, and the third line of (24) appears inconsistent with the current notation.

Author response:
Thanks for pointing out this. The correct forms of the first two lines of our Eq. (24) are

$ E_{x_t \sim p_{\theta}(x_{t})} [E_{x_{t-1} \sim  p_{\theta}(x_{t-1} \mid x_{t}) } [\ln ( \frac{p_{\theta}(x_{t-1} \mid x_{t})} {q(x_{t-1} \mid x_{t})} \frac{q(x_{t} \mid x_{t-1})} {p_{\theta}(x_{t-1} \mid x_{t})} ) ] ] = E_{x_t \sim p_{\theta}(x_{t})} [E_{x_{t-1} \sim  p_{\theta}(x_{t-1} \mid x_{t}) } [\ln  \frac{p_{\theta}(x_{t-1} \mid x_{t})} {q(x_{t-1} \mid x_{t})}  ] ] + E_{x_t } [E_{x_{t-1}} [ \ln \frac{q(x_{t} \mid x_{t-1})} {p_{\theta}(x_{t-1} \mid x_{t})} ] ]$.

We also fix the last line of Eq. (24) as

$E_{x_t \sim p_{\theta}(x_{t})} [ D_{KL} (p_{\theta}(x_{t-1} \mid x_{t}) || q(x_{t-1} \mid x_{t})) ] + E_{x_t \sim p_{\theta}(x_{t})} [E_{x_{t-1} \sim  p_{\theta}(x_{t-1} \mid x_{t}) } [ \ln \frac{q(x_{t} \mid x_{t-1})} {p_{\theta}(x_{t-1} \mid x_{t})} ] ]$.

We will add the above corrections to the revised version.

Reviewer comment (Experiments):
Hyper-Parameter Selection: The authors should provide further clarification on the process for selecting hyper-parameters. Notably, the weight assigned to the regularization term is a critical factor that requires elucidation. How to select the hyper-parameters? One such important factor is the weight of the regularization term, could the authors please clarify?

Author response:
We applied the grid search to hyper-parameter selection. For the weights of regularization terms, we respectively construct candidate sets $[ 0.2, 0.4, 0.6, 0.8 ]$ and $[ 1 \times 10^{-3}, 3 \times 10^{-3}, 6 \times 10^{-3}, 9 \times 10^{-3} ]$ for $\lambda^{reg}$ and $\rho$. As mentioned in Sec. 6.1 of our paper, the experiments turned out that this combination $\lambda^{reg} = 0.2, \rho = 3 * 10^{-3}$ performed the best.

Reviewer comment (Writing):
The notation for conditional expectation should be corrected throughout the paper to $E_{X_{t}|X_{t+1}}$.

Author response:
Thanks for raising these points. We will accept your suggestions about the notation and the typo in the revised version.

For your comment “x < 4”, we prove that $D_{t}^{cumu} < 4$ such that $x = D_{t}^{cumu} < 4$. Recall that the definition of $D_{t}^{cumu}$ is as

$|E_{p_{\theta}(x_{t-1})}[\phi(x_{t-1})] - E_{q(x_{t-1})}[\phi(x_{t-1})]|^2$,

where $|\phi| = \sup_{x} |\phi(x)| < 1$. According to the Triangle Inequality, we have

$D_{t}^{cumu} \le (|E_{p_{\theta}(x_{t-1})}[\phi(x_{t-1})]|+ |E_{q(x_{t-1})}[\phi(x_{t-1})]|)^2$,

Since $|x|$ is a convex function, we can apply Jensen's inequality to the above equation:

$D_{t}^{cumu} \le (E_{p_{\theta}(x_{t-1})}[|\phi(x_{t-1})|]+ E_{q(x_{t-1})}[|\phi(x_{t-1})|])^2 < (E_{p_{\theta}(x_{t-1})}[1]+ E_{q(x_{t-1})}[1])^2 = (1 + 1)^2 = 4$,

showing that our claim holds.

Besides the above proof, the experiments in Fig. 2 of our paper also empirically confirmed that $D_{t}^{cumu} < 4$.

Reviewer comment (Writing):
In Section 3.2, second paragraph, "get ride of" should be corrected to "get rid of".

Author response:
Thanks for raising these points. We will accept your suggestions about the notation and the typo in the revised version.

For your comment “x < 4”, we prove that $D_{t}^{cumu} < 4$ such that $x = D_{t}^{cumu} < 4$. Recall that the definition of $D_{t}^{cumu}$ is as

$|E_{p_{\theta}(x_{t-1})}[\phi(x_{t-1})] - E_{q(x_{t-1})}[\phi(x_{t-1})]|^2$,

where $|\phi| = \sup_{x} |\phi(x)| < 1$. According to the Triangle Inequality, we have

$D_{t}^{cumu} \le (|E_{p_{\theta}(x_{t-1})}[\phi(x_{t-1})]|+ |E_{q(x_{t-1})}[\phi(x_{t-1})]|)^2$,

Since $|x|$ is a convex function, we can apply Jensen's inequality to the above equation:

$D_{t}^{cumu} \le (E_{p_{\theta}(x_{t-1})}[|\phi(x_{t-1})|]+ E_{q(x_{t-1})}[|\phi(x_{t-1})|])^2 < (E_{p_{\theta}(x_{t-1})}[1]+ E_{q(x_{t-1})}[1])^2 = (1 + 1)^2 = 4$,

showing that our claim holds.

Besides the above proof, the experiments in Fig. 2 of our paper also empirically confirmed that $D_{t}^{cumu} < 4$.

Reviewer comment (Evaluation):
The proposed method requires significant compute overhead, so gains need to be weighed against this increase in compute.

Author response:
Thanks for raising this point. A reminder is that **the y-axis in Fig. 4 of our paper actually represents the time cost of $200$ training steps**, rather than just $1$ step. We apologize if this mistake has misled you to the impression that our method is very inefficient.

As you suggested, **we have compared our model with the two most efficient baselines** (i.e., DDPM and DDIM) on CIFAR-10 with the same training time. For Consistent DM and FP-Diffusion, our model performs better than them and has a lower time cost per training step. The experiment results of FID scores and time costs are as follows.

| Model                  | Training Steps | Total Training Time | FID Score |
|------------------------|----------------|---------------------|-----------|
| DDPM                   | 200K           | 10hrs               | 3.61      |
| ADM-IP                 | 190K           | 10hrs               | 3.29      |
| **DDPM w/ Our Method** | **98K**            | **10hrs**               | **3.16**      |
| DDPM w/ Our Method     | 200K           | 20hrs               | 2.93      |

From the above table, we can see that **our proposed method still outperforms the baselines with the same amount of training time**. For example, our method reduces the FID score of DDPM on CIFAR-10 by 12.47%, using no more than half of its training steps. Besides, while leading to a higher time cost per training step, **our method has no impact on the inference speed of diffusion models, which is more important in practical applications**.

Reviewer comment (Evaluation):
The proposed method has significant computational overhead (Figure 4). How does the comparison to baselines (Table 1) change at equal wall-clock time?

Author response:
Thanks for raising this point. A reminder is that **the y-axis in Fig. 4 of our paper actually represents the time cost of $200$ training steps**, rather than just $1$ step. We apologize if this mistake has misled you to the impression that our method is very inefficient.

As you suggested, **we have compared our model with the two most efficient baselines** (i.e., DDPM and DDIM) on CIFAR-10 with the same training time. For Consistent DM and FP-Diffusion, our model performs better than them and has a lower time cost per training step. The experiment results of FID scores and time costs are as follows.

| Model                  | Training Steps | Total Training Time | FID Score |
|------------------------|----------------|---------------------|-----------|
| DDPM                   | 200K           | 10hrs               | 3.61      |
| ADM-IP                 | 190K           | 10hrs               | 3.29      |
| **DDPM w/ Our Method** | **98K**            | **10hrs**               | **3.16**      |
| DDPM w/ Our Method     | 200K           | 20hrs               | 2.93      |

From the above table, we can see that **our proposed method still outperforms the baselines with the same amount of training time**. For example, our method reduces the FID score of DDPM on CIFAR-10 by 12.47%, using no more than half of its training steps. Besides, while leading to a higher time cost per training step, **our method has no impact on the inference speed of diffusion models, which is more important in practical applications**.

Reviewer comment (Theory):
The assumption in the core theorem is absolutely wrong. "suppose that the output of neural network $\epsilon_\theta$ follows a standard Gaussian", which cannot be true. Because the noise-pred model corresponds to the denoising score matching loss, it is proved that the ground truth of such model is propotional to the score function of the distribution, i.e., $\nabla_{x_t} \log q_t(x_t)$. For a small $t$, such score function is quite complex and cannot be a simple and single-mode Gaussian distribution, and is far different.

Author response:
Thanks for raising this point. In this answer, **we first explain the role of your mentioned assumption in Theorem 3.1 and then replace it with a much weaker assumption**, which is derived from a theoretical perspective and still makes the theorem hold. **Lastly, we perform experiments to show that our new assumption is indeed valid in practice**.

Since neural network $\epsilon_{\theta}(x_t, t)$ is tasked to fit Gaussian noise $\epsilon$ in the loss function $L^{nll}$, we previously assumed that the output distribution of $\epsilon_{\theta}(x_t, t)$ follows a standard Gaussian. In our proof to Theorem 3.1, **that assumption is applied to indicate that $E[\| \epsilon_{\theta}(x_t, t) \|^2] = K$** ($K$ is the vector dimension). However, considering Eq. (27) and Eq. (28) in the proof, **our theorem still holds for $E[\| \epsilon_{\theta}(x_t, t) \|^2] \ge K$**.

To derive a new assumption, we first set a term 

$r_t = E[ \| \epsilon - \epsilon_{\theta}(\sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon, t) \|^2 ]$, 

which indicates the prediction error of neural network $\epsilon_{\theta}(x_t, t)$. The term will vanish to $0$ if $\epsilon_{\theta}(x_t, t)$ is fully accurate in backward denoising. According to the Triangle Inequality, we then have 

$r_t \ge E[ \| \epsilon \|^2 - \| \epsilon_{\theta}(\cdot) \|^2  ] = E [\| \epsilon \|^2] - E[ \| \epsilon_{\theta}(\cdot) \|^2 ] $. 

Because the second moment of Gaussian distribution $E[\| \epsilon \|^2]$ is $K$, we further have 

$(1 / K) E[ \| \epsilon_{\theta}(\sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon, t) \|^2 ] \ge 1 - (r_t / K) $. 

Finally, consider the fact [1] that $x_t$ can be reparameterized as $\sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon, \epsilon \sim N(0, I) $ and let the prediction error $r_t$ vanish, **the above equation motivates us to make the following new assumption**: 

$(1 / K) E[ \| \epsilon_{\theta}(x_t, t) \|^2 ] \ge 1$, 

**which is much weaker than the previous assumption but still makes our Theorem 3.1 hold**.

We perform experiments on 2 datasets to verify our new assumption. For every dataset, we sample $10000$ trajectories $x_{0:T}$ from a well-trained diffusion model. In this process, we collect $10000$ outputs from neural network $\epsilon_{\theta}(x_t, t)$ at every step $t$ and use them to estimate $(1 / K) E[ \| \epsilon_{\theta}(x_t, t) \|^2 ]$. Below are the experiment results.

| Dataset | t=1000 | t=800 | t=600 | t=400 | t=200 | t=1   |
|---------|--------|-------|-------|-------|-------|-------|
| CIFAR-10 | 1.17   | 1.19  | 1.15  | 1.11  | 1.08  | 1.05  |
| CelebA  | 1.25   | 1.21  | 1.23  | 1.17  | 1.11  | 1.09  |

From the above table, we can see that **the estimated value of $(1 / K) E[ \| \epsilon_{\theta}(x_t, t) \|^2 ]$ is consistently larger than $1$ at every backward iteration for both datasets**. Therefore, our new assumption $(1 / K) E[ \| \epsilon_{\theta}(x_t, t) \|^2 ] \ge 1$ has strong empirical support.

Reviewer comment (Theory):
Remark 3.2 is not rigorous. The proof requires $T$ goes to infinity, but it is not true in practice.

Author response:
Thanks for pointing out this. We address your concern by **introducing an alternative precondition that is weaker than the previous one but still makes our proposition hold**. Importantly, our new precondition is supported by current works [1,2].

 In our proof to Proposition A.1, the use of that precondition $T \rightarrow \infty$ is to let $q(x_T | x_0) \rightarrow N(x_T; 0, I)$. Considering the fact [1] that 

$q(x_t | x_0) = N(x_t; \sqrt{\bar{\alpha}_t} x_0, \sqrt{1 - \bar{\alpha}_t} \epsilon)$, 

**we can achieve the goal with a much weaker assumption: $\lim_{t \rightarrow T} \bar{\alpha_t} \rightarrow 0$**. Notably, **this new assumption is a standard configuration in current diffusion models** (e.g., DDPM [1] and SGM [2]), which lets $x_T$ contain no information about $x_0$.

Reviewer comment (Reproducibility):
Lack of detailed settings of experiments: what is the sampling algorithm for obtaining the FID results? What is the detailed network structure (e.g., layer structure and number of hidden neurons) and amount of parameters?

Author response:
For computing the FID scores, we adopt the same sampling procedure as DDPM (see Algo. 2 of [1]) and generate $50000$ samples with $1000$ backward iterations. For the network structure, we typically use U-Net with $4$ layers, $192$ channels, and channel multipliers as $\\{1,2,3,4\\}$, leading to a model size of 300M. We use Adam as the optimization algorithm with a learning rate of $10^{-4}$ and set the dropout ratio as $0.2$. 

Sec. 6.1 of our paper has some other details of the experiment setup. We will merge this answer into that section in the revised version.

Reviewer comment (Experiments):
The authors briefly address the trade-off between increased training time and reduced error propagation (resulting in better FID) for the 32x32 images of CIFAR and ImageNet but do not mention their CelebA experiments on 64x64 images. It is not clear if the benefits scale with the image sizes without an increased overhead as it is possible that the error estimate requires more samples or a larger sampling length $L$.

Author response:
Thanks for pointing out this. As a reminder, **the y-axis in Fig. 4 of our paper in fact represents the time cost of $200$ training steps** (not just $1$ step). We apologize if this mistake might mislead you to the impression that our method is very inefficient.

To address your concern, we have performed an additional trade-off study (showing how the FID score and the time cost change w.r.t. increasing bootstrapping steps $L$) on CelebA. The results are in the following:

| Bootstrapping Steps $L$ | 2      | 3      | 4      | 5      | 6      | 7      | 8      |
|-------------------------|--------|--------|--------|--------|--------|--------|--------|
| FID Score                  | 1.51   | 1.39   | 1.29   | 1.22   | 1.18   | 1.16   | 1.15   |
| **Decrement of FID Score**  | N.A.   | 0.12   | 0.10   | 0.07   | 0.04   | 0.02   | 0.01   |
| Time Cost Per Training Step | 0.37s | 0.45s | 0.55s | 0.63s | 0.72s | 0.81s | 0.92s |

From the above table, we can see the FID scores tend to converge after $L=7$, which is similar to the situations of CIFAR-10 and ImageNet. **The results imply that the performance improvements contributed by our proposed regularization scales with the image size**. Importantly, while our method incurs an extra time cost per training step, it has no impact on the sampling speed of diffusion models, which is more important in practice.

Reviewer comment (Experiments):
Would it be possible to fine-tune pre-trained diffusion models with the regularization term to mitigate this error propagation a-posteriori? If the training time of adding the regularization makes training larger models prohibitive it would be interesting to explore whether it is possible to tune the denoiser network after having trained with just the $L^{nll}$ loss.

Author response:
Thanks for providing such an interesting idea. To verify the feasibility of your idea, we have trained a diffusion model on CIFAR-10 in a way that **only the last 10% training steps involve our proposed regularization**. The results are as below.

| Model                   | All Training Steps | Steps with Regularization | Training Time | FID Score |
|-------------------------|---------------------------|---------------------------|----------------|-----------|
| DDPM                    | 200K |0                         | 10hr           | 3.61      |
| **DDPM w/ Our Proposed Reg.** | **200K** | **20K**                     | **11hr**          | **3.21**      |
| DDPM w/ Our Proposed Reg.      | 200K |200K                      | 20hr           | 2.93      |


From the above table, we can see that the diffusion model is still improved much (i.e., a reduction of the FID score by 11.08%) by partly applying our method, with a minor increase in training time. Therefore, **your idea indeed works empirically and is very useful in improving the efficiency of our proposed regularization**. 

We will include the above experiment and discussion in the revised version.