Diffusion Models Revisited (VIII): Optimal Diffusion Variance Estimation (Part 2)
In the previous post, Diffusion Models Revisited (VII): Optimal Diffusion Variance Estimation (Part 1), we introduced and derived the optimal variance estimation result from Analytic-DPM for diffusion models — an analytic estimate of the optimal variance given an already-trained generative diffusion model. Experiments showed that this estimate indeed improves generation quality.
In this post we continue by introducing the upgraded version of Analytic-DPM, from a paper by the same author team, Estimating the Optimal Covariance with Imperfect Mean in Diffusion Probabilistic Models, referred to in the official GitHub repository as "Extended-Analytic-DPM" — a name we'll adopt here as well.
Recap of Results
The previous post, building on DDIM, derived that the optimal variance for the DDIM generative process should be
\begin{equation}\sigma_t^2 + \gamma_t^2\bar{\sigma}_t^2\end{equation}
where $\bar{\sigma}_t^2$ is the variance of the distribution $p(\boldsymbol{x}_0|\boldsymbol{x}_t)$, and it has the following estimate (here we take the result from "Variance Estimate 2"):
\begin{equation}\bar{\sigma}_t^2 = \frac{\bar{\beta}_t^2}{\bar{\alpha}_t^2}\left(1 - \frac{1}{d}\mathbb{E}_{\boldsymbol{x}_t\sim p(\boldsymbol{x}_t)}\left[ \Vert\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\Vert^2\right]\right)\label{eq:basic}\end{equation}
Looking back, the estimation idea isn't actually that hard. Suppose
\begin{equation}\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t) = \frac{1}{\bar{\alpha}_t}\left(\boldsymbol{x}_t - \bar{\beta}_t \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\right)\label{eq:bar-mu}\end{equation}
has already accurately predicted the mean vector of the distribution $p(\boldsymbol{x}_0|\boldsymbol{x}_t)$. Then, by definition, the covariance is given by
\begin{equation}\begin{aligned} \boldsymbol{\Sigma}(\boldsymbol{x}_t)=&\, \mathbb{E}_{\boldsymbol{x}_0\sim p(\boldsymbol{x}_0|\boldsymbol{x}_t)}\left[\left(\boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\right)\left(\boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\right)^{\top}\right] \\ =&\, \frac{1}{\bar{\alpha}_t^2}\mathbb{E}_{\boldsymbol{x}_0\sim p(\boldsymbol{x}_0|\boldsymbol{x}_t)}\left[\left(\boldsymbol{x}_t - \bar{\alpha}_t\boldsymbol{x}_0\right)\left(\boldsymbol{x}_t - \bar{\alpha}_t\boldsymbol{x}_0\right)^{\top}\right] - \frac{\bar{\beta}_t^2}{\bar{\alpha}_t^2} \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)^{\top}\\ \end{aligned}\label{eq:full-cov}\end{equation}
Averaging both sides over $\boldsymbol{x}_t\sim p(\boldsymbol{x}_t)$ to remove the dependence on $\boldsymbol{x}_t$,
\begin{equation} \boldsymbol{\Sigma}_t = \mathbb{E}_{\boldsymbol{x}_t\sim p(\boldsymbol{x}_t)}[\boldsymbol{\Sigma}(\boldsymbol{x}_t)] = \frac{\bar{\beta}_t^2}{\bar{\alpha}_t^2}\left(\boldsymbol{I} - \mathbb{E}_{\boldsymbol{x}_t\sim p(\boldsymbol{x}_t)}\left[ \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)^{\top}\right]\right)\label{eq:uncond-var-2}\end{equation}
Finally, averaging the diagonal elements to turn it into a scalar (i.e., treating the covariance as a multiple of the identity matrix), i.e., $\bar{\sigma}_t^2 = \text{Tr}(\boldsymbol{\Sigma}_t)/d$, gives us the estimate $\eqref{eq:basic}$.
How to Improve
Before formally introducing Extended-Analytic-DPM, let's first think about where Analytic-DPM still has room for improvement.
A little thought reveals quite a few possibilities. For instance, Analytic-DPM assumes that the covariance matrix of the normal distribution used to approximate $p(\boldsymbol{x}_0|\boldsymbol{x}_t)$ is designed as $\bar{\sigma}_t^2\boldsymbol{I}$, i.e., a diagonal matrix with identical diagonal entries. A direct improvement, then, is to allow the diagonal entries to differ, i.e., $\text{diag}(\bar{\boldsymbol{\sigma}}_t^2)$ (here we adopt the convention that vector multiplication is performed via the Hadamard product, e.g., $\boldsymbol{x}^2=\boldsymbol{x}\otimes \boldsymbol{x}$). The corresponding result only needs to consider the diagonal part of $\boldsymbol{\Sigma}_t$, so starting from equation $\eqref{eq:uncond-var-2}$, we can obtain the corresponding estimate
\begin{equation}\bar{\boldsymbol{\sigma}}_t^2 = \frac{\bar{\beta}_t^2}{\bar{\alpha}_t^2}\left(\boldsymbol{1}_d - \mathbb{E}_{\boldsymbol{x}_t\sim p(\boldsymbol{x}_t)}\left[ \boldsymbol{\epsilon}_{\boldsymbol{\theta}}^2(\boldsymbol{x}_t, t)\right]\right) \end{equation}
where $\boldsymbol{1}_d$ is the $d$-dimensional all-ones vector. A further improvement is to retain the dependence of $\bar{\boldsymbol{\sigma}}_t^2$ on $\boldsymbol{x}_t$, i.e., to consider $\bar{\boldsymbol{\sigma}}_t^2(\boldsymbol{x}_t)$ — this is analogous to $\boldsymbol{\mu}(\boldsymbol{x}_t)$, requiring a model that takes $\boldsymbol{x}_t$ as input to learn it.
So could we consider the full $\boldsymbol{\Sigma}_t$? In theory, yes; in practice, essentially not, because the full $\boldsymbol{\Sigma}_t$ is a $d\times d$ matrix, and for image data, $d$ is the total number of pixels in the image — even for CIFAR-10 this is already $d=32^2\times 3=3072$, not to mention higher-resolution images. So given the practical experimental setting, the storage and computation costs of a $d\times d$ matrix are simply too high.
Beyond this, there's an issue that perhaps many readers haven't noticed: all the analytic derivations above rely on $\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t) = \mathbb{E}_{\boldsymbol{x}_0\sim p(\boldsymbol{x}_0|\boldsymbol{x}_t)}[\boldsymbol{x}_0]$, but in fact $\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)$ is learned by a model, and it need not be exactly equal to the true mean $\mathbb{E}_{\boldsymbol{x}_0\sim p(\boldsymbol{x}_0|\boldsymbol{x}_t)}[\boldsymbol{x}_0]$. This is precisely what "Imperfect Mean" in the title of the Extended-Analytic-DPM paper refers to. Improving the estimation result under the Imperfect Mean setting would be much more practically meaningful.
Maximum Likelihood
Suppose the mean model $\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)$ has already been trained beforehand. Then the only remaining parameter of the distribution $\mathcal{N}(\boldsymbol{x}_0;\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t),\bar{\sigma}_t^2\boldsymbol{I})$ to be determined is $\bar{\sigma}_t^2$, and the corresponding negative log-likelihood is
\begin{equation}\begin{aligned} &\, \mathbb{E}_{\boldsymbol{x}_t\sim p(\boldsymbol{x}_t)}\mathbb{E}_{\boldsymbol{x}_0\sim p(\boldsymbol{x}_0|\boldsymbol{x}_t)}\left[-\log \mathcal{N}(\boldsymbol{x}_0;\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t),\bar{\sigma}_t^2\boldsymbol{I})\right] \\ =&\, \frac{\mathbb{E}_{\boldsymbol{x}_t,\boldsymbol{x}_0\sim p(\boldsymbol{x}_t|\boldsymbol{x}_0)\tilde{p}(\boldsymbol{x}_0)}\left[\Vert\boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\Vert^2\right]}{2\bar{\sigma}_t^2} + \frac{d}{2}\log \bar{\sigma}_t^2 + \frac{d}{2}\log 2\pi \\ \end{aligned}\label{eq:neg-log}\end{equation}
Solving for the minimum, we find it is exactly
\begin{equation}\bar{\sigma}_t^2 = \frac{1}{d}\mathbb{E}_{\boldsymbol{x}_t,\boldsymbol{x}_0\sim p(\boldsymbol{x}_t|\boldsymbol{x}_0)\tilde{p}(\boldsymbol{x}_0)}\left[\Vert\boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\Vert^2\right]\end{equation}
The key feature here is that $\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)$ need not be the exact mean, so the second equality in equation $\eqref{eq:full-cov}$ no longer holds — only the first equality does. Substituting equation $\eqref{eq:bar-mu}$ in, we get
\begin{equation}\bar{\sigma}_t^2 = \frac{\bar{\beta}_t^2}{\bar{\alpha}_t^2 d}\mathbb{E}_{\boldsymbol{x}_0\sim \tilde{p}(\boldsymbol{x}_0),\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0},\boldsymbol{I})}\left[\left\Vert\boldsymbol{\varepsilon} - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\bar{\alpha}_t\boldsymbol{x}_0 + \bar{\beta}_t\boldsymbol{\varepsilon}, t)\right\Vert^2\right]\end{equation}
Of course, here we've only analyzed the simple case where the covariance matrix is $\bar{\sigma}_t^2\boldsymbol{I}$; we can also consider the more general diagonal covariance case, i.e., $\mathcal{N}(\boldsymbol{x}_0;\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t),\text{diag}(\bar{\boldsymbol{\sigma}}_t^2))$, for which the corresponding result is
\begin{equation}\bar{\boldsymbol{\sigma}}_t^2 = \frac{\bar{\beta}_t^2}{\bar{\alpha}_t^2 }\mathbb{E}_{\boldsymbol{x}_0\sim \tilde{p}(\boldsymbol{x}_0),\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0},\boldsymbol{I})}\left[\left(\boldsymbol{\varepsilon} - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\bar{\alpha}_t\boldsymbol{x}_0 + \bar{\beta}_t\boldsymbol{\varepsilon}, t)\right)^2\right]\end{equation}
Conditional Variance
If we want to obtain the covariance $\text{diag}(\bar{\boldsymbol{\sigma}}_t^2(\boldsymbol{x}_t))$ conditioned on $\boldsymbol{x}_t$, this is equivalent to computing each component independently, which amounts to dropping the averaging step $\mathbb{E}_{\boldsymbol{x}_t\sim p(\boldsymbol{x}_t)}$:
\begin{equation}\bar{\boldsymbol{\sigma}}_t^2(\boldsymbol{x}_t) = \mathbb{E}_{\boldsymbol{x}_0\sim p(\boldsymbol{x}_0|\boldsymbol{x}_t)}\left[(\boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t))^2\right] = \frac{\bar{\beta}_t^2}{\bar{\alpha}_t^2}\mathbb{E}_{\boldsymbol{x}_0\sim p(\boldsymbol{x}_0|\boldsymbol{x}_t)}\left[\left(\boldsymbol{\epsilon}_t - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\right)^2\right] \end{equation}
where $\boldsymbol{\epsilon}_t = \frac{\boldsymbol{x}_t - \bar{\alpha}_t \boldsymbol{x}_0}{\bar{\beta}_t}$. Just as in the previous post, using
\begin{equation}\mathbb{E}_{\boldsymbol{x}}[\boldsymbol{x}] = \mathop{\text{argmin}}_{\boldsymbol{\mu}}\mathbb{E}_{\boldsymbol{x}}\left[\Vert \boldsymbol{x} - \boldsymbol{\mu}\Vert^2\right]\label{eq:mean-opt}\end{equation}
we obtain
\begin{equation}\begin{aligned} \bar{\boldsymbol{\sigma}}_t^2(\boldsymbol{x}_t) =&\, \frac{\bar{\beta}_t^2}{\bar{\alpha}_t^2}\mathbb{E}_{\boldsymbol{x}_0\sim p(\boldsymbol{x}_0|\boldsymbol{x}_t)}\left[\left(\boldsymbol{\epsilon}_t - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\right)^2\right] \\ =&\, \frac{\bar{\beta}_t^2}{\bar{\alpha}_t^2}\mathop{\text{argmin}}_{\boldsymbol{g}}\mathbb{E}_{\boldsymbol{x}_0\sim p(\boldsymbol{x}_0|\boldsymbol{x}_t)}\left[\left\Vert\left(\boldsymbol{\epsilon}_t - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\right)^2-\boldsymbol{g}\right\Vert^2\right] \\ =&\, \frac{\bar{\beta}_t^2}{\bar{\alpha}_t^2}\mathop{\text{argmin}}_{\boldsymbol{g}(\boldsymbol{x}_t)}\mathbb{E}_{\boldsymbol{x}_t\sim p(\boldsymbol{x}_t)}\mathbb{E}_{\boldsymbol{x}_0\sim p(\boldsymbol{x}_0|\boldsymbol{x}_t)}\left[\left\Vert\left(\boldsymbol{\epsilon}_t - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\right)^2-\boldsymbol{g}(\boldsymbol{x}_t)\right\Vert^2\right] \\ =&\, \frac{\bar{\beta}_t^2}{\bar{\alpha}_t^2}\mathop{\text{argmin}}_{\boldsymbol{g}(\boldsymbol{x}_t)}\mathbb{E}_{\boldsymbol{x}_t,\boldsymbol{x}_0\sim p(\boldsymbol{x}_t|\boldsymbol{x}_0)\tilde{p}(\boldsymbol{x}_0)}\left[\left\Vert\left(\boldsymbol{\epsilon}_t - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\right)^2-\boldsymbol{g}(\boldsymbol{x}_t)\right\Vert^2\right] \\ \end{aligned}\end{equation}
This is the "NPR-DPM" scheme for learning conditional variance in Extended-Analytic-DPM. The original paper also proposes an "SN-DPM" scheme, which is based on the Perfect Mean assumption rather than Imperfect Mean. However, the paper's experimental results show that SN-DPM actually outperforms NPR-DPM — meaning that although the paper claims to be solving the Imperfect Mean problem, the experiments show that the Perfect Mean assumption scheme works better. This in turn suggests that the Perfect Mean assumption is actually a very good fit for practice; in other words, the Imperfect Mean problem can effectively be regarded as nonexistent.
Two Stages
Readers might wonder: didn't we say earlier that the learnable variance in Improved Denoising Diffusion Probabilistic Models increases training difficulty? So why does Extended-Analytic-DPM go back to building a trainable variance model?
We know that DDPM offers two variance schemes, $\sigma_t = \frac{\bar{\beta}_{t-1}}{\bar{\beta}_t}\beta_t$ and $\sigma_t = \beta_t$, and these two simple schemes already work quite well. This indirectly suggests that fine-tuning the variance more precisely has little effect on generation quality (at least for full $T$-step diffusion) — what really matters is learning $\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)$, with the variance playing only a "icing on the cake" role. If we treat the variance as a learnable parameter or model and train it jointly with the mean model $\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)$, then the variance, changing as training progresses, will seriously interfere with the learning process of the mean model $\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)$, violating the principle that "$\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)$ comes first, variance is secondary."
The cleverness of Extended-Analytic-DPM lies in its two-stage training scheme: first train the mean model $\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)$ using the original fixed-variance setup, then freeze that model and reuse most of its parameters to learn a variance model. This achieves three benefits at once:
First, it reduces parameter count and training cost;
Second, it allows reuse of an already-trained mean model;
Third, the training process is more stable.
Personal Reflections
At this point, the introduction to Extended-Analytic-DPM is essentially complete. Attentive readers may feel that, if the previous post's Analytic-DPM result was "dazzling," this post's Extended-Analytic-DPM feels rather ordinary by comparison — nothing especially striking. One could say Extended-Analytic-DPM is a fairly straightforward extension of Analytic-DPM; although the experimental results show it does bring a decent improvement, overall it comes across as somewhat unremarkable. Of course, this is largely because Analytic-DPM "set the bar so high beforehand" that this work looks a bit dimmer by comparison — in itself, it's still a fairly solid piece of work.
Furthermore, as mentioned above, the experimental results show that SN-DPM, based on the Perfect Mean assumption, outperforms NPR-DPM, which is based on the Imperfect Mean assumption. This result also makes the original paper's title somewhat of a misnomer — since the experiments show the Perfect Mean scheme works better, this conversely implies that the Imperfect Mean problem can be regarded as nonexistent. The original paper doesn't further analyze or comment on this result. The author wonders whether this might be related to the bias in variance estimation: as we know, directly estimating variance using the "divide by $n$" formula is biased, and NPR-DPM is built precisely on this operation, whereas SN-DPM directly estimates the second moment, and second-moment estimation is unbiased. This feels somewhat reasonable, but it doesn't fully explain things either — a bit puzzling.
Finally, I wonder if readers share the same question I have: given $\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)$, why not directly use a negative log-likelihood loss like equation $\eqref{eq:neg-log}$ to learn the variance, rather than redesigning the two MSE-form losses of NPR-DPM or SN-DPM? Is there something special about MSE-form losses? I haven't yet found an answer to this either.
Summary
This post introduced the optimal variance estimation result for diffusion models from the upgraded version of Analytic-DPM — "Extended-Analytic-DPM." It primarily derives results for the imperfect-mean setting and proposes a scheme for learning conditional variance.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.