Algorithm- and Data-Dependent Generalization Bounds for Diffusion Models

We thank the reviewer for their insightful comments. In the following, we address all the remarks and questions.

Weakness 1 and Q1

We emphasize that Theorems 4.1 and 4.2 are not only consequences of existing results, but also from our theory, developed in Section 3. Indeed, the conclusion of Section 3 is that we boiled down the problem of score approximation to a well-defined statistical learning problem (the generalization error of Theorem 3.1), where a loss function $\ell$ is well-defined. Then, our theorems should be thought as a continuation of Section 3 by exploiting this setting to involve the benefits of statistical learning (i.e. the dynamics of the considered optimization algorithm) for score-based diffusion models.

We thank the reviewer for the question of satisfying the assumptions of Theorems 4.1 and 4.2. Let us first note that the sub-Gaussian or bounded loss assumptions are very classical in the learning theory literature [HNK+20,MWZZ18,NHD+19]. Verifying these assumptions will depend simultaneously on the dataset and the architecture of the score network, as they refer to the denoising score-matching loss. While it might not be a direct consequence of the assumptions required by our theory in Section 3, we expect that these conditions on the denoising score matching loss can be fulfilled in simple settings, like image datasets. In the final version of the paper, we will add a discussion on satisfying these assumptions for particular data structures and leave a more precise understanding of this matter to future work.

Weakness 2 and Q2

Thank you for this insightful remark. These topological complexities have been used in several previous works (see [ADS+24] and references therein). They intuitively represent different notions of ``complexity'' of the trajectory of the optimizer (eg, SGD, ADAM, ...) near a local minimum, as is briefly explained in the appendix. They essentially compute the "intrinsic dimension" of the local minimum found by the optimization algorithm (by looking at the topological properties of the trajectory), hence, they are expected to perform better than bounds that are based on the parameter count, which is always larger than the intrinsic dimension.

In the final version of the paper, we will use the additional space to provide a more intuitive explanation of the meaning of this quantity.

Concerning how these quantities should be interpreted within the context of diffusion models, let us detail this. Intuitively, the trajectory of the optimizer should be more complex when the learning task is more difficult. In the context of diffusion models, it means that several parameters should affect it like, for instance, the discretization step of the forward process simulation, as the loss is more sensitive for small times. In the final version, we will add a discussion about these aspects. Thank you for raising this and help us improving the quality of our work.

Weakness 3

One of our goal is to relate the performance of diffusion models to the optimization dynamics of the score network. For instance, this allows us to connect (part of) the score estimation error to the norms of the gradients encountered during training. To the best of our knowledge, such theoretical observations are new in the context of diffusion models. Therefore, we believe that this provides insights towards better training of diffusion models. For instance, Theorem 4.1 is supporting the idea that the generalization error of diffusion models should improve when the model converges to a flat minimum in the parameter space, as it is commonly found in the supervised learning literature [1,2].

Q3

The batch size (as well as some other hyperparameters) does not appear explicitly in our theoretical contributions. However, note that these hyperparameters appear implicitly through the complexity measures that we use in Theorems 4.1 and 4.2 (through the gradient norms and the topological complexities). Figure 1 reports the Wasserstein metric and the FID calculated on the test set. We see in Figure 1 that both metrics decrease, stabilize and then increase for higher learning rates. This illustrates the existence of a trade-off, in the sense that the test error depends on both the optimization error and the generalization error, whose analysis is not the aim of this paper. Finally, Figure 1 was obtained using very simple models and datasets and its purpose is to illustrate that the learning algorithm affects the test performance of diffusion models, hence, justifying our approach and our main decomposition.

Q4

Incorporating the noise schedule is definitely an interesting suggestion by the reviewer.

In our setting and under our assumptions (especially the finite Fisher information), we believe that it is non trivial to explicitly factor the noise schedule in our contributions and we leave this fr future work. However, the reviewer is right to mention that the noise schedule is known to affect the performance. In particular, as the noise schedule affects the properties of the loss function, should have an impact on the optimization dynamics and, hence, on the generalization error term (i.e., the second term $G_\lambda$ in our decomposition). Therefore, the effect of the noise schedule is implicitly taken into account by the complexity measures studied in Section 4. In the final version, we will add a discussion about this remark from the reviewer.

Q5

It has been a long standing topic in learning theory that what is a good notion of "model complexity" for neural networks. It has been shown in numerous works that the good notion is definitely not the number of parameters of the model.

Within this context, the results in Section 4 suggest two notions of model complexity, that are implicitly measured by the quantities appearing in Theorems 4.1 (gradient norm) and 4.2 (topological complexity). Moreover, our framework is compatible with more classical measures of the complexity of the model class, such as the Rademacher complexity, so that they could be incorporated in our analysis of the second term of the decomposition (the generalization error). However, we believe that such complexity metric would have little empirical interest but we are open to any suggestion of the reviewer regarding other pertinent complexity metrics. To sum up, one of the goals of our paper is to introduce a notion of model complexity through the analysis of the second term on our decomposition. Our theoretical contributions allow us to do that by upper bounding the other term in the decomposition.

Q6

We thank the reviewer for the opportunity to further clarify the precise role of the complexity measures introduced in our work and their relationship with the optimization process.

Because of a recent update to Neurips guidelines, we are not allowed to include figures in our rebuttal and have enough characters for only one table. However, we have produced additional tables to satisfy the reviewer's request, tracking the denoising score-matching (DSM) loss against our complexity measures. Indeed, the DSM loss is equal to the empirical explicit score-matching (ESM) loss up to a constant independent of model parameters. We will also include these experiments, along with more figures, in the final version of the paper.

Across these experiments, contrary to the generalization gap, we observe no simple relationship between the train ESM loss and our complexity measures, as clearly exemplified in Table 1, where a low complexity value is observed concurrently with high and low train loss values. In Table 1, where all hyperparameters achieve a similarly low ESM loss, we see non-trivial generalization patterns showing a clear positive correlation with our topological or gradient norms-based complexity measures. This confirms that we have not collected spurious observations from trivial regimes and that our complexity measures are relevant quantities

Table 1: ADAM optimizer on the flowers dataset. Train/test Denoising Score-Matching loss (DSM) vs several complexity metrics and corresponding learning rate and batch size.

[1] Kaiyue Wen, Zhiyuan Li, and Tengyu Ma. Sharpness Minimization Algorithms Do Not Only Minimize Sharpness To Achieve Better Generalization. In Conference on Neural Information Processing Systems, 2023.

[2] Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos J. Storkey. Three Factors Influencing Minima in SGD. arXiv, abs/1711.04623, 2017.

We thank the reviewer for their detailed review, which will help improve our work.

Weakness 1 and question 1, the algorithm-dependent generalization analysis in the paper focuses primarily on optimizers, and the effects of model architectures do not appear to be discussed.

Thank you for this insightful remark. Indeed, our bounds (Theorems 4.1 and 4.2) have no explicit dependence on the model architecture, but an implicit one. More precisely, our theory does not require explicit assumption on the structure of the model class, which makes it applicable to a broad range of architectures, even the ones that are used in practical applications, such as the UNet.
The generality of this approach stems from the generalization bounds we used, derived from PAC-Bayes theory (Alquier 2024), which allow considering any architecture by relocating the intrinsic properties of the model within a KL divergence term, translated here by the gradient norms (Thm 4.1) and the topological complexity measures (Thm 4.2). These terms take implicitly into account the model architecture in our generalization bounds, as it is clear that the model architecture will have an impact on the gradient norms, for instance. In the next version of the paper, we will add a precise discussion about this matter.

Weakness 2

Thank you for catching this issue regarding the axis of some figures in the appendix, we will correct it in the final version of the paper. This formatting issue is due to the fact that in these experiments, the positive magnitude with scale $\sqrt{n}$ saturates to its maximum value, so all points on the y axis have approximately the same value. The fact that the positive magnitude can saturate when the scale parameter is too high has been documented in previous works, see for instance [ADS+24] and references therein. We reported these bottom-right figures for the sake of completeness; they do not yield significant experimental insights.

Weakness 3 and Question 2

We thank the reviewer for pointing out several typos, we will correct all of them in the final version of the paper. Here are a few additional comments.

• The notation $B(0,D)$ indeed refers to the ball of radius $D$, we will add this definition to the assumption.
• L669 $\sigma$ is indeed a notation for $\sqrt{2}$ it is a typo and will be replaced by $\sqrt{2}$.
• L715 and L716, by Hölder's inequality we mean that we use the inequality $(a + b)^j \leq 2^{j-1}(a^j + b^j)$, which is also a consequence of Jensen's inequality. We will clarify it.

Reference: Alquier 2024 User-friendly Introduction to PAC-Bayes Bounds

Dear reviewer, If not already, please take a look at the authors' rebuttal, and discuss if necessary. Thanks, -AC

We thank the reviewer for their feedback. We believe that we have addressed all the raised issues. We hope that the reviewer could reconsider their score if they believe their concerns are resolved.

Weakness 1 We emphasize that the contributions of Section 4 are more diverse than direct adaptations of existing results, let us develop this below.

The key contribution of our work is showing that these results, which were originally developed in the context of supervised learning, can be meaningfully applied to the analysis of diffusion models through the lens of score matching. This adaptation is enabled by our theoretical reduction (Theorem 3.1), which bridges the gap between generative modeling and the generalization of score functions. While the bounds themselves are not new, their application in the generative setting is, and we believe this is an important and nontrivial contribution. Hence, our results are original within this context.

To be more precise, in Theorem 3.1 (our main decomposition), we unveil the impact of three terms. The main theoretical contribution of this work regards the decomposition itself and our bound on the third term (the data-dependent constant), most of our theoretical efforts are focused on this part. Thus, while Theorems 4.1 and 4.2 rely on existing results, Section 4 brings up original results in two distinct ways.

Theorems 4.1 and 4.2 draw links between diffusion models and statistical learning. Our theorems should be thought as a continuation of the theoretical results of Section 3. Indeed, the conclusion of Section 3 is that we boiled down the problem of score approximation in diffusion models to a well-defined statistical learning problem (the generalization error of Theorem 3.1), where a loss function $\ell$ is well-defined. Then, the results of Section 4 exploits this setting to involve the benefits of statistical learning (i.e. the dynamics of the considered optimization algorithm) for score-based diffusion models.

Theorems 4.1 and 4.2 allow novel optimization-based correlation measures. Theorems 4.1 and 4.2 should also be considered as ways to obtain novel and practical correlation measures of the generalization ability of diffusion models. To our knowledge, it is the first time that the benefits of the optimization algorithm are involved as a measure of generalization of diffusion models. We insist that to highlight this strength, we computed practical experiments showing the relevance of these correlation measures on several datasets.

To sum up, while Theorems 4.1 and 4.2 exploit existing generalization bounds, they should be thought simultaneously as the final development of our theoretical contributions (expanded in Section 3) and the first step towards practical experiments, by unveiling original correlation measures that are tailored for diffusion models.

Weakness 2

As the reviewer mentions, our theory makes use of a finite dataset sampled i.i.d. from the data distribution (hence, a Dirac sampling from the data distribution). This is a classical setting in learning theory, as in practice, models are trained on a finite dataset. As it can be seen in Theorem 3.1, this empirical distribution of the dataset plays an important role in every term of the proposed decomposition of the score approximation error, as the reviewer suggested. Let us make it more clear by providing a more intuitive explanation of these terms.

The first term (empirical explicit score matching loss) corresponds to the quantity that is being minimized by the learning algorithm on the dataset, and it is the lower bound of the empirical denoising score-matching loss.

The third term represents (the data-dependent constant) a statistical bias that only depends on the data distribution and the dataset, i.e., not on the score network. Therefore, the existence of this term is a consequence of the Dirac sampling of the data distribution. It is a characteristic of the learning problem and providing an upper bound on this term is another main contribution of this work.

Finally, the remaining term (second term) is the generalization error associated with the denoising loss. Again, this term is non-trivial because of the presence of the empirical distribution of the dataset in the optimization process. This allows us to exploit the rich literature on generalization bounds and to apply, for the first time, classical generalization bounds to diffusion models. Note that we used as our examples existing bounds for SGLD and ADAM, but other generalization bounds would be applicable to our framework too, hence, highlighting its generality.

In the final version of the paper, we will add more details on the interpretation of these terms, as asked by the reviewer.

Dear reviewer, If not already, please take a look at the authors' rebuttal, and discuss if necessary. Thanks, -AC

We thank the reviewer for their insightful feedback. We address below the raised concerns and questions.

Question 1

We thank the reviewer for raising this important point. Indeed, in our work, we assume that the data distribution has compact support since it simplifies several technical derivations. As the reviewer suggests, we expect that it is possible to weaken this assumption. Indeed, the compact support assumption is mainly used for two reasons in our proof: (i) to apply Hoeffding inequality in the proof of Lemma 3.1 and (ii) to obtain estimates of the moments of the norm of the score function. In both cases, it should be enough, at the cost of additional derivations, to only make assumptions on the growth of the moments of the data distribution, e.g., that it is sub-Gaussian. In the final version of the paper, we will add a discussion about this aspect.

Question 2

The reviewer is right to notice that Figure 3 support the findings of [MWZZ] and others, in the context of generative modeling (notice that these bounds were developed for supervised learning), hence it serves as a sanity check to illustrate that the bounds are still meaningful in the generative modeling context.

We would like to underline that our work is the first, to our knowledge, to provide a sound theoretical framework to exploit these previous results (crafting the generalization error of Theorem 3.1 is not obvious) and exploit them in concrete experiments for diffusion models. Beyond this original connection, we believe that Figure 3 is necessary to assess that these generalization bounds have empirical relevance when applied to score-based generative models. Indeed, such experiments has not been done in the existing literature.

Question 3 and Weaknesses

We thank the reviewer for assessing the novelty of our proposed decomposition. The key novelty of Section 4 is not in deriving entirely new bounds from scratch, but in showing that existing generalization results—originally developed for supervised learning—can be repurposed to analyze generative models. This repurposing is made possible by the new connection we establish in Section 3, which reduces the analysis of generation algorithms to the score generalization term $G_\lambda$. To our knowledge, this is the first time such a transfer of theoretical tools from supervised learning to generative modeling has been made explicit. We believe this is a substantial contribution, as it opens up a rich and previously untapped body of results for use in the generative setting.

Thus, while Theorems 4.1 and 4.2 rely on existing results, Section 4 brings up original results in two distinct ways.

Theorems 4.1 and 4.2 draw links between DMs and statistical learning. Our theorems should be thought as a continuation of the theoretical results of Section 3. Indeed, the conclusion of Section 3 is that we boiled down the problem of generalisation in DMs to a well-defined statistical learning problem (the generalization error of Theorem 3.1), where a loss function $\ell$ is well-defined. Then, the results of Section 4 exploits this setting to involve the benefits of statistical learning (i.e, the dynamics of the considered optimization algorithm) for score-based diffusion models.

Theorems 4.1 and 4.2 allow novel, optimization-based, correlation measures. Theorems 4.1 and 4.2 should also be considered as ways to obtain novel and practical correlation measures of the generalization ability of diffusion models. To our knowledge, it is the first time that the benefits of the optimization algorithm are involved as a measure of generalization of diffusion models. To highlight this strength, we computed practical experiments showing the relevance of these correlation measures on several datasets.

To sum up, while Theorems 4.1 and 4.2 exploit existing generalization bounds, they should be thought simultaneously as the final development of our theoretical contributions (expanded in Section 3) and the first step towards practical experiments, by unveiling original diffusion models-tailored correlation measures.