Title: On the Generalization Properties of Diffusion Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: diffusion models

training dynamics

generalization gap

modes shift
Abstract: Diffusion models are a class of generative models that serve to establish a stochastic transport map between an empirically observed, yet unknown, target distribution and a known prior. Despite their remarkable success in real-world applications, a theoretical understanding of their generalization capabilities remains underdeveloped. This work embarks on a comprehensive theoretical exploration of the generalization attributes of diffusion models. We establish the theoretical estimates of the generalization gap that evolves in tandem with the training dynamics of score-based diffusion models, suggesting a polynomially small generalization error ($O(n^{-2/5}+m^{-4/5})$) on both the sample size $n$ and the model capacity $m$, evading the curse of dimensionality (i.e., independent of the data dimension) when *early-stopped*. Furthermore, we extend our quantitative analysis to a *data-dependent* scenario, wherein target distributions are portrayed as a succession of densities with progressively increasing distances between modes. This precisely elucidates the *adverse* effect of "*modes shift*'' in ground truths on the model generalization. Furthermore, these estimates are not solely theoretical constructs but have also been confirmed through numerical simulations. Our findings contribute to the rigorous understanding of diffusion models' generalization properties and provide insights that may guide practical applications.

On the Generalization Properties of Diffusion Models

Abstract

Diffusion models are a class of generative models that serve to establish a stochastic transport map between an empirically observed, yet unknown, target distribution and a known prior. Despite their remarkable success in real-world applications, a theoretical understanding of their generalization capabilities remains underdeveloped. This work embarks on a comprehensive theoretical exploration of the generalization attributes of diffusion models. We establish theoretical estimates of the generalization gap that evolves in tandem with the training dynamics of score-based diffusion models, suggesting a polynomially small generalization error ( O ( n -2 / 5 + m -4 / 5 ) ) on both the sample size n and the model capacity m , evading the curse of dimensionality (i.e., not exponentially large in the data dimension) when early-stopped . Furthermore, we extend our quantitative analysis to a datadependent scenario, wherein target distributions are portrayed as a succession of densities with progressively increasing distances between modes. This precisely elucidates the adverse effect of ' modes shift ' in ground truths on the model generalization. Moreover, these estimates are not solely theoretical constructs but have also been confirmed through numerical simulations. Our findings contribute to the rigorous understanding of diffusion models' generalization properties and provide insights that may guide practical applications.

1 Introduction

As an emerging family of deep generative models, diffusion models (DMs; [16, 45]) have experienced a surge in popularity, owing to their unparalleled performance in a wide range of applications ([22, 39, 68, 15, 3, 35, 4, 60, 33, 69, 64, 63, 53]). This has led to notable commercial successes, such as DALL·E ([34]), Imagen ([38]), and Stable Diffusion ([36]). Mathematically, diffusion models learn an unknown underlying distribution through a two-stage process: (i) first, successively and gradually injecting random noises (forward process); (ii) then reversing the forward process through denoising for sampling purposes (reverse process). To achieve this, an equivalent formulation of diffusion models called score-based generative models (SGMs; [50, 52]) is employed. SGMs implement the aforementioned two-stage process via the continuous dynamics represented by a joint group of coupled stochastic differential equations (SDEs) ([54, 20]).

Despite their impressive empirical performance, the theoretical foundation of DMs/SGMs remains underexplored. Generally, fundamental theoretical questions can be categorized into several aspects. By considering machine learning models as mathematical function classes from certain spaces, one can identify three central aspects: approximation, optimization and generalization. At the forefront lies the generalization problem, which aims to characterize the learning error between the learned and ground truth distributions.

The development of generalization theory for diffusion models is pressing due to both theoretical and practical concerns:

In theory, the generalization issues of generative modeling (or learning for distributions) may exhibit as the memorization phenomenon, if the modeled distribution is eventually trained to converge to the empirical distribution only associated with training samples. Intuitively, memorization arises from two reasons: (i) it is useful for the hypothesis space to be large enough to approximate highly complex underlying target distributions (universal convergence; [65]); (ii) the underlying distribution is unknown in practice, and one can only use a dataset with finite samples drawn from the target distribution. Rigorous mathematical characterizations of memorization are developed for bias potential models and GANs in [66] and [67], respectively. A natural question is, does a similar phenomenon occur for diffusion models? To answer this, a thorough investigation of generalization properties for DMs/SGMs is required.

In practice, the generalization capability of diffusion models is also an essential requirement, as the memorization can lead to potential privacy and copyright risks when models are deployed. Similar to other generative models and large language models (LLMs) [5, 70, 19, 6], diffusion models can also memorize and leak training samples [5, 46], hence can be subsequently attacked using specific procedures and algorithms [28, 18, 62]. Although there are defense methods developed to meet privacy and copyright standards ([11, 14, 58]), these approaches are often heuristic, without providing sufficient quantitative understandings particularly on diffusion models. Therefore, a comprehensive investigation of the generalization foundation of diffusion models, including both theoretical and empirical aspects, is of utmost importance in improving principled tutorial guidance in practice.

The current work develops the generalization theory of diffusion models in a mathematically rigorous manner. Our main results include the following:

We derive an upper bound of the generalization gap for diffusion models along the training dynamics. This result suggests, with early-stopping , the generalization error of diffusion models scales polynomially small on the sample size ( O ( n -2 / 5 ) ) and the model capacity ( O ( m -4 / 5 ) ). Notably, the generalization error also escapes from the curse of dimensionality.

This 'uniform' bound is further extended to a data-dependent setting, where a sequence of unidimensional Gaussian mixtures distributions with an increasing modes' distance is considered as the ground truth. This result characterizes the effect of ' modes shift ' quantitatively, which implies that the generalization capability of diffusion models is adversely affected by the distance between high-density regions of target distributions.

The theoretical findings are numerically verified in simulations.

The rest of this paper is organized as follows. In Section 2, we discuss the related work on the convergence and training fronts of diffusion models and also the generalization aspects of other generative modeling methods. Section 3 is the central part, which includes the problem formulation, main results, and consequences. Section 4 includes numerical verifications on synthetic and real-world datasets 3 . All the details of proofs and experiments are found in the appendices.

2 Related Work

We review the related work on diffusion models concerning the central results in this paper.

First, on the convergence theory, [7, 9] established elaborate error estimates between the modeled and target distribution given the discretization and time-dependent score matching

α(0)

c(T)

αc(0)

dc =f(α,t) - g²(t)Vα log Pt(α)|dt + g(t)dWt

α(T)

st,0(x) :=

Ao(W+Ue(t),0=A

1

m

Target: Finitely-supported prob.& Gaussian mixtures

Notations:

t:SDE time

T:maximal SDE time

Loss: Time-dependent score matching (Eq. (7))

PT ≥ π: a known prior

Algorithm:Gradientflow

T: training time

tolerance. Compared to the present work, they did not evolve the concrete training dynamics since the setting therein focuses on the properties of optimizers.

Second, on the training front, [51] proposed a set of techniques to enhance the training performance of score-based generative models, scaling diffusion models to images of higher resolution, but without any characterization of possible generalization improvements. Similar to [7, 9], [49] also provided error estimates between the modeled and target distributions in a point-wise sense and again did not evolve the detailed training dynamics.

Third, on the generalization and memorization side, corresponding theories are developed for bias potential models and GANs in [66] and [67], respectively, where the modeled distribution learns the ground truth with early-stopping and diverges or converges to the empirical distribution only associated with training samples after sufficiently long training time. The current work extends the mathematical analysis to the case of diffusion models under a data-dependent setting.

As a supplement, we also discuss related literature regarding low-density learning. [50] illustrated the difficulty of learning from low-density regions with toy formulations and simulations, which motivates the sampling method of annealed Langevin dynamics as the predecessor of (score-based) diffusion models. [21] restudied similar problems under the pure score matching regime (without the denoising or time-dependent dynamics) and attributed the difficulty to increasing isoperimetry of distributions with modes shift. [41] defined the Hardness score and numerically justified that decreasing manifold densities leads to the increasing Hardness score, and consequently applied the Hardness score as the regularization to the sampling process to enhance synthetic images from low-density regions. As a comparison, this work establishes a mathematically rigorous estimate on the generalization gap that quantitatively depends on the range of low-density regions (between modes) in target distributions for (score-based) diffusion models that requires denoising .

3 Formulation and Results

In this section, we first introduce the problem setup. Next, we state the main theoretical results, subsequent consequences, and possible connections. The numerical illustration is provided at last.

3.1 Problem Formulation

Since DMs/SGMs have already grown into a large family of generative models with an enormous number of variants, there are various ways to define the parameterization of diffusion models. Here, we adopt (one of) the most fundamental architectures proposed in [54], where the forward perturbation and reverse sampling process are both implemented by a joint group of coupled (stochastic) differential equations. See Figure 1 for an illustration of the problem formulation.

Forward perturbation. We start with the setting of unsupervised learning. Given an unlabeled dataset D x = { x i } n i =1 ⊂ R d with the sample x i i.i.d. ∼ p 0 ( x ) , where p 0 denotes the underlying (ground truth or target) distribution, the forward diffusion process is defined as

Here, the drift coefficient f ( · , t ) : R d ↦→ R d is a time-dependent vector-valued function, and the diffusion coefficient g ( · ) : R ≥ 0 ↦→ R is a scalar function, and W t denotes the standard Wiener process (a.k.a., Brownian motion). The SDE (1) has a unique, strong solution under certain regularity conditions (i.e., globally Lipschitz coefficients in both state and time; see [31]). From now on, we denote by p t ( x ( t )) the marginal distribution of x ( t ) , and let p t | s ( x ( t ) | x ( s )) be the (perturbation) transition kernel from x ( s ) to x ( t ) , 0 ≤ s < t ≤ T < ∞ with T as the time horizon. By appropriately selecting f and g , one can force the SDE (1) to converge to a prior distribution (typically a Gaussian). Common examples include the (time-rescaled) Ornstein-Uhlenbeck (OU) process, 4 which is a special case of the linear version of (1):

Reverse sampling. According to [1] and [54], both the following reverse-time SDE and probability flow ODE share the same marginal distribution as the forward-time SDE (1):

where ¯ W t is a standard Wiener process when time flows backwards from T to 0 , and dt is an infinitesimal negative time step. With the initial condition x ( T ) ∼ p T ≈ π , where π is a known prior distribution such as the Gaussian noise, one can (numerically) solve (3) or (4) to transform noises into samples from p 0 , which is exactly the goal of generative modeling.

Loss objectives. The only remaining task is to estimate the unknown (Stein) score function ∇ x log p t ( x ) . This is achieved by minimizing the following weighted sum of denoising score matching ([57]) objectives:

with θ ∗ := arg min θ L ( θ ; λ ( · )) , where U (0 , T ) denotes the uniform distribution over [0 , T ] , and λ ( t ) : [0 , T ] ↦→ R + is a weighting function, which is typically selected as

according to ([54]). The score function s t, θ : R d ↦→ R d is time-dependent and can be parameterized as a neural network (encoded with the time information) such as the U-net([37]) architecture commonly applied in the field of image segmentation. Alternatively, one can also define the timedependent score matching loss

which is equivalent to (5) up to a constant independent of θ by [57, 49].

In practice, expectations in the objective (5) can be respectively estimated with empirical means over time steps in [0 , T ] , data samples from p 0 and p t | 0 , which is efficient when the drift coefficient f ( · , t ) is linear. Specifically, if the forward-time SDE takes the form of (2), the transition kernel p t | 0 has a closed form ([40])

where N ( x ; µ , Σ ) denotes the multivariate Gaussian distribution evaluated at x with the expectation µ and covariance Σ , and r ( t ) := e ∫ t 0 f ( ζ ) dζ , v ( t ) := √ ∫ t 0 g 2 ( ζ ) r 2 ( ζ ) dζ .

Training. We aim to investigate the gradient flow training dynamics over the empirical loss

where ˆ L n is the Monte-Carlo estimation of L defined in (5) on the training dataset, 5 with an auxiliary gradient flow over the population loss

In both cases, the weighting function λ ( · ) is selected as in (6). Denote the score function learned at the training time τ evaluated at the SDE time t with respect to the empirical loss and population loss as s t, ˆ θ n ( τ ) ( x ( t )) and s t, θ ( τ ) ( x ( t )) , respectively. The corresponding density functions, denoted by p t, ˆ θ n ( τ ) ( x ( t )) and p t, θ ( τ ) ( x ( t )) , are obtained by solving

Score networks. We parameterize the score function s t, θ as the following random feature model

where σ is the ReLU activation function, A = ( a 1 , . . . , a m ) ∈ R d × m is the trainable parameter, while W = ( w 1 , . . . , w m ) ⊤ ∈ R m × d and U = ( u 1 , . . . , u m ) ⊤ ∈ R m × d e are randomly initialized and frozen during training, and e : R ≥ 0 ↦→ R d e is the embedding function concerning the time information. Assume that a i , w i and u i are i.i.d. sampled from an underlying distribution ρ . Then, as m →∞ , we get

with a ( w , u ) := 1 ρ 0 ( w , u ) ∫ R d a ρ ( a , w , u ) d a and ρ 0 ( w , u ) := ∫ R d ρ ( a , w , u ) d a . By the positive homogeneity property of the ReLU activation, we can assume that ∥ w ∥ 1 + ∥ u ∥ 1 ≤ 1 w.l.o.g.

One can view ¯ s t, ¯ θ ( x ) as a continuous version of the random feature model. Correspondingly, the optimal solution is denoted as ¯ θ ∗ when replacing the parameterized score function s t, θ ( x ) in the loss objective (5) or (7) by ¯ s t, ¯ θ ( x ) . Define the kernel k ρ 0 ( x , x ′ ) := E ( w , u ) ∼ ρ 0 [ σ ( w ⊤ x + u ⊤ e ( t )) σ ( w ⊤ x ′ + u ⊤ e ( t )) ] , and let H k ρ 0 be the induced reproducing kernel Hilbert space (RKHS; [2]), we have ¯ s t, ¯ θ ∈ H k ρ 0 if the RKHS norm ∥ ∥ ¯ s t, ¯ θ ∥ ∥ 2 H kρ :=

H kρ 0 m F m i =1

0 E ( w , u ) ∼ ρ 0 [ ∥ a ( w , u ) ∥ 2 2 ] = ∥∥ a ∥ 2 ∥ 2 L 2 ( ρ 0 ) < ∞ , and the corresponding discrete version can be defined by the empirical average, i.e., ∥ s t, θ ∥ 2 := 1 ∥ A ∥ 2 = 1 ∑ m ∥ a ( w i , u i ) ∥ 2 2 .

Remark 1. There are more modern and complex mathematical tools such as neural tangent kernels (NTKs) and mean fields that can be selected as the score networks. Employing these modern tools is valuable at least for theoretical completeness and we leave these as the future work. 6

The goal is to measure and bound the generalization error evolving with the gradient flow training dynamics (9) between the learned distribution and target distribution, using the common Kullback-Leibler (KL) divergence.

Definition 1 (KL divergence) . Given two distributions p and q , the KL divergence from q to p is defined as D KL ( p || q ) = ∫ R d p ( x ) log ( p ( x ) q ( x ) ) d x .

Based on the above definitions, the generalization gap along the gradient flow training dynamics (9) is formulated as D KL ( p 0 ∥ p 0 , ˆ θ n ( τ ) ) , which is only a function of the training time τ . The goal is to estimate D KL ( p 0 ∥ p 0 , ˆ θ n ( τ ) ) .

3.2 Main Results

In this section, we state the main results of the generalization capability of the (score-based) diffusion models and how it evolves as the training proceeds. Based on the formulation in Section 3.1, we theoretically derive several upper bounds to estimate D KL ( p 0 ∥ p 0 , ˆ θ n ( τ ) ) under different settings.

The results cover both the positive and negative aspects: in the data-independent setting where the target distribution has finite support, the generalization gap is proved to be small; while in the data-dependent setting where the target distribution possesses shift modes, the generalization is adversely affected by the modes' distance.

3.2.1 Data-Independent Generalization Gap

In this section, we provide the characterization of the generalization capability for diffusion models given a target distribution defined on a finite domain. Generally, the KL divergence from the learned distribution p 0 , ˆ θ n ( τ ) at the training time τ to the target distribution p 0 can be estimated as follows.

Theorem 1. Suppose that the target distribution p 0 is continuously differentiable and has a compact support set, i.e., || x || ∞ is uniformly bounded, and there exists a reproducing kernel Hilbert space (RKHS) H (:= H k ρ 0 ) such that ¯ s 0 , ¯ θ ∗ ∈ H . Assume that the initial loss, trainable parameters, the embedding function e ( t ) and weighting function λ ( t ) are all bounded. Then for any δ > 0 , δ ≪ 1 , with the probability of at least 1 -δ , we have

where ≲ hides the term d log( d +1) , the polynomials of log(1 /δ 2 ) , finite RKHS norms and universal positive constants only depending on T .

Remark 2. Since p 0 is compactly supported, the target score function s 0 ( x ) = ∇ x log p 0 ( x ) is also defined on a compact domain. According to [12, 13], s 0 is contained in the Barron function space with a finite Barron norm, and hence in a certain RKHS with a finite RKHS norm. Therefore, it is reasonable to require that the global minimizer s 0 , θ ∗ or ¯ s 0 , ¯ θ ∗ is also contained in some RKHS.

Proof sketch. Theorem 1 is proved via the following procedure.

According to Theorem 1 in [49], the KL divergence on the left-hand side can be upper bounded by the population loss of the trained model up to a small error. That is,

We use the model trained with respect to the population loss (10) to perform the decomposition:

where we omit the weighting function g 2 ( · ) for simplicity. Here, ¯ ˜ L is the loss objective obtained by replacing s t, θ in ˜ L (defined in (7)) by ¯ s t, ¯ θ , and I 3 summarizes the resulting Monte Carlo error.

I 3 can be estimated via a similar argument as in [24] (Lemma 48).

I 2 can be upper bounded via a standard analysis on the gradient flow dynamics over convex objectives.

I 1 can be reduced as the norm product of s 0 , ˆ θ n ( τ ) , s 0 , θ ( τ ) and their gap, then

(a) the former can be bounded with a square-root rate growth via a general norm estimate of parameters trained under the gradient flow dynamics;

(b) the latter can be estimated by the Rademacher complexity (see e.g., Chapter 26 in [43]).

Combining all above gives the desired result. The detailed proof is found in Appendix A.1.

Discussion on error bounds. The three error terms are further analyzed as follows.

The first term is the main error, which implies an early-stopping generalization gap. In fact, if one selects an early-stopping time τ es as τ es = Θ ( n 2 5 ) , and let m ∼ n , we have

The second term is m -dependent and corresponds to the approximation error, which is o (1) when m ≫ 1 . In fact, the random feature model is a universal approximator to Lipschitz continuous functions on a compact domain (Theorem 6 in [17]). See more details in Appendix A.1 (the last paragraph).

The third term is exponentially small in T since π (e.g. the Gaussian density) is log-Sobolev, according to a classical result in e.g. [55] (Theorem 3.20, Theorem 3.24 and Remark 3.26).

Remark 3. In practice, it is common to use the test error to evaluate the generalization performance. For diffusion models, a straightforward approach is to compute the negative log-likelihood (averaged in bits/dim; equivalent to the KL divergence) on the test dataset during training with the instantaneous change-of-variable formula ([8]) and probability flow ODE (defined in (4)), where the true score function ∇ x log p t ( x ) is replaced by s t, θ ( τ ) ( x ) .

Remark 4. Previous literature has established similar bounds for bias potential models ([66]) and GANs ([67]). Theorem 1 extends the corresponding results to the setting of diffusion models. Furthermore, this upper bound is finer in the sense that it incorporates the information regarding the model capacity (the hidden dimension m in this case), which shows that more parameters benefit the learning and generalization, as expected.

3.2.2 Data-Dependent Generalization Gap

In Section 3.2.1, we derive estimates on the generalization error for diffusion models along the training dynamics, where the target distribution is assumed to be finitely supported. In reality, this is often not the case, where target distributions usually possess distant multi-modes, from simple Gaussian mixtures to complicated Boltzmann distributions of physical systems ([29, 30]). Under these settings, the above analysis in Section 3.2.1 can not directly apply since the data domain is unbounded. It remains a problem to quantitatively characterize the generalization behavior of diffusion models given these target distributions with distant multi-modes or modes shift.

To provide a fine-grained demonstration of the generalization capability of diffusion models when applied to learn distributions with distant multi-modes, as an illustrating example, the Gaussian mixture with two modes is selected as the target distribution.

Theorem 2. Suppose the target distribution p 0 is a one-dimensional 2-mode Gaussian mixture: p 0 ( x ) = q 1 N ( x ; -µ, 1) + q 2 N ( x ; µ, 1) , where µ > √ log(1 /δ 2 ) , q 1 , q 2 > 0 with q 1 + q 2 = 1 are all constants. Under the conditions of Theorem 1 (except the uniform boundness of inputs), we have

where τ ≥ 1 , ≲ hides the polynomials of log(1 /δ 2 ) , finite RKHS norms and universal positive constants only depending on T .

Proof sketch. Theorem 2 is proved following a similar procedure with Theorem 1, except the input data x does not have a uniform bound here. This problem mainly affects the last step (5 (b)) in the proof sketch of Theorem 1, and can be handled by using the fact

| x | ∈ [ µ -√ log(1 /δ 2 ) , µ + √ log(1 /δ 2 )] = Θ( µ ) (17) given the target Gaussian mixture distribution. The detailed proof is found in Appendix A.2.

Remark 5. Theorem 2 indicates that, even for a simple target distribution (e.g. a one-dimensional 2-mode Gaussian mixture), the generalization error of diffusion models can be polynomially large regarding the modes' distance. Although Theorem 2 provides only an upper bound, the modes shift effect holds due to the model-target inconsistency (the last paragraph in Appendix A.2) and the following consistent experiments (Section 4.1.2).

4 Numerical Verifications

In this section, we numerically verify the previous theoretical results and insights (early-stopping generalization and modes shift effect) on both synthetic datasets and real-world datasets.

4.1 Simulations on Synthetic Datasets

4.1.1 Early-Stopping Generalization

First, we illustrate the early-stopping generalization gap established in Theorem 1. We select the one-hidden-layer neural network with Swish activations as the score network, which is trained using the SGD optimizer with a fixed learning rate 0 . 5 . The target distribution is set to be a onedimensional 2-mode Gaussian mixture with the modes' distance equalling 6, and the number of data samples is 1000. We measure the KL divergence from the trained model to the target distribution D KL ( p 0 ∥ p 0 , ˆ θ n ( τ ) ) along with the training epochs.

From Figure 3, one can observe that the KL divergence achieves its minimum at approximately the 800 -th training epoch, and it starts to increase after this turning point. The experimental results are consistent with Theorem 1 (over multiple runs), which states that there exist early-stopping times when diffusion models can generalize well, indicating the effectiveness of our upper bound. Further, the KL divergence begins to oscillate after the minimum point, which may suggest a phase transition in the training dynamics, and the transition point is around the (optimal) early-stopping time.

1.4

Repetition1

Repetition2

1.2

Repetition 3

1.0

KL divergence

0.8

0.6

0.4

0.2

0

250

500

750

1000

1250

1500

1750

4.1.2 Modes Shift Effect

Next, we numerically test the relationship between the modes' distance and generalization (density estimation) performance. All the configurations remain the same as Section 4.1.1, except that the target Gaussian mixtures have different modes' distances.

In Figure 4, the modeled distributions exhibit the following two-stage dynamics: (i) first gradually fitting the two modes (epoch = 100 → 1000 ); (ii) then diverging (epoch = 1000 → 1900 ). This aligns with the KL divergence dynamics (Figure 3) and again verifies the corresponding theoretical results (Theorem 1). However, Figure 5 shows that when the modes are distant from each other, there is difficulty in the learning process. In Figure 5, the optimal generalization is achieved at epoch = 100 , but is still far from well generalizing. As the training proceeds, the learned model is almost always a single-mode distribution (epoch = 1000 , 1900 ). This phenomenon is aligned with the results established in Theorem 2, which states that when there are distant modes in the target distribution, the generalization performance is relatively poor.

0.175

0.150

0.125

epoch = 100

Trained

0.200

0.175

0.150

0.125

epoch = 1000

Target

Trained

0.175

0.200

0.150

0.125

epoch = 1900

Target

Trained

Aursuaa

0.100

0.100

0.100

0.075

0.075

0.075

0.050

0.050

0.050

0.025

0.025

0.025

0.000

10.07.55.0

2.5

0.0

2.5

5.0

7.5

10.0

0.000

10.0

7.5

5.0

2.5

0.0

2.5

5.0

7.5

10.0

0.000

10.07.55.0

2.5

0.0

2.5

5.0

7.5

10.0

epoch = 100

epoch = 1000

epoch = 1900

0.25

Target

Trained

0.200

Trained

Target

0.30

Target

Trained

0.175

0.20

0.150

0.25

Density

0.15

KqIsu

0.125

Aisuaa

0.20

0.100

0.15

0.10

0.075

0.10

0.050

0.05

0.025

0.05

00°0

OE-

20

10

0

10

20

30

0.000

=30

20

-10

0

10

20

30

0.00

30

20

-10

0

10

20

30

4.2 Simulations on Real-World Datasets

In this subsection, we verify our results on the MNIST dataset using the standard U-net architecture as the score network, which suggests that the adverse effect of modes shift on the generalization performance of diffusion models also appears in general .

̸

The setup is as follows. First, we perform a K -means clustering on D ( D denote the MNIST dataset) to get D = ⋃ K k =1 D k , and ¯ x k as the center of D k , k = 1 , 2 , · · · , K . Let ( i ∗ , j ∗ ) := arg max i = j ∥ ¯ x i -¯ x j ∥ 2 , and D farthest := D i ∗ ⋃ D j ∗ . D nearest is similarly constructed by arg min indices. Then, by randomly selecting the same number of data samples and using the same configuration, we train two separate diffusion models on D farthest and D nearest, respectively, and then perform inference (sampling). The training loss curves and sampling results are shown in Figure 6 and Figure 7, respectively. One can observe a significant performance gap: the diffusion model trained on D farthest appears a higher learning loss and worse sampling quality compared to those of D nearest .

180

Nearest

160

Farthest

140

Training loss

120

100

80

60

40

20

0

200

400

600

800

1000

Epoch

2

8

2

8

3

S

3

3

5

C

3

8

3

9

3

8

2

8

3

3

3

3

2

8

8

3

3

2

2

2

8

3

3

3

8

3

5

3

2

2

S

2

了

3

3S283

2

4.3 Discussion

We compare the results developed in this work with former corresponding literature as follows:

The previous work [21] also studied the adverse effect of modes shift, which particularly reported a contrastive simulation indicating the degraded performance when modeling Gaussian mixtures with the increasing distance between modes (see Figure 2 in [21] and compare with Figure 4 and Figure 5). However, the results therein are established and tested under the 'pure' score matching setting, without the denoising or time-dependent dynamics. As a comparison, Theorem 2 establishes a theoretical estimate on the generalization gap for diffusion models that requires denoising , and this upper bound directly depends on the distance between modes of target distributions, instead of a circuitous characterization in [21] to attribute the difficulty of learning modes shift to increased isoperimetry of corresponding target distributions.

Similar adverse effect of modes shift has also been theoretically analyzed and numerically verified on recurrent neural networks (RNNs), see e.g., [23, 24]. There, the modes shift is understood as a type of long-term memory. This is the phenomenon of the 'curse of memory': When there is long-term memory in the target, it requires a large number of parameters for the approximation. Meanwhile, the training process will suffer from severe slowdowns. Both of these effects can be exponentially more pronounced with increasing memory or modes shift.

The very recent work [42] also considered the problem of learning Gaussian mixtures using the denoising diffusion probabilistic model (DDPM) objective, but under a teacher-student setting. That is, given the Gaussian mixtures target, [42] parametrized the score network model in the same form of the score function target, with the goal to identify true parameters. Consequently, there are all positive convergence results developed in [42], despite that the time and sample complexity increases with the distance between modes. As a comparison, Theorem 2 adopts a pre-selected score network model without incorporating any information from the ground truth, and hence establish negative results concerning the modes shift. In fact, if the goal is to identify only true positions of the target Gaussian mixture (this is exactly the setting of [42]), the teacher-student setup seems not necessary (see Figure 5, where the true position is also learned efficiently using the one-hidden-layer Swish neural network as the score network, but the modes are always weighed incorrectly). In addition, [42] did not take the whole denoising dynamics into account. That is, the gradient descent (GD) analysis therein was performed on the denoising score matching objective successively at only two time stages: a larger t 1 ('high noise') and a smaller t 2 ('low noise'), which is often not the case in practice.

5 Conclusion

In this paper, we provide a theoretical analysis of the fundamental generalization aspect of training diffusion models under both the data-independent and data-dependent settings, and early-stopping estimates of the generalization gap along the training dynamics are derived. Quantitatively, the data-independent results indicate a polynomially small generalization error that escapes from the curse of dimensionality, while the data-dependent results suggest the adverse effect of modes shift in target distributions. Numerical simulations have illustrated and verified these theoretical analyses. This work forms a basic starting point for understanding the intricacies of modern deep generative modeling and corresponding central concerns such as memorization, privacy, and copyright arising from practical applications and business products. More broadly, the approach here may have the potential to be extended to other variants in the diffusion models family, including a general SDEbased design space ([20]), consistency models ([48]), rectified flows ([27, 26]), Schrödinger bridges ([59, 56, 10, 44, 47, 25]), etc. These are certainly worthy of future exploration.

References

Brian D. O. Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications , 12(3):313-326, 1982.

Nachman Aronszajn. Theory of reproducing kernels. Transactions of the American Mathematical Society , 68(3):337-404, 1950.

Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems , 34:17981-17993, 2021.

Omri Avrahami, Dani Lischinski, and Ohad Fried. Blended diffusion for text-driven editing of natural images. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 18208-18218, 2022.

Nicholas Carlini, Jamie Hayes, Milad Nasr, Matthew Jagielski, Vikash Sehwag, Florian Tramer, Borja Balle, Daphne Ippolito, and Eric Wallace. Extracting training data from diffusion models. In USENIX Security Symposium , pages 5253-5270, 2023.

Nicholas Carlini, Daphne Ippolito, Matthew Jagielski, Katherine Lee, Florian Tramer, and Chiyuan Zhang. Quantifying memorization across neural language models. In International Conference on Learning Representations , 2023.

Hongrui Chen, Holden Lee, and Jianfeng Lu. Improved analysis of score-based generative modeling: User-friendly bounds under minimal smoothness assumptions. International Conference on Machine Learning , 202:4735-4763, 2023.

Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K. Duvenaud. Neural ordinary differential equations. Advances in Neural Information Processing Systems , 31:6571-6583, 2018.

Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru Zhang. Sampling is as easy as learning the score: Theory for diffusion models with minimal data assumptions. In International Conference on Learning Representations , 2023.

Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion Schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems , 34:17695-17709, 2021.

Tim Dockhorn, Tianshi Cao, Arash Vahdat, and Karsten Kreis. Differentially private diffusion models. Transactions on Machine Learning Research , 2023.

Weinan E, Chao Ma, and Lei Wu. A priori estimates of the population risk for two-layer neural networks. Communications in Mathematical Sciences , 17(5):1407-1425, 2019.

Weinan E, Chao Ma, and Lei Wu. The Barron space and the flow-induced function spaces for neural network models. Constructive Approximation , 55(1):369-406, 2022.

Sahra Ghalebikesabi, Leonard Berrada, Sven Gowal, Ira Ktena, Robert Stanforth, Jamie Hayes, Soham De, Samuel L Smith, Olivia Wiles, and Borja Balle. Differentially private diffusion models generate useful synthetic images. In International Workshop on Trustworthy Federated Learning in Conjunction with IJCAI , 2023.

Shansan Gong, Mukai Li, Jiangtao Feng, Zhiyong Wu, and Lingpeng Kong. Diffuseq: Sequence to sequence text generation with diffusion models. In International Conference on Learning Representations , 2023.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Daniel Hsu, Clayton H. Sanford, Rocco Servedio, and Emmanouil Vasileios VlatakisGkaragkounis. On the approximation power of two-layer networks of random ReLUs. Conference on Learning Theory , 134:2423-2461, 2021.

Hailong Hu and Jun Pang. Membership inference of diffusion models. arXiv preprint arXiv:2301.09956 , 2023.

Matthew Jagielski, Om Thakkar, Florian Tramer, Daphne Ippolito, Katherine Lee, Nicholas Carlini, Eric Wallace, Shuang Song, Abhradeep Guha Thakurta, Nicolas Papernot, and Chiyuan Zhang. Measuring forgetting of memorized training examples. In International Conference on Learning Representations , 2023.

Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. Advances in Neural Information Processing Systems , 35:26565-26577, 2022.

Frederic Koehler, Alexander Heckett, and Andrej Risteski. Statistical efficiency of score matching: The view from isoperimetry. In International Conference on Learning Representations , 2023.

Haoying Li, Yifan Yang, Meng Chang, Shiqi Chen, Huajun Feng, Zhihai Xu, Qi Li, and Yueting Chen. Srdiff: Single image super-resolution with diffusion probabilistic models. Neurocomputing , 479:47-59, 2022.

Zhong Li, Jiequn Han, Weinan E, and Qianxiao Li. On the curse of memory in recurrent neural networks: Approximation and optimization analysis. In International Conference on Learning Representations , 2021.

Zhong Li, Jiequn Han, Weinan E, and Qianxiao Li. Approximation and optimization theory for linear continuous-time recurrent neural networks. Journal of Machine Learning Research , 23(42):1-85, 2022.

Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos Theodorou, Weili Nie, and Anima Anandkumar. I 2 SB: Image-to-image Schrödinger bridge. International Conference on Machine Learning , 202:22042-22062, 2023.

Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209.14577 , 2022.

Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In International Conference on Learning Representations , 2023.

Tomoya Matsumoto, Takayuki Miura, and Naoto Yanai. Membership inference attacks against diffusion models. In IEEE Security and Privacy Workshops , pages 77-83, 2023.

Laurence Illing Midgley, Vincent Stimper, Gregor N. C. Simm, Bernhard Schölkopf, and José Miguel Hernández-Lobato. Flow annealed importance sampling bootstrap. In International Conference on Learning Representations , 2023.

Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science , 365(6457):eaaw1147, 2019.

Bernt Øksendal. Stochastic Differential Equations . Springer, 2003.

Jakiw Pidstrigach. Score-based generative models detect manifolds. Advances in Neural Information Processing Systems , 35:35852-35865, 2022.

Vadim Popov, Ivan Vovk, Vladimir Gogoryan, Tasnima Sadekova, and Mikhail Kudinov. Gradtts: A diffusion probabilistic model for text-to-speech. International Conference on Machine Learning , 139:8599-8608, 2021.

Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 , 2022.

Kashif Rasul, Calvin Seward, Ingmar Schuster, and Roland Vollgraf. Autoregressive denoising diffusion models for multivariate probabilistic time series forecasting. International Conference on Machine Learning , 139:8857-8868, 2021.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10684-10695, 2022.

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention , pages 234-241. Springer, 2015.

Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Raphael Gontijo-Lopes, Burcu Karagol Ayan, Tim Salimans, Jonathan Ho, David J. Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems , 35:36479-36494, 2022.

Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J. Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. IEEE Transactions on Pattern Analysis and Machine Intelligence , 45(4):4713-4726, 2023.

Simo Särkkä and Arno Solin. Applied Stochastic Differential Equations , volume 10. Cambridge University Press, 2019.

Vikash Sehwag, Caner Hazirbas, Albert Gordo, Firat Ozgenel, and Cristian Canton. Generating high fidelity data from low-density regions using diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11492-11501, 2022.

Kulin Shah, Sitan Chen, and Adam Klivans. Learning mixtures of gaussians using the DDPM objective. Advances in Neural Information Processing Systems , 36, 2023.

Shai Shalev-Shwartz and Shai Ben-David. Understanding Machine Learning: From Theory to Algorithms . Cambridge University Press, 2014.

Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion Schrödinger bridge matching. Advances in Neural Information Processing Systems , 36, 2023.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. International Conference on Machine Learning , 37:2256-2265, 2015.

Gowthami Somepalli, Vasu Singla, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Diffusion art or digital forgery? investigating data replication in diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 6048-6058, 2023.

Vignesh Ram Somnath, Matteo Pariset, Ya-Ping Hsieh, Maria Rodriguez Martinez, Andreas Krause, and Charlotte Bunne. Aligned diffusion Schrödinger bridges. Conference on Uncertainty in Artificial Intelligence , 216:1985-1995, 2023.

Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. International Conference on Machine Learning , 202:32211-32252, 2023.

Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in Neural Information Processing Systems , 34:14151428, 2021.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems , 32, 2019.

Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in Neural Information Processing Systems , 33:12438-12448, 2020.

Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. Conference on Uncertainty in Artificial Intelligence , 115:574-584, 2020.

Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with score-based generative models. In International Conference on Learning Representations , 2022.

Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations , 2021.

Ramon Van Handel. Probability in high dimension. Lecture Notes (Princeton University) , 2014.

Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving Schrödinger bridges via maximum likelihood. Entropy , 23(9):1134, 2021.

Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Computation , 23(7):1661-1674, 2011.

Nikhil Vyas, Sham M. Kakade, and Boaz Barak. On provable copyright protection for generative models. International Conference on Machine Learning , 202:35277-35299, 2023.

Gefei Wang, Yuling Jiao, Qian Xu, Yang Wang, and Can Yang. Deep generative learning via Schrödinger bridge. International Conference on Machine Learning , 139:10794-10804, 2021.

Jay Zhangjie Wu, Yixiao Ge, Xintao Wang, Stan Weixian Lei, Yuchao Gu, Yufei Shi, Wynne Hsu, Ying Shan, Xiaohu Qie, and Mike Zheng Shou. Tune-a-video: One-shot tuning of image diffusion models for text-to-video generation. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 7623-7633, 2023.

Lei Wu and Weijie J. Su. The implicit regularization of dynamical stability in stochastic gradient descent. International Conference on Machine Learning , 202:37656-37684, 2023.

Yixin Wu, Ning Yu, Zheng Li, Michael Backes, and Yang Zhang. Membership inference attacks against text-to-image generation models. arXiv preprint arXiv:2210.00968 , 2022.

Tian Xie, Xiang Fu, Octavian-Eugen Ganea, Regina Barzilay, and Tommi S. Jaakkola. Crystal diffusion variational autoencoder for periodic material generation. In International Conference on Learning Representations , 2022.

Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations , 2022.

Hongkang Yang. A mathematical framework for learning probability distributions. Journal of Machine Learning , 1(4):373-431, 2022.

Hongkang Yang and Weinan E. Generalization and memorization: The bias potential model. Mathematical and Scientific Machine Learning , 145:1013-1043, 2022.

Hongkang Yang and Weinan E. Generalization error of GAN from the discriminator's perspective. Research in the Mathematical Sciences , 9(1):1-31, 2022.

Ruihan Yang, Prakhar Srivastava, and Stephan Mandt. Diffusion probabilistic modeling for video generation. Entropy , 25(10):1469, 2023.

Jongmin Yoon, Sung Ju Hwang, and Juho Lee. Adversarial purification with score-based generative models. International Conference on Machine Learning , 139:12062-12072, 2021.

Chiyuan Zhang, Daphne Ippolito, Katherine Lee, Matthew Jagielski, Florian Tramèr, and Nicholas Carlini. Counterfactual memorization in neural language models. Advances in Neural Information Processing Systems , 36, 2023.

A Technical Results and Proofs

A.1 Data-Independent Generalization Gap

To derive the theorem for the generalization error of this score-based generative model, we first give the following lemmas.

Lemma 1 (Forward perturbation estimates) . Consider the forward diffusion process with the linear drift coefficient (2). For any δ > 0 , δ ≪ 1 , with the probability of at least 1 -δ , we have

where C T := max t ∈ [0 ,T ] { r ( t ) , r ( t ) v ( t ) } .

Proof. When the drift coefficient f ( · , t ) : R d ↦→ R d is linear to x , i.e., f ( x , t ) = f ( t ) x , the transition kernel p t | 0 has a closed form (8)

with r ( t ) := e ∫ t 0 f ( ζ ) dζ , v ( t ) := √ ∫ t 0 g 2 ( ζ ) r 2 ( ζ ) dζ . Hence, we have

For any ϵ ∼ N (0 , 1) , c > 1 , we have

Hence, for any δ ∈ (0 , 1) with δ ≪ 1 , with the probability of at least 1 -δ , we have

with C T := max t ∈ [0 ,T ] { r ( t ) , r ( t ) v ( t ) }

Lemma 2 (Theorem 1 in [49]) . We have

Lemma 3. Both of the loss objectives ˜ L ( θ ; λ ( · )) and ¯ ˜ L ( ¯ θ ; λ ( · )) are quadratic (and hence convex).

Proof. The convexity arises from the following points: (i) The score matching loss objectives ˜ L ( θ ; λ ( · )) (defined in (7)) and ¯ ˜ L ( ¯ θ ; λ ( · )) are L 2 -metrics between the score network model and target score function; (ii) The score networks are defined as random feature models (see (12) and (13)) that are linear to trainable parameters. Therefore, using some trace techniques and basic variational calculations, it is not hard to derive the fact that the loss objectives are quadratic and hence convex with respect to trainable parameters.

˜ 1

Since for any h , ¯ h , B , we have

we further get where

Here, B 1 is a positive semi-definite matrix, since v ⊤ B 1 v = E t E x ( t ) [( v ⊤ h 1 ( x ( t ) , t )) 2 ] ≥ 0 for any v . Notice that for any A , B ,

where ⊗ denotes the Kronecker product. Hence

is a quadratic function. It is straightforward to show that the eigenvalues of B 1 ⊗ I are the same as B 1 but with multiplicity, 7 implying that B 1 ⊗ I is also positive semi-definite. Therefore,

is positive semi-definite, i.e., the loss is convex with respect to trainable parameters.

(2) For ¯ ˜ L ( ¯ θ ; λ ( · )) , notice that

Then for any ϕ ∈ L 2 ( ρ 0 ) , by symmetry we have

where

This yields and

hence the functional ¯ ˜ L ( ¯ θ ; λ ( · )) is convex with respect to ¯ θ (given any positive weighting function λ ( · ) ). The proof is completed.

Since the weighting function is fixed in our analysis, we omit the notation λ ( · ) without ambiguity in the following contents.

Lemma 4. For any τ > 0 and θ , ¯ θ , we have

Proof. (1) For the loss objective ¯ ˜ L , we define the Lyapunov function

Then, we have

where the last inequality holds since

By convexity, for any τ 1 , τ 2 , it holds that

hence d dτ ¯ E ( τ ) ≤ 0 . We conclude that ¯ E ( τ ) ≤ ¯ E (0) , or equivalently

Therefore, note that ∥ ∥ ¯ s 0 , ¯ θ ∥ ∥ 2 H := E ρ 0 [ ∥ a ∥ 2 2 ] (let H := H k ρ 0 ), we obtain

which gives the desired estimate.

(2) For the loss objective ˜ L , the argument is almost the same, except replacing the Lyapunov function by

and ⟨· , ·⟩ L 2 ( ρ 0 ) by ⟨· , ·⟩ . Note that θ = vec ( A ) / √ m and ∥ s 0 , θ ∥ 2 H = ∥ A ∥ 2 F /m , we obtain

which completes the proof.

Lemma 5. Suppose that the loss objectives ˜ L , ˜ L ( n ) , ¯ ˜ L , ¯ ˜ L ( n ) are bounded at the initialization, then for any τ > 0 , we have

Proof. It's sufficient to prove ∥ ∥ ∥ ¯ s 0 , ¯ θ ( τ ) ∥ ∥ ∥ H ≲ √ τ , and the rest part follows similarly. Since

applying Cauchy-Schwartz inequality yields

Thus, again by Cauchy-Schwartz inequality, for any τ > τ 0 ≥ 0 , we have

By choosing τ 0 = 0 , we have ∥ a τ ∥ L 2 ( ρ 0 ) ≲ ∥ ∥ ∥ ¯ s 0 , ¯ θ (0) ∥ ∥ ∥ H + √ τ , hence ∥ ∥ ∥ ¯ s 0 , ¯ θ ( τ ) ∥ ∥ ∥ H ≲ ∥ ∥ ∥ ¯ s 0 , ¯ θ (0) ∥ ∥ ∥ H + √ τ , which completes the proof.

Lemma 6 (Monte Carlo estimates) . Define the Monte Carlo error

Suppose that ∥ x (0) ∥ ∞ ≤ 1 , and the trainable parameter a and embedding function e ( · ) are both bounded. Then, given any ¯ θ , for any δ > 0 , δ ≪ 1 , with the probability of at least 1 -δ , there exists θ such that

where ≲ hides universal positive constants only depending on T .

Proof. Fix any ¯ θ . According to Lemma 1, for any δ > 0 , δ ≪ 1 , with the probability of at least 1 -δ , we have

Based on the representation (13), for any W = ( w 1 , . . . , w m ) ⊤ ∈ R m × d , U = ( u 1 , . . . , u m ) ⊤ ∈ R m × d e with ( w i , u i ) ∼ ρ 0 , i = 1 , · · · , m , let A := ( a 1 , . . . , a m ) ∈ R d × m with a i := a ( w i , u i ) for i = 1 , · · · , m , and

then E W , U [ s t, θ ( x )] = E ( w , u ) ∼ ρ 0 [ a ( w , u ) σ ( w ⊤ x + u ⊤ e ( t )) ] = ¯ s t, ¯ θ ( x ) . For k = 1 , · · · , d , let

If ( ˜ W , ˜ U ) is different from ( W , U ) at only one component indexed by i , we have

where (i) is from the triangle inequality, (ii) is due to the fact that | σ ( y ) | = | ReLU ( y ) | ≤ | y | for any y ∈ R , the triangle and Hölder's inequality and (27), (iii) follows from the positive homogeneity property of the ReLU activation, and (iv) is due to the boundness of the trainable parameter a and

embedding function e ( · ) . By McDiarmid's inequality (see e.g., Lemma 26.4 in [43]), for any δ > 0 , with the probability of at least 1 -δ , we have

Since

̸

where (v) is due to Fubini's theorem, and (ii), (iii), (iv) is the same as before. By the triangle inequality, Jensen's inequality, (29) and (31), we obtain

̸

which gives

Obviously, 1 m ∥ A ∥ 2 F = 1 m ∑ m i =1 ∥ a ( w i , u i ) ∥ 2 2 = ∥ s t, θ ∥ 2 H kρ 0 < ∞ . The proof is completed.

Now we are ready to prove the main theorem.

Proof of Theorem 1. Based on Lemma 2, we have

To bound the term ˜ L ( ˆ θ n ( τ ); g 2 ( · )) , we use the following decomposition:

According to Lemma 4, we obtain

The term I 3 can be further divided into

where θ ∗ is the Monte Carlo estimator of ¯ θ ∗ . By Lemma 6, for any δ > 0 , δ ≪ 1 , with the probability of at least 1 -δ , it holds that

and similarly,

Hence, for any δ > 0 , δ ≪ 1 , with the probability of at least 1 -δ , it holds that

For the term I 1 , we have

Notice that which gives

Here, the triangle inequality, Cauchy-Schwartz inequality, the fact that | σ ( y ) | = | ReLU ( y ) | ≤ | y | for any y ∈ R , Hölder's inequality, the positive homogeneity property of the ReLU activation, and

the boundness of the input data, embedding function e ( t ) and weighting function λ ( t ) . Thus, we get

where the last inequality follows from Lemma 4. We further deduce that

where p ( n ) t denotes the empirical distribution of p t . Note that

by Lemma 5 we get ∥ θ ( τ ) ∥ 2 = ∥ ∥ s 0 , θ ( τ ) ∥ ∥ H ≲ ∥ ∥ s 0 , θ 0 ∥ ∥ H + √ τ/m . For any t ∈ [0 , T ] , define the function space

where

Then, according to Theorem A.5 in [61], for any δ ∈ (0 , 1) , with the probability at least 1 -δ over the choice of the dataset D x = { x i } n i =1 , it holds that

where ̂ Rad n denotes the (empirical) Rademacher complexity of F t on D x = { x i } n i =1 , and all the inequalities hold in the element-wise sense.

(i) For J 1 , according to Lemma A.6 in [61], we have

Note that which yields

and hence

where C ′ T,δ := max x ( t ) ∈ [ -C T,δ ,C T,δ ] d ∣ ∣ ∇ x ( t ) log p t ( x ( t )) ∣ ∣ . This gives

Let

according to Lemma 26.6 in [43], we get ̂ Rad n ( F 1 ,t ) ≤ ̂ Rad n ( F ′ 1 ,t ) . Since

where the { ξ i } n i =1 are independent random variables with the distribution P ( ξ i = 1) = P ( ξ i = -1) = 1 / 2 , and obviously,

we have

and by symmetry

i.e., ̂ Rad n ( F 1 ,t ) ≤ ̂ Rad n ( F ′ 1 ,t ) ≲ 2 ( ∥ ∥ s 0 , θ 0 ∥ ∥ H + √ τ/m ) ̂ Rad n ( F 2 ,t ) . According to Lemma 26.9 (contraction lemma) and Lemma 26.11 in [43], we have

Combining above, we obtain

(ii) For J 2 , by (36) and (37), we get

which gives

(iii) For J 3 , we similarly have

hence by (36) and (37), we get

Combining (i), (ii) and (iii), we obtain

which gives

where ≲ hides universal positive constants only depending on T .

Combining all above estimates yields

If one only focuses on the ( m,n ) -dependence, i.e. the dependence on the model capacity and sample size, this upper bound can be further simplified as

where we assume τ ≥ 1 for simplicity, and ( vi ) ≲ hides the term d log( d + 1) , the polynomials of log(1 /δ 2 ) and finite RKHS norms ∥ · ∥ H . The proof is completed.

Approximation errors. For the (universal) approximation, we discuss in two points:

First, the approximation is a separate problem that can be analyzed independently of training and generalization. While it is beyond the scope of current work, we have included the approximation error in the final estimates. When the approximation by random feature models fails, the generalization error is supposed to be significant.

In addition, the random feature model can approximate Lipschitz continuous functions on a compact domain (Theorem 6 in [17]). Notice that the forward diffusion process defines a random path ( x ( t ) , t ) t ∈ [0 ,T ] contained in a rectangular domain R T,δ := [ -C T,δ , C T,δ ] d × [0 , T ] ⊂ R d +1 with C T,δ := C T ( C x + √ log(1 /δ 2 )) (use Lemma 1 and the boundness of inputs), one can apply Theorem 6 in [17] to bound (7) on the domain R T,δ in R d +1 to obtain approximation results for Lipschitz continuous target score functions.

A.2 Data-Dependent Generalization Gap

Lemma 7 (Forward perturbation estimates, Gaussian mixtures) . Suppose that x is sampled from a one-dimensional 2-mode Gaussian mixture: p 0 ( x ) = q 1 N ( x ; -µ, 1) + q 2 N ( x ; µ, 1) , where µ > 0 , q 1 , q 2 > 0 with q 1 + q 2 = 1 are all constants. Then for any δ > 0 , δ ≪ 1 , µ > √ log(1 /δ 2 ) , with the probability of at least 1 -δ , we have

Proof. It is straightforward to verify that

where the last inequality applies (21). That is, for any δ > 0 , δ ≪ 1 , with the probability of at least 1 -δ , we have

Hence, Lemma 1 gives

which completes the proof.

Lemma 8 (Monte Carlo estimates, Gaussian mixtures) . Define the Monte Carlo error

Suppose that the trainable parameter a and embedding function e ( · ) are both bounded, and x is sampled from a one-dimensional 2-mode Gaussian mixture: p 0 ( x ) = q 1 N ( x ; -µ, 1) + q 2 N ( x ; µ, 1) , where µ > 0 , q 1 , q 2 > 0 with q 1 + q 2 = 1 are all constants. Then, given any ¯ θ , for any δ > 0 , δ ≪ 1 , µ > √ log(1 /δ 2 ) , with the probability of at least 1 -δ , there exists θ such that

where ≲ hides universal positive constants only depending on T .

Proof. According to Lemma 7, we just need to follow the proof of Lemma 6 by replacing C T,δ by C T,µ,δ . Notably, based on (33) in the proof of Lemma 6, one can finally derive

which gives the desired estimates.

Proof of Theorem 2. We decompose the loss ˜ L ( ˆ θ n ( τ )) in the same way as in the proof of Theorem 1. In fact, Theorem 2 is similarly proved by replacing C T,δ in the proof of Theorem 1 by C T,µ,δ . Note that C T,µ,δ = C T ( µ + √ log(1 /δ 2 ) ) ≲ µ for µ ≫ 1 , the proof is completed.

Remark 6. A standard variance is used here for convenience. In general, if var = o ( µ ) as µ → + ∞ (e.g. a bounded var), similar analysis and results are supposed to hold. However, this is different when var = Θ( µ ) as µ → + ∞ , since the modes are not separated in this case, and we can not characterize modes shift by simply varying µ .

Remark 7. Simply scaling down inputs seems not to resolve the adverse effect of modes shift, since the ground truth µ is unknown. One can use the input scale to approximate µ on toy datasets, but it is not that trivial in practice, particularly for real-world applications with multiple high-dimensional modes in varied scales.

The model-target inconsistency. Informally, there is inconsistency between the score network model and target score function. In fact, given the target Gaussian mixture p 0 ( x ) = q 1 N ( x ; -µ, 1) + q 2 N ( x ; µ, 1) , the target score function is

which gives s 0 ( x ) ≈ -x + µ , x ≥ 0 and s 0 ( x ) ≈ -x -µ , x ≤ 0 . While the score network model is

where " ≈ " holds since | x | = Θ( µ ) by (40) ( µ ≫ 1 ). Both s 0 and s 0 ,θ are linear functions, but they have unmatched scales in slopes and intercepts: O (1) and O ( µ ) for s 0 , but both O ( a ) 's for s 0 ,θ . That is, modeling Gaussian mixtures with large modes distances as random feature models is inconsistent.

Remark 8. Recall that the goal of this work is to estimate D KL ( p 0 ∥ p 0 , ˆ θ n ( τ ) ) , which aims at the density estimation characterization. This is different from [32], where the equivalence of learned and target data manifolds (support sets) is studied.

B Additional Experiments

B.1 Early-Stopping Generalization

Weillustrate the early-stopping generalization gap established in Theorem 1 using the Adam optimizer. All the configurations remain the same as Section 4.1 except that the learning rate is now 10 -3 .

From Figure 8, one can observe that the KL divergence achieves its minimum at the 1000th training epoch, and it starts to increase after this turning point. The plot is aligned with Theorem 1, which states that there exists an optimal early-stopping time when the model can generalize well, indicating the effectiveness of the upper bound. Further, the KL divergence begins to oscillate after the minimum point (1000th training epoch), which may suggest a phase transition in the KL divergence dynamics, and the transition point is around the optimal early-stopping time. The finding is aligned with SGD setting.

0.8

0.7

0.6

divergence

0.5

0.4

0.3

0.2

0.1

0.0

0

250

500

750

1000

1250

1500

1750

B.2 Modes Shift Effect

We further test the relationship between the modes' distance and the generalization performance using the Adam optimizer under the same configurations.

In Figure 9 and Figure 10, it is shown that the training of modeled distributions exhibits the same two-stage dynamics as the SGD setting, indicating that the modes shift effect holds not particularly for a certain optimizer .

epoch = 100

epoch = 1000

epoch = 1900

0.16

Trained

0.16 -

Target

Trained

0.30

Target

Trained

0.14

0.14

0.25

0.12

0.12

0.20

0.10

Aisua

0.10

0.08

aa

0.08

0.15

0.06

0.06

0.10

0.04

0.04

0.02

0.02

0.05

0.00

10.0

7.5

-5.0

2.5

0.0

2.5

5.0

7.5

10.0

0.00

10.0

7.5-5.0

2.5

0.0

2.5

5.0

7.5

10.0

0.00

10.0

7.5

-5.0

2.5

0.0

2.5

5.0

7.5

10.0

epoch = 100

epoch = 1000

epoch = 1900

0.30

Target

0.35

Target

Target

Trained

Trained

0.35

Trained

0.25

0.30

0.30

0.20

0.25

0.25

Aisuaa

0.15

0.20

0.15

0.15

0.10

0.10

0.10

0.05

0.05

0.05

0.00

30

20

-10

0

10

20

30

0.00

30

20

-10

0

10

20

30

0.00

30

20

-10

0

10

20

30

B.3 Model Capacity Dependency

We also numerically study the dependency of generalization on the model capacity. Following the same configurations in Section 4.1.1, we conduct experiments for different hidden dimensions varying from 2 1 to 2 11 .

10000

1.2

1.0

8000

divergence

Training epoch

0.8

6000

0.6

4000

0.4

0.2

2000

0.0

22

24

26

28

210

22

24

26

28

210

m

m

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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. \n- 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).\n- 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.\n\nWEAKNESSES:\n- 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.\n- 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}}$.\n- 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).\n- 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?\n\nQUESTIONS:\n- What is the role of the time embedding function $\\mathbf{e}$? Does it have any meaningful effect on the analysis?\n- 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?\n- Please clarify the second point from the \"Weaknesses\" section.\n- 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?\n- 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.\n- 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?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

comments: 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?

LIMITATIONS:
yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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}$ \nwhere the parameter $\\hat{\\theta}_n$  is a minimizer of the empirical loss function obtained by running gradient descent up to some time $\\tau.$ \n\nThe 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}.$\n \nThe 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].\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\nAside 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.\n\nQUESTIONS:\n1/ 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.\n2/ 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.\n3/ 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.\n\nLIMITATIONS:\nThe paper's proof might be incorrect.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

comments: 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)

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-18