High variance score function estimates help diffusion models generalize

Thanks for reading the paper and offering comments. Hopefully the new version of the paper (to be posted before AOE Nov 22) addresses your concerns. See below for detailed responses:

Weaknesses:

1. The mathematical derivation in the paper looks like heuristics instead of rigorous proof. It is hard to tell the correctness of the arguments.

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.

2. No experiments are provided to justify the high variance arguments.

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.

3. 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.

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).

4. The writing can be improved.

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.

Questions:

1. 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?

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.

2. The most important issue with this paper is that I cannot tell how the high variance is connected to the generalization of diffusion models.

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

3. Can you provide some numerical experiments to support your claims in the paper?

Various numerical experiments have been added in order to both (i) convince the reader that the high variance phenomenon is real, and (ii) exhibit the inductive biases it produces.

Thanks for reading the paper and offering comments. Hopefully the new version of the paper (to be posted before AOE Nov 22) addresses your concerns. See below for detailed responses:

Weaknesses:

The main concern is that the paper uses some simple diffusion models (1)-(2), as well as the Gaussian mixture training distribution. This limits the applicability of the results. For instance, there are more advanced diffusion models, e.g. VE, VP, sub-VP... What happens in these cases? Is it possible to identify all theses diffusion models by a suitable kernel density problem? The authors may want to comment or explain.

A substantial amount of effort has been put into generalizing the results. In the newest version, two assumptions are no longer required: the training distribution need not be a Gaussian mixture, and score estimator parameters need not be learned independently for different times.

The score estimator is still assumed to be linear in a set of feature maps (which are now allowed to be both state- and time-dependent), but this is actually a fairly general class of estimators and includes many arbitrarily expressive function approximators (e.g., a Fourier basis, Gaussian basis functions, orthogonal polynomials). Importantly, the key assumption is that the score is linear in these features, not linear in states. This is a point of confusion that has been clarified in the text.

One case which is somewhat realistic that is covered now is the case of neural networks in the neural tangent kernel (NTK) regime. Because these are (to very good approximation, in the infinite width limit) linear in their parameters, they are now covered by the main result of the newest version. This takes the main result closer to SOTA models, although obviously there is a ways to go (e.g., what about neural networks in the so-called "rich" learning regime?). We still think this is a substantial theoretical contribution that could serve as a foundation for future work to generalize in various ways.

VP, VE, and sub-VP models are all covered by the main result, and some discussion of their differences has been added to the main text. We agree that this is a practically relevant case worthy of explicit discussion.

Yes, all cases can be identified with a suitable kernel density problem. One nuance is that, in the new (more general) result, where feature maps are allowed to be time-dependent, the result generically includes an effective noise term in the probability flow ODE (as opposed to just a convolution at the end; see below).

Also it is not clear whether the "theoretical" computations can carry over to other distributions than Gaussian mixture (or a mixture of delta mass). The authors may want to comment on this.

This assumption has been removed in the newest version, and turns out to be completely unimportant.

For completeness' sake, here is the main theorem in the newest version:

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

Thanks for reading the paper and offering comments. Hopefully the new version of the paper (to be posted before AOE Nov 22) addresses your concerns. See below for detailed responses:

Weaknesses

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

I thank the authors for their efforts in revising the paper. I highly suggest that the authors highlight their modifications in a different color, especially when the revision is of this significance. Since a lot of the suggested changes are not incorporated in the updated version, I cannot recommend acceptance. Nevertheless, I agree that understanding the generalization behavior of diffusion models is an important topic, and I hope to read the completed paper in the future.

Thanks for reading the paper and offering comments. Hopefully the new version of the paper (to be posted before AOE Nov 22) addresses your concerns. See below for detailed responses:

Weaknesses:

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. 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. I feel in the main result they are getting variation because they do not have enough sample to train each single score function.

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.

In addition, 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. Perhaps the authors can make their results more persuasive by giving several examples that have linear score functions?

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.

Questions:

1. Is there a way to estimate the feature maps $\phi$ when they are not known a priori?

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.

2. How accurately can a linear estimator learn the score function in typical situations? Maybe the authors can comment a little bit on that.

See above answers. A variety of arbitrarily expressive feature maps can be used. Some clarification related to this has been added to the text.

3. 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?

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.

4. If a different sampling distribution is employed, how does the results change?

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.