On the Generalization Properties of Diffusion Models

Summary

This paper studies theoretical bounds on score-matching diffusion models' generalization ability, in terms of KL divergence between the true distribution and the learned distribution generated by the model. It is assumed that the score model is parametrized as a time-dependent (time refers to $t$ in the diffusion process) 2-layer random feature model, and that the training process is the continuous-time gradient flow with respect to the score matching loss $\tilde{\mathcal{L}}$. The authors characterize the generalization bound as a decreasing function of $n$ (the sample size) and $m$ (the model capacity; the hidden layer size), and shows that the error decreases polynomially in $n$ with appropriately chosen early-stopping time $\tau$, under the setup where the target distribution has compact support. They also study the effect of distance between the modes on the theoretical bound in the toy case of fitting a 2-Gaussian-mixture.

Strengths

- The paper successfully builds the setup where the generalization of diffusion models is theoretically analyzable in terms of sample size, model capacity and training time, using common machineries used in ML theory. - The analysis carefully deals with the necessary technical components for bounding the generalization error for diffusion models and the proofs seem sound (legitimate arguments assigned to each component of the analysis). - Overall, I think the paper provides a good starting point for the theoretical study on the generalization of diffusion models and identifies some interesting questions on diffusion model training such as dimension dependency or the choice of optimal early stopping time.

Weaknesses

- Of course, the investigated setup is limited (dealing with 2-layer random feature models) and is far away from what is being used in practice. To some extent this is understandable, considering that the similar limitation is shared among most of the relevant works in the field. However, one can still argue that the setup is obsolete as there are now more 'modern' tools like the neural tangent kernels or mean field limits for performing theoretical analyses on neural networks. - I think the paper has insufficient discussion regarding the optimal solution $\bar{\boldsymbol{\theta}}^\ast$. It does not discuss, e.g., the universal approximation property of the random feature parametrization with respect to the score-matching loss $\tilde{\mathcal{L}}$ which is a sort of modified $L^2$ norm. So it is confusing whether the authors are assuming that the true score function is representable as a continuous random feature model not (so, if $\bar{\tilde{\mathcal{L}}}(\bar{\boldsymbol{\theta}}^\ast)=0$ or not); the term $\bar{\tilde{\mathcal{L}}}(\bar{\boldsymbol{\theta}}^\ast)$ is present in Theorems 1, 2 while absent in Corollaries 1, 2. It is also confusing whether $\mathcal{H}$ in Theorem 1 coincides with the RKHS induced by $k_{\rho_{0}}$. - The value of Section 3.2.2 is unclear to me. It deals only with the toy case of 1D Gaussian mixture with two modes. Additionally, it provides an upper bound that is adversely affected by $\mu$, but there is no guarantee of tightness so it does not sufficiently "explain" the modes shift effect (to do so, one might rather need some sort of lower bound). - Experiments are minimal and loosely designed. It is performed over a single toy case (1D 2-Gaussian-mixture) and it is not reported how many repetitions were tried and whether the results were consistent. Also, it is inconsistent with the theory, which deals with the continuous-time gradient flow. Why use Adam optimizer for the experiments?

Questions

- What is the role of the time embedding function $\mathbf{e}$? Does it have any meaningful effect on the analysis? - In Lemma 3, the authors use the convexity of the diffusion loss $\bar{\tilde{\mathcal{L}}}$ with respect to the trainable parameter $\bar{\mathbf{a}}$, and maybe this is why they consider the random feature model. But won't similar analyses be possible in more general setups using e.g., neural tangent kernels? - Please clarify the second point from the "Weaknesses" section. - In Theorem 2, how does variance of the Gaussian distributions comprising the mixture affect the bound? If the variance is not important, does simply scaling down all the data resolve the adverse modes shift effect claimed in Section 3.2.2? - In the equation right before Line 511 in the appendix, is the minus sign missing in front of $\frac{d}{d\tau} \tilde{\mathcal{L}}(\boldsymbol{\theta}(s))$? I think the subsequent lines using the integral Cauchy-Schwarz inequality also share the sign issue. - In the experiments, do the authors observe that their theory provide quantitative predictions on the optimal early-stopping time (the "transition point" in the training dynamics), or is it just a post-hoc observation? Were the results consistent over multiple runs?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A

Summary

This paper provides some generalization bounds for Diffusion Models when seen as score-based models. These bounds are similar to previous literature for GANs and consider two scenarios. In the first one, the target distribution has compact support, and in the second case, it corresponds to a one-dimensional 2-mode Gaussian mixture.

Strengths

This paper is well-written and provides some generalization bounds on a general and toy problem and a numerical experiment. The sketches of the proofs are well detailed and explained.

Weaknesses

Some of the motivations of the paper can seem obscure and a bit misleading. For instance, from the introduction, it feels that the authors will provide a general result for mode shift distribution while they only provide a theorem for the case of a 1-dimensional mixture of 2 Gaussian. It seems fine to study that toy problem, but it should be advertised as is. Also, the first two paragraphs on Page 2 were a bit obscure.

Questions

"early-stopping estimates" better generalize, do you have a practical approach to early stopping here?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

yes

Summary

The paper aims to characterize the generalization gap of score-based diffusion models. The setting is as follows. The input is $n$ samples from some unknown distribution $p_0$, and the goal is to approximate $p_0$ by a distribution $p_{0,\hat{\theta}_n}$ where the parameter $\hat{\theta}_n$ is a minimizer of the empirical loss function obtained by running gradient descent up to some time $\tau.$ The loss function is the expected difference between the gradients of the log density function of distributions along the trajectories of 2 SDEs (see equation (2)), one initialized at $p_0$ and the other at $p_{0,\hat{\theta}_n}.$ The paper proves a bound on the KL divergence between the true distribution $p_0$ and the learned distribution $p_{0,\hat{\theta}_n}$ in terms of the training time $\tau$ for gradient descent, the sample size $n$ and model parameter $m.$ The bound is obtained by using [Song, Durkan, Murray, Ermon’21]’s main result to bound the aforementioned KL divergence in terms of the empirical loss plus a negligible term, then bound the empirical loss by the population loss. The paper states a general result (theorem 1), then shows how to apply it to the specific case when the target distribution is a mixture of two 1-dimensional Gaussians with the same covariance. They quantitatively show that the error grows when the distance between the two modes of the Gaussians increases, the KL divergence between the true distribution and the learned one becomes worse, which matches the qualitative result given in [Koehler, Heckett, Risteski—ICLR’23].

Strengths

The paper explicitly bounds the generalization error, which is defined as the KL divergence between the true distribution and the learned one, in the model parameters and the sample size. Similar bounds are shown before by [Koehler, Heckett, Risteski—ICLR’23] but this previous work doesn’t make explicit the dependency on the sample size and the model parameters, and further requires that the learned distribution. The results appear to be interesting, if correct.

Weaknesses

The paper might be incorrect, or at least currently doesn’t do a good job of convincing me otherwise. I believe the statement of theorem 1/theorem 3 in the appendix lacks the important assumption that the loss function is convex (see the top line of page 15 of the supplement (proof of Lemma 3) and my 1st question for details). However, the loss function of mixtures of Gaussians appears non-convex (or at least it’s unclear to me why it should be so), so lemma 3 couldn’t be used in the proof of theorem 2 as the paper currently does. Aside from that, the paper is quite hard to understand, and the definition of many terms are not easy to find. For ease of understanding, the author should make clear the difference between $T$, the time to run the diffusion models so that $p_T$ is close to the known prior distribution $\pi$ (e.g. Gaussian), and $\tau,$ the time to train the neural net for gradient estimation. These definitions, along with the definition of $p_T$ and $\pi$, should be clearly stated in the statement of theorem 1. The preconditions of the main theorem (theorem 1) should be stated upfront in the main text instead of in the appendix of the supplement.

Questions

1/ In proof of lemma 3, the top line of page 15 of the supplement, does “by convexity” refer to the fact that the loss function is a convex function? If so, this assumption should be clearly stated as part of theorem 1 or theorem 3. 2/ In the neural net setup, lines 156 and 157 of the supplement, does A essentially play the role of theta? If so, this should be clearly stated. 3/ If lemma 3 requires the assumption that the loss function is convex, then can the authors justify that the loss function for Gaussian mixtures is convex? Else, theorem 2 cannot use lemma 3 like it currently does.

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

1 poor

Contribution

2 fair

Limitations

The paper's proof might be incorrect.

Summary

This paper proves generalization bounds for diffusion models. With specific stopping, they show that the generalization error goes to zero, with a specific upper bound on the rate that scales polynomially with the sample size and the model capacity. Theorem 1 is the main convergence result, while in Theorem 2 they extend their results to the data-dependent setting. The paper is concluded with some supporting experiments.

Strengths

- well-motivated problem and novel formulation - nice theoretical results - enough citations to prior works

Weaknesses

- while it's generally well-written, it could've been better (e.g., a figure could've been used for problem formulation to make it better readable)

Questions

This is a nice theoretical result. Here are some comments/questions. - Equation 3: there are two $dt$'s there, is that right, or it's a typo? Please make it clear - Theorem 1: do you need to know that RKHS to run the algorithm? Or do you just use it in the proofs? This is not that much clear from the context - it is stressed that the upper bound on the generalization is dimension-free, while there is dependence on $d$ in the bound (Equation 18)

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A

Response

I acknowledge the response provided by the authors, specifically the idea of adding new figures. Please also add the clarifications suggested in my comments to the new version of the paper (as promised here). I decided to keep my score unchanged.

A remaining question on specifying the role of approximation error

I have read the rebuttal, and I thank the authors for the detailed response, including additional experiments to appropriately support the theory. I find most of the responses adequate. One remaining question is: although I acknowledge that universal approximation may be a separate problem, in the statements of Corollary 1 and Corollary 2, the bound (right hand side of $\lesssim$) only involves $(\log(d+1)/n)^{\frac{1}{4}}$ and lacks the term $\bar{\tilde{\mathcal{L}}}(\bar{\boldsymbol{\theta}}^*)$ (and also the $m$-dependent term and KL divergence term between $p^T$ and $\pi$). This is confusing, and it is the reason why I initially thought the authors are implicitly assuming universal approximation and infinite width, etc.. Therefore, unless I am missing something, I insist that the authors should carefully clarify what they are omitting in Corollaries 1 and 2, and precisely under which conditions/assumptions one can expect the *KL divergence between $p^0$ and the generated $p^{0, \hat{\boldsymbol{\theta}}_n (\tau^{es})}$ actually goes to zero, polynomially in $n$*.