Paper: Soft Mixture Denoising: Beyond the Expressive Bottleneck of Diffusion Models (ICLR-2024)

Reviewer comment (Experiments):
Is SMD compatible with fast samplers such as EDM [1]? If it is, can the authors provide results? If not, can the authors suggest how SMD could be modified to be compatible with such fast samplers?

Author response:
Yes, our proposed SMD is compatible with fast samplers---including the stochastic sampler of EDM [1], which applied Langevin dynamics to correct the predictions of denoising neural networks. We perform experiments on CIFAR-10 and LSUN-Conference to show how the EDM sampler and SDM interact. The results are as follows.

| Model                           | CIFAR-10 | LSUN-Conference |
| ------------------------------- | -------- | --------------- |
DDPM | 3.78 | 4.15
DDPM w/ EDM Sampler | 3.57 | 3.95
| DDPM w/ SMD                     | 3.13     | 3.52            |
| DDPM w/ SMD, EDM Sampler        | 2.95     | 3.21            |


We observe that the stochastic sampler of EDM further improves DDPM and our model (i.e., DDPM w/ SMD) on both datasets. For example, the FID score of our model on LSUN-Conference is reduced by 8.81%. Perhaps more importantly, we see that SMD improves the FID for both datasets, with and without EDM. This indicates the advantage of SMD is separate from that of a specific sampler, and that a better sampler can still benefit from SMD.

Reviewer comment (Experiments):
How does the performance of SMD vary as we change the size of $g_\xi$ and $f_\phi$ networks? Does SMD work better if we use larger networks or is it sufficient to use small networks?

Author response:
Thanks for raising this point. We built $g_{\xi}, f_{\phi}$ both as feedforward networks.  Empirically, we found that larger networks $g_{\xi}, f_{\phi}$ lead to better performances, but the marginal improvements decrease as the networks get bigger. Take $g_{\xi}$ with $128$ hidden dimensions as an example, the FID scores of our model (e.g., DDPM w/ SMD) for different network layers as follows:

| Layers of $g_{\xi}$ | CIFAR-10 |
| ------------------- | -------- |
| 2                   | 3.25     |
| **4 (main paper)**       | **3.13**     |
| 6                   | 3.06     |
| 8                   | 3.02     |

We see that larger networks indeed result in better performance, but improvement increases become rather marginal.

Reviewer comment (Theory):
The critique leveled against existing diffusion models seems to be somewhat overstated. These models have achieved considerable success across various applications and can represent complex distributions effectively. The alleged expressive bottleneck is contingent upon the noise scheduling strategy deployed. For instance, in typical diffusion models, a small value of $\beta_t$, such as 0.0001, is assumed to be used as the initial. As indicated in Equation (25), the transition probability $q(x_{t-1} | x_t)$ approaches a Gaussian distribution as $\beta_t$ tends toward zero, which contradicts the claim of an inherent expressive limitation.

Author response:
Diffusion models are indeed performing well and have many successful applications (e.g., Midjourney). We also agree that a mixture model becomes unnecessary when $T\rightarrow \infty$. In practice, however, large $T$ requires more computations and slows down sampling. Thus, current diffusion models are not perfect. For example, **recent works (e.g., DDIM [2] and DPM [3]) show that their performances degrade significantly for few backward iterations**. Our paper aims to solve this problem by introducing a new backward denoising paradigm that makes diffusion models more expressive. **As shown in our paper (Fig. 1 and Fig.3), diffusion models with the proposed SMD approach are much less affected by the number of backward iterations**. Consequently, we believe our proposed SMD has addressed a serious weakness of current diffusion models and can improve sampling speed and reduce computational cost.

Regarding your concerns about “expressive bottleneck being contingent upon the noise scheduling strategy” and $\beta_t$, we note that **our Theorem 3.1, which shows the denoising error $M_t$ is uniformly unbounded, is independent of the variance schedule $\beta_t$. In other words, regardless of the scale of $\beta_t$, $q(x_{t-1} | x_t)$ can be arbitrarily too complex for $p_{\theta}(x_{t-1} | x_t)$ to approximate**. Additionally, $\beta_t$ is only initially small and gradually becomes non-negligible as $t$ increases from $1$ to $T$. As evidence, Fig. 1 and Fig. 3 of our paper show that vanilla diffusion models perform poorly with few backward denoising iterations, where $\beta_t$ becomes large even for small $t$.

Reviewer comment (Theory):
There appears to be a flaw in the derivation of local denoising error $M_t$. The associated loss term in $L_{t-1}$ is predicated on the KL divergence $KL[q(x_{t-1} | x_t, x_0) || p_\theta(x_{t-1} | x_t)]$. Here, $q(x_{t-1} | x_t, x_0)$, which is a known Gaussian distribution, should not be conflated with $q(x_{t-1} | x_t)$, which represents an unknown distribution. The validity of Theorems 3.1 and 3.2 is reliant on the accurate definition of $M_t$.

Author response:
Thank you for raising this concern. We suspect that your concern stems from Eq. (3) of our paper, which includes the term $L_{t-1}$  *but also other terms*. When combined with the other terms, the full DDPM loss is $L_{p_{\theta}} = E_{q}[-\ln \frac{p_{\theta}(x_{0:T})}{q(x_{1:T}\vert x_0)}]$ (see also Eq. 3 of [1]). With the following proposition, we will show that $M_t$ as defined in the paper (with $q(x_{t-1} | x_t)$ and **not** $q(x_{t-1} | x_t, x_0)$) is appropriate for studying this loss.

**Proposition: For a perfectly optimized diffusion model $p_{\theta} = \arg\min_{p_{\theta}'} L_{p_{\theta}'}$, the equality $p_{\theta}(x_{t-1} | x_t) = q(x_{t-1} | x_t)$ holds for every denoising iteration $t \in [1, T]$. Here the two terms $p_{\theta}, L_{p_{\theta}}$ respectively represent a diffusion model $[ p_{\theta}(x_{t-1} | x_t) ]_{t \in [1, T]}$ and its loss**.

**Proof**: Following DDPM [1], we formulate the loss function of diffusion models as

$L_{p_{\theta}} = E_{q}[-\ln \frac{p_{\theta}(x_{0:T})}{q(x_{1:T}|x_0)}] = D_{KL}(q(x_{0:T}) || p_{\theta}(x_{0:T})) - E_q[q(x_0)]$.

Note that the second expectation term $E_q[\cdot]$ is a constant and the first KL-divergence term $D_{KL}(\cdot)$ reaches its minimum $0$ when $p_{\theta}(x_{0:T})$ equals $q(x_{0:T})$. Therefore, for a perfectly optimized diffusion model $p_{\theta} = \min_{p_{\theta'}} L_{p_{\theta'}}$, we have $p_{\theta}(x_{0:T}) = q(x_{0:T})$. Then, for every iteration $t \in [1, T]$, we have

$p_{\theta}(x_t, x_{t-1}) = \int p_{\theta}(x_{0:T}) dx_{1:t-2} dx_{t+1:T} = \int q(x_{0:T}) dx_{1:t-2} dx_{t+1:T} = q(x_t, x_{t-1})$.

Similarly, we get $p_{\theta}(x_t) = \int p_{\theta}(x_t, x_{t-1}) dx_{t,t-1} = \int q(x_t, x_{t-1}) dx_{t,t-1} = q(x_t)$. Based on the above two equations, we finally derive 

$p_{\theta}( x_{t-1} \vert x_t) = \frac{p_{\theta}( x_{t-1}, x_t)}{p_{\theta}(x_t)} = \frac{q( x_{t-1}, x_t)}{q(x_t)} = q( x_{t-1} \vert x_t)$, 

proving the proposition.

**From the above proposition, we can see that $p_{\theta}(x_{t-1} | x_t)$ is in fact optimized towards $q(x_{t-1} | x_t)$ for minimizing $L_{p_{\theta}}$. Therefore, it is appropriate to define the denoising error $M_t$ of $p_{\theta}(x_{t-1} | x_t)$ as its distributional gap to $q(x_{t-1} | x_t)$, and not $q(x_{t-1} | x_t, x_0)$**.

Reviewer comment (Writing):
There seems to be a typographical error involving a comma in the superscript at the end of Equation (3).

Author response:
**Response Q1**: We have now removed this. Thanks!

Reviewer comment (Experiments):
Could you detail the noise schedule utilized in your algorithm? The experiments section suggests that the original DDPM scheduling is retained while only the model component is modified. Considering that your paper emphasizes the importance of shorter chains and the expressiveness issue within them, it would be beneficial to see experimentation with significantly fewer steps to underscore the advantages of your proposed Soft Mixture Denoising (SMD).

Author response:
**Response Q2**: For the noise schedule, we have followed the linear schedule of the standard DDPM [1], increasing from $\beta_1 = 10^{-4}$ to $\beta_T = 0.02$. For your suggested “experimentation with significantly fewer steps”, please check **Fig. 1 and Fig. 3 of the paper, which both show our proposed approach significantly improves current diffusion models in the case of few backward iterations**. For example, Fig. 1 shows that the diffusion model with our approach achieves an FID score of $6.85$ with only $T = 100$ backward iterations, which is almost half of the FID performance (i.e., $11.29$) of the vanilla diffusion model with the same number of backward iterations.

Reviewer comment (Novelty):
The SMD approach bears resemblance to a Variational Autoencoder (VAE) in its structure. Could you confirm if this observation is accurate or elaborate on the distinctions between SMD and VAE?

Author response:
**Response Q3**: Our proposed approach and VAE both have a concept of latent variables, but they are very different. While VAE builds mappings between real samples and latent variables, SMD resorts to a latent variable $z_t$ to only capture the potential mixture structure $q(z_t | x_t)$ of $q(x_{t-1} | x_{t})$. Therefore, in contrast to VAEs, the SMD $z_t$ does not encode a full sample, nor do we put a prior on its distribution (cf. the usually Gaussian prior used in VAEs).