Paper: High variance score function estimates help diffusion models generalize (ICLR-2024)

Reviewer comment (Theory):
The mathematical derivation in the paper looks like heuristics instead of rigorous proof. It is hard to tell the correctness of the arguments.

Author response:
We agree that this was a weakness of the original version, and the mathematical content has been substantially reorganized (and made more rigorous) to fit a more standard theorem/proof structure. The original result included in the paper has also been substantially generalized to remove two hypotheses; the result no longer requires a gaussian mixture training distribution or that score function parameters are learned separately for different times.

Reviewer comment (Experiments):
No experiments are provided to justify the high variance arguments.

Author response:
New experiments have been done to justify the high variance arguments (although we still refer heavily to previous work). Two important pieces of empirical evidence are that (i) the score function learned by models is usually not the same as the "empirical" score of the data set, and (ii) models trained using what we call the "naive" objective exhibit more memorization.

Reviewer comment (Theory):
Although the starting point of the paper is on the generalization of diffusion models, it is unclear how the high variance of the score estimators helps model generalization.

Author response:
We agree that this could have been better explained in the original version. Two new sections have been added to the newest version to remedy this: one presenting some intuition about generalization, and another discussing concrete examples of how the additional variance discussed in the text implements an inductive bias that encourages interpolation (filling in gaps between training data), extrapolation (extending patterns in the training data), and feature blending (generating samples which include both feature $X$ and feature $Y$ even when training examples only involve one of the two features).

Reviewer comment (Writing):
The writing can be improved.

Author response:
The aforementioned new sections have been added, and the discussion was clarified and reorganized in various places in order to try to remedy this issue.

Reviewer comment (Theory):
The paper only considers the linear score estimator and derives the optimal closed-form solution. How do you know the ground truth score function is linear? For a general score function, can we still have high variance parameters?

Author response:
It is important to note that the score estimators considered in this paper are linear in *features*, not linear in *states*. This means that a variety of arbitrarily expressive function approximators (e.g., a Fourier basis, Gaussian basis functions, orthogonal polynomials) are permitted. This point of confusion has been clarified in the text. 

Another case of interest is that of neural networks in the neural tangent kernel (NTK) regime. These are effectively linear, and the main result of the paper holds for them also. (A brief discussion of this has been added to the paper.) So additional effective variance due to noisy parameter estimates is a phenomenon that appears to happen for a fairly large class of models, although for reasons of mathematical tractability we are unable to discuss all possible models (e.g., neural networks in the "rich" learning regime) here.

Reviewer comment (Theory):
The most important issue with this paper is that I cannot tell how the high variance is connected to the generalization of diffusion models.

Author response:
See above answer. Two new sections have been added to discuss this explicitly.

Reviewer comment (Theory):
My biggest complaint about the paper is that I am not convinced by the paper's argument of the high variance score function estimates at small $t$, which also serves as the paper's motivation. I am not sure why in Eq.6, we are looking at the covariance over $x_t$, since the network is current-position-aware (the network has $x_t$ and $t$ as input). The network prediction targets are often regularized to have a unit norm [1], so the network predicted scale stays the same across $t$, which is not considered in the paper. Also, people have tried to directly predict $x$ through a different parameterization, which will make the prediction targets very stable at small $t$. This would invalidate the paper's analysis, and we still do not see perfect memorization happen.

Author response:
> My biggest complaint about the paper is that I am not convinced by the paper's argument of the high variance score function estimates at small $t$, which also serves as the paper's motivation. I am not sure why in Eq.6, we are looking at the covariance over 
$x_t$, since the network is current-position-aware (the network has and $x_t$ and $t$ as input). 

This brief discussion is just meant to provide some intuition about a crucial difference between what we are calling the "naive" objective and the DSM objective. What actually matters is the covariance of the target as a function of all variables, i.e., with respect to $p(\mathbf{x}_ t | \mathbf{x}_ 0, t) p(\mathbf{x}_ 0) \lambda(t)$. Integrating wrt $x_ 0$ doesn't change anything, and integrating with respect to $t$ (assuming $\lambda(t) \sim \sigma_ t^2$) produces the integral of a constant, and hence proportional to $t_{max}$. 

Other authors (like Song et al. 2021) have explicitly considered the covariance of the target with respect to $\mathbf{x}_t$ alone, so I think it's a reasonable thing to look at. Moreover, in the (now more general) central result of the paper, this variance is indeed the key ingredient that creates what we are calling a type of generalization.

> The network prediction targets are often regularized to have a unit norm [1], so the network predicted scale stays the same across $t$, which is not considered in the paper. Also, people have tried to directly predict $x$
 through a different parameterization, which will make the prediction targets very stable at small $t$. This would invalidate the paper's analysis, and we still do not see perfect memorization happen.

I think this is a misunderstanding, although it is a subtle and important one. The situation is somewhat confusing because different authors use different notation, sometimes using the same symbol to mean very different things. A new appendix has been added that discusses this particular point, since it is potentially a major source of confusion.

Adding to the confusion was the previous choice of learning parameters separately for different times. In the newest version, this restriction has been removed, so feature maps are allowed to depend on both state and time in a fairly arbitrary way (with some restriction, like square-integrability with respect to the relevant distributions, and smoothness).

The situation considered in the current version of the paper is (we will argue) very similar to current SOTA models. The time-sampling distribution used in the main result is
$$ \lambda_ *(t) = \frac{\sigma_ t^2}{\int_0^{t_ {max}} \sigma_ t^2 \ dt} = \frac{\sigma_ t^2}{Z_ {\sigma}} .$$
In the newest version of the paper, we are explicit about scaling the model as $\mathbf{\epsilon}_ {\boldsymbol{\theta}}(\mathbf{x}_ t, t) = \sigma_ t \hat{\mathbf{s}}_ {\mathbf{\theta}}(\mathbf{x}_ t, t)$ (i.e., the function approximator $\mathbf{\epsilon}$ learns the true score multiplied by $\sigma_ t$). With this change, the objective becomes
$$
J_{1}(\boldsymbol{\theta}) = \frac{1}{2 Z_ {\sigma}} \int \ \Vert \boldsymbol{\epsilon}_ {\boldsymbol{\theta}}(\alpha_ t \mathbf{x}_ 0 + \sigma_ t \boldsymbol{\epsilon}, t) - \boldsymbol{\epsilon} \Vert_2^2 \ \mathcal{N}(\boldsymbol{\epsilon}; \mathbf{0}, \mathbf{I}) p(\mathbf{x}_ 0) \ d\boldsymbol{\epsilon} d\mathbf{x}_ 0 dt   $$
This is exactly the objective used by, e.g., Stable Diffusion (Rombach et al 2022; see Eq. 1) and the PNDM paper (Ho et al 2020; see Eq. 14). It turns out that the way we handled the independent-time case hid this scaling (although it was still implicitly there). Some discussion of the variance normalization is now present in the revised version.

Reviewer comment (Theory):
In the introduction, the paper rules out many candidates for generalization, purely based on prior works, intuitions, and observations. In order to make these claims, careful study is required. For example, the authors mention that modern architectures are flexible enough to in principle learn the optimal score. The authors need to back up the claims carefully, because that implies that you can perfectly reconstruct the dataset with these networks, which does not happen.

Author response:
> In the introduction, the paper rules out many candidates for generalization, purely based on prior works, intuitions, and observations. In order to make these claims, careful study is required. For example, the authors mention that modern architectures are flexible enough to in principle learn the optimal score. The authors need to back up the claims carefully, because that implies that you can perfectly reconstruct the dataset with these networks, which does not happen.

We have softened our claims here, since their precise form is just for intuition. However, we maintain that modern architectures are indeed flexible enough to in principle learn the optimal score. What we mean is that approximators like neural networks can approximate arbitrary functions; this is just mathematically true. Indeed, it is empirically true that real diffusion models do not perfectly reconstruct the data set. All we are saying is that this cannot be due to a lack of expressivity on the part of the function approximator. 

Our new theoretical result also covers some related cases, as mentioned above. For example, a fully-connected ReLu network with one hidden layer in the infinite width limit is arbitrarily expressive, and is covered by our main theorem. If trained using the naive objective, we prove theoretically that it *does* perfectly reconstruct the data set.

Reviewer comment (Experiments):
Since the issue considered in this paper is really an empirical one (because the optimal solution to diffusion model training is memorizing all training samples), I highly suggest the authors do some experiments. Related issue: the presentation can be significantly improved with figures. The authors should have more than enough space in the current version.

Author response:
> Since the issue considered in this paper is really an empirical one (because the optimal solution to diffusion model training is memorizing all training samples), I highly suggest the authors do some experiments. Related issue: the presentation can be significantly improved with figures. The authors should have more than enough space in the current version.

We have added some small experiments to the paper, as well as a few small figures. However, it is worth noting that what is probably lacking in the diffusion field right now is not more experiments, but a deeper theoretical understanding of why these models generalize. In this sense, the issue cannot totally be resolved via numerical experiments.

Reviewer comment (Theory):
Where is Eq.8 used in practice? Or, equivalently, which diffusion models used in practice actually underweight small $t$ during training? I see the authors also cite [1], which used a log-normal distribution, and they basically train with more samples at small $t$ compared to large $t$. This fact directly counters one of the assumptions listed in the last paragraph in Sec.3: “the choice of a time sampling distribution that underweights small times…”

Author response:
> Where is Eq.8 used in practice? Or, equivalently, which diffusion models used in practice actually underweight small 
 during training? I see the authors also cite [1], which used a log-normal distribution, and they basically train with more samples at small $t$ compared to large $t$. This fact directly counters one of the assumptions listed in the last paragraph in Sec.3: “the choice of a time sampling distribution that underweights small times…”

You appear to have this backwards. In the revised version, we have made the time-dependence of the DSM objective explicit since we are no longer forcing ourselves to learn the score at different times separately:
$$J_ 1 = \frac{1}{2} \int \lambda(t) \ \Vert \hat{\mathbf{s}}_ {\boldsymbol{\theta}}(\mathbf{x}_ t, t) - \nabla_ {\mathbf{x}_ t} \log p(\mathbf{x}_ t | \mathbf{x}_ 0) \Vert_ 2^2 \ p(\mathbf{x}_ t | \mathbf{x}_ 0) p(\mathbf{x}_ 0) \ d\mathbf{x}_ t d\mathbf{x}_ 0 dt$$

The standard choice (used by Song et al 2021, PNDM, Stable Diffusion...) is to take $\lambda(t) \propto \sigma_t^2$, and this is the choice we consider in this paper. Since $\sigma_t^2 \to 0$ for small $t$ (and in particular, $\sigma_t^2 \propto \Delta t$ for small $t$ in all cases we consider, including VP, VE, sub-VP...the kinds of models defined by Eq 1-3), small times are by construction undersampled. 

The EDM paper identifies time and noise scale, so $\sigma_t = t$. They separately consider the distribution $p_{train}(\sigma)$ of noise scales and an effective weighting factor $\lambda(\sigma) c_ {out}(\sigma)^2$ (see near Eq. 8); these two things can be combined into one effective noise sampling distribution. We have
$$p_{eff}(\sigma) = p_{train}(\sigma) \lambda(\sigma) c_{out}(\sigma)^2  .$$
They recommend $\lambda(\sigma) \sim  c_{out}(\sigma)^2$ (pg 9), and $p_{train}(\sigma)$ to be log-normal. Then
$$p_{eff}(\sigma) = p_{train}(\sigma)$$ 
is log-normal. Log-normal distributions actually approach zero fairly quickly, although not as quickly as $t$---hence, small times / noise scales are indeed undersampled, quite drastically. See Figure 5A from the EDM paper for some illustration of this. Note also, in the same panel, that the loss for the score function at small times is somewhat large after training (partly as a consequence of this undersampling!). Not only does this not counter our assumption, it is in fact evidence in our favor. Maybe the confusion comes from the meaning of "small"; here, we really mean 'close to zero', rather than intermediate times / noise scales. We have tried to make this clearer in the revision.

Reviewer comment (Theory):
I have some doubts on the asymptotic regime considered in this paper. I feel it is not very efficient to use only data in a small time window to train the score function at that time point.

Author response:
This is a good point, and it previously bothered us that we were only able to obtain a result for this (admittedly artificial) case. In the newest version, a substantial amount of work has been put in to generalize this result. Two restrictive hypotheses have been removed from the main result: sample splitting is no longer required (i.e., the feature maps are allowed to depend on time), and the training distribution need not be a Gaussian mixture. Even in the case of time-dependent feature maps, additional variance survives in the limit of many samples (as long as, as in the original result, the number of samples $N$ and the time step $\Delta t$ used to simulate the probability flow ODE scale in a comparable fashion). We think, with this change, the result is much more interesting.

Reviewer comment (Theory):
As far as I am concerned, $\lambda(t)$ in the past papers was introduced to impose weights on the loss function instead of sample splitting. I think the authors should elaborate more on why the sample splitting scheme is a reasonable one.

Author response:
This is a good point, and it previously bothered us that we were only able to obtain a result for this (admittedly artificial) case. In the newest version, a substantial amount of work has been put in to generalize this result. Two restrictive hypotheses have been removed from the main result: sample splitting is no longer required (i.e., the feature maps are allowed to depend on time), and the training distribution need not be a Gaussian mixture. Even in the case of time-dependent feature maps, additional variance survives in the limit of many samples (as long as, as in the original result, the number of samples $N$ and the time step $\Delta t$ used to simulate the probability flow ODE scale in a comparable fashion). We think, with this change, the result is much more interesting.

Reviewer comment (Experiments):
I feel in the main result they are getting variation because they do not have enough sample to train each single score function.

Author response:
This is a good point, and it previously bothered us that we were only able to obtain a result for this (admittedly artificial) case. In the newest version, a substantial amount of work has been put in to generalize this result. Two restrictive hypotheses have been removed from the main result: sample splitting is no longer required (i.e., the feature maps are allowed to depend on time), and the training distribution need not be a Gaussian mixture. Even in the case of time-dependent feature maps, additional variance survives in the limit of many samples (as long as, as in the original result, the number of samples $N$ and the time step $\Delta t$ used to simulate the probability flow ODE scale in a comparable fashion). We think, with this change, the result is much more interesting.

Reviewer comment (Theory):
I feel using only linear score estimator is a bit restricted, as score function is defined as the gradient of the log density, I would expect that it is in general non-linear.

Author response:
This is a point of confusion other reviewers have had, so we have clarified it in the text. Our score function estimator need not be linear in *state* (i.e., $x_t$); it only need to be linear in some set of *feature maps*. Indeed, score functions are generally not linear in state (and the score function that *is* linear in state is that of a Gaussian). 

The class of score estimators linear in some set of feature maps is actually quite broad, and includes many arbitrarily expressive function approximators (e.g., a Fourier basis, Gaussian basis functions, orthogonal polynomials). An important example of an estimator which is linear in this sense is a neural network in the neural tangent kernel (NTK) regime, since it is effectively linear in its (learned) weights; a brief discussion of this example has been added to the paper.

Some more explicit examples of different possible feature maps have been added in order to make this point clearer.

Reviewer comment (Experiments):
Perhaps the authors can make their results more persuasive by giving several examples that have linear score functions.

Author response:
This is a point of confusion other reviewers have had, so we have clarified it in the text. Our score function estimator need not be linear in *state* (i.e., $x_t$); it only need to be linear in some set of *feature maps*. Indeed, score functions are generally not linear in state (and the score function that *is* linear in state is that of a Gaussian). 

The class of score estimators linear in some set of feature maps is actually quite broad, and includes many arbitrarily expressive function approximators (e.g., a Fourier basis, Gaussian basis functions, orthogonal polynomials). An important example of an estimator which is linear in this sense is a neural network in the neural tangent kernel (NTK) regime, since it is effectively linear in its (learned) weights; a brief discussion of this example has been added to the paper.

Some more explicit examples of different possible feature maps have been added in order to make this point clearer.

Reviewer comment (Theory):
Is there a way to estimate the feature maps $\phi$ when they are not known a priori?

Author response:
We consider the problem of choosing 'good' feature maps out of the scope of this paper, since this is extremely similar to the question of which neural network architectures are best for diffusion models (and this is a difficult question for which the answer is not quite known, although it might involve transformers are convolutional neural nets). For our purposes, it suffices to use arbitrary expressive ones (e.g., an NTK neural net, or Fourier basis). Some examples of arbitrarily expressive feature maps have been added to the main text.

Reviewer comment (Experiments):
If we use a better estimator than taking the sample average to estimate the score function, do we get a better result? Will it hurt or improve the generalization ability?

Author response:
Since sample averages are the provably optimal result of (full batch) gradient descent, we have decided to restrict the scope to this case. It is an interesting question for future work what happens when something different is used, e.g., gradient descent using mini-batches. A note of this has been added to the discussion.

Reviewer comment (Experiments):
If a different sampling distribution is employed, how do the results change?

Author response:
Some discussion of different time sampling distributions $\lambda(t)$ (used during training) has been added to the main text. The particular choice of $\lambda(t)$ considered in this paper (and used in various SOTA models, see e.g. Song et al 2021 and Karras et al 2022) is special because of its small-time behavior; sampling distributions with $\lambda(dt) \approx (dt)^n$ for some $n > 1$ do not generalize in the sense described in this paper. Some discussion of this point has been added to the text.

Reviewer comment (Theory):
In practice, the number of samples used for training will be very large, hence $c$ would also be large. In this case, the variance according to the theorem will be small. Do we even expect generalization to happen in such a large-sample situation?

Author response:
In practical situations, we would indeed expect generalization to happen for the following reason. We want to integrate the probability flow ODE accurately in order to sample from the learned distribution; this requires using small $\Delta t$ time steps. Heuristically, using smaller $\Delta t$ means we would like to learn the score function accurately for a larger number of values, which is possible when the number of samples $N$ becomes larger. Theoretically, we want $N$ and $\Delta t$ to a scale in a comparable fashion (i.e., $N \Delta t$ = constant) in order for the model to learn something other than the optimal score. Intuitively, what happens is that you get more and more information to learn an ever larger number of values, so it can never be done completely perfectly. 

In practice, $N$ is not infinite and $\Delta t$ is not zero. However, usually $N$ is somewhat large and $\Delta t$ is somewhat small, making this limit not totally unreasonable. The situation is somewhat similar to the infinite width limit of neural networks: a width does not have to be *that* large in order to resemble an 'infinitely' wide network. 

---------- 


For completeness' sake, here is the main theorem in the newest version. It involves no sample splitting in general, getting rid of the previous major weakness.

**Theorem 1.** [**Linear score estimators trained via DSM asymptotically generalize**] Suppose that the parameters of a linear score estimator are optimized according to the DSM objective using $N$ independent samples from $\lambda_0(t) p(\mathbf{x}_0) p(\mathbf{x}_t | \mathbf{x}_0, t)$. Consider the result of reverse diffusion using this estimator by Euler-integrating the probability flow ODE with a small time step $\Delta t$. If $N \to \infty$ and $\Delta t \to 0$ with $N \Delta t = c$ held constant, then sampling from $\mathbb{E}[ q(\mathbf{x}_0 | \mathbf{x}_T, \boldsymbol{\theta}) ]$ is equivalent to simulating the backwards-time (Ito-interpreted) SDE

$$\dot{\mathbf{x}}_ t = -\beta_ t \mathbf{x}_ t - \frac{1}{2} {g_ t}^2 \mathbf{s}_ *(\mathbf{x}_ t, t) + \boldsymbol{\xi}(\mathbf{x}_ t, t) $$

from $t = t_{max}$ to $t = 0$ with initial condition $\mathbf{x}(t_{max}) = \mathbf{x}_T$. The noise term $\boldsymbol{\xi}(\mathbf{x}_t, t)$ is generically correlated across different times, and has

$$\text{Cov}_ {t, t', \mathbf{x}_ t | t, \mathbf{x}_ {t'} | t'}[ {\xi}_ {i} (\mathbf{x}_ t, t), {\xi}_ {j} (\mathbf{x}_ {t'}, t') ] = V_ {ij} (\mathbf{x}_ {t}, \mathbf{x}_ {t'}, t, t')$$

where we define the $D \times D$  "V kernel" $\mathbf{V}$ via
$$
V_ {ij} := \frac{\delta_ {ij}}{{g_ 0}^2 Z_ {\sigma} c} \  \left( \frac{{g_ {t}}^2}{2 \sigma_ t} \cdot \frac{{g_ {t'}}^2}{2 \sigma_ {t'}} \right)  \ \boldsymbol{\phi}(\mathbf{x}_ t, t)^T \bar{\mathbf{K}}^{-1} \bar{\mathbf{K}}(0) \bar{\mathbf{K}}^{-1} \boldsymbol{\phi}(\mathbf{x}_ {t'}, t') .
$$