Improved Sample Complexity Bounds for Diffusion Model Training

We appreciate your thorough review and valuable suggestions on the presentation of the paper. We appreciate the time and effort you have invested in providing this feedback.

Implications of our work

Our work is the first to show a polynomial sample complexity for learning a diffusion models using a deep network under a reasonable assumption -- that the scores of the distribution can be represented using a $P$-parameter neural network. From a theoretical point of view, we believe this is a significant contribution -- it was not clear prior to our work whether minimizing the score-matching objective is sufficient to learn a diffusion model using a small number of samples.

More broadly, we hope that this style of assumption will inspire future work in understanding diffusion models and other deep learning models. %We would like to highlight the concrete new insights our research contributes to the field of diffusion models.

Responses to specific questions

Can the authors develop more why and when the dependence on $\theta$, $\gamma$, $D$ and $P$ is more important than the dependence on $\varepsilon$? Can a numerical estimate be provided here?

To illustrate the dependence on these parameters, we provide some natural choices of parameters:

• TV error $\varepsilon = 0.01$.
• Wasserstein error $\gamma = 0.01$.
• Space dimension $d = 12288 = 64 \times 64 \times 3$.
• Total Parameter Number $P = 270M = 2.7 \times 10^8$.
• Network Depth $D = 12$.
• Parameter Weight Range $\Theta = 5$.

For parameter number and the space dimension, we used the parameters in stable diffusion[1]. Our bound gives $\frac{d^2 PD}{\varepsilon^3} \cdot \log \Theta \cdot \log^3 \left(\frac{1}{\gamma}\right) \approx 7.7 \times 10^{25}$, compared to the previous bound of $\frac{d^{5/2}}{\gamma^3 \varepsilon^2} \left(\Theta^2 P\right)^D \sqrt{D} \approx 5.2 \times 10^{138}$.

While we do not claim that $\theta$, $\gamma$, $D$, and $P$ are universally more important than $\varepsilon$, our results demonstrate a significant exponential improvement in $D$ and $\gamma$ with only a minor loss in $\varepsilon$, resulting in a substantially better overall bound.

Is it true that assumption A2 implies $\inf_f \mathbb{E}_{X \sim q_t}[\| f(x) - s_t(x) \|_2] = 0$? Does this assumption implicitly assume that the score function is "simple" enough to belong to the family?

Actually, we only need the $L^2$ error to be polynomially small, namely $\frac{\delta \cdot \varepsilon^3}{N^2 \sigma^2_{T - t_k} \log \frac{d}{\gamma}}$, as specified in Theorem C.3. We will put this in the main body of the paper. Intuitively, this assumption states that neural networks can effectively represent the score function. This assumption is both reasonable (neural networks have proven to be surprisingly powerful function approximators) and necessary (if the distribution's score isn't well approximated by a neural network, diffusion won't work).

What are $\gamma$ and $R$ in line 41?

In line 41, $R$ represents the radius of the support size of the distribution, and $\gamma$ is a parameter used to specify the Wasserstein error.

In Setting 1.1, does $P$ denote the number of parameters per layer?

In Setting 1.1, $P$ denotes the total number of parameters, not just per layer. We will clarify this in the final version to avoid any ambiguity.

[1] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

We appreciate your comments and are glad that you like and support the paper.

The assumption that networks can represent the score is somehow strong. We agree that it is somewhat strong, but view this more as a statement about data; analogous to "the data is sparse in X basis". Over the past decade neural networks have proven to be surprisingly effective function approximators. We ask: supposing this is true, how many samples do you need to accurately learn a generative model?

(And if it's not true, then trying to learn the score with a neural network is doomed to fail regardless of how many samples you use.)

Q1 (Applying Lemma 1.4 for all N): That's right, Theorem 1.2 applies Lemma 1.4 in a union bound over the N time steps to invoke Lemma 1.3. Since the dependence on probability $\delta_{train}$ is logarithmic, this loses at most a $\log N$ factor. Not good if $N$ is really $\infty$, but fine for more reasonable values of $N$.

That said, there is an interesting question here. In our analysis, it doesn't matter if the image samples are chosen independently or jointly across different scales. In practice, they are typically chosen jointly (so each sample image is then used at all noise scales). It seems plausible to us that jointly sampling images could improve the sample complexity in Theorem 1.2 from $d^2$ down to $d$, by correlating the "bad" events across scales.

Q2 (Quantitative Assumption 2): We give the precise bound in Theorem C.3 in the appendix. It's a polynomial bound, namely $\frac{\delta \varepsilon^3}{N^2 \sigma^2 \log \frac{d}{\gamma}}$.

Q3 (Other function classes): Our proof actually applies to any Lipschitz activation function, we should have specified this. We think measuring a neural network by its parameter count / weight / depth is a simple way to bound its complexity. One can certainly analyze other classes -- Lemma A.2 applies to arbitrary finite function classes, and one can take a net for other classes -- but we're not sure of any more general statement that applies cleanly to neural networks.

(Some work (e.g. BMR20) tries to use Rademacher complexity to get a general statement, but the Rademacher complexity of neural networks is exponential in their depth so this is unsatisfying.)

Comparison to previous work

We would like to clarify that most of the prior work in the literature, including the works you mentioned, ask about the approximation power of neural networks for representing the score of arbitrary distributions, and/or make strong assumptions on the distribution. Typically, for $d$ dimensional distributions, the number of parameters necessary is exponential in $d$ -- indeed, the bounds shown in all the works you mentioned suffer from this exponential dependence. But in practice, neural networks are surprisingly powerful function approximators of real-world functions, and they seem to be able to represent $d$-dimensional distributions with much less than $\exp(d)$ parameters.

To circumvent this issue, we take a different route: we assume that a $P$ parameter neural network can approximate the score of the distribution that we are trying to learn. Note that such an assumption is necessary: if our network cannot even represent the scores, there is no hope of learning it.

We then ask: if $P$ parameters suffice, how many samples do we need to learn this approximation? We show that the sample complexity for this problem is only polynomial in all relevant parameters, and even logarithmic in the Wasserstein error, unlike all prior works. We believe that this question and our answer more accurately capture the situation in practice than the one above.

We appreciate the feedback to provide a table -- please see the general response.

Note that [HRX24], which is focused on the two-layer network case in the NTK regime, actually makes the assumption that gradient descent in RKHS using a sufficient number of samples is sufficient to learn the score (Assumption 3.11). Our paper proves that a small number of samples is sufficient to learn the score under the reasonable assumption that the network can represent it; moreover, we provide a quantitative bound that is polynomial in all relevant parameters.

Technical contributions

Here is a brief summary of our technical contributions:

• We exploit the fact that recent works on sampling only require an error of $\varepsilon^2/\sigma_t^2$ for each $t$ to show a sample complexity bound independent of $\sigma_t$ for score estimation with this error. That is, the accuracy required fortuitously matches the accuracy achievable at each scale $\sigma_t$, with a sample complexity independent of $\sigma_t$.

• This lets us run until $\sigma_t$ is extremely small; this leads to a sample complexity logarithmic in the final Wasserstein error.

• To show our score estimation result for a single time $t$, instead of making the strong assumption that the distribution is bounded as in [BMR20], we observe that with high probability the score at time $t$ is bounded in terms of $\sigma_t$. This observation, combined with a technical argument allows us to remove the $R^3$ dependence in the sample complexity in [BMR20].

• Other contributions that we will refrain from describing in detail here in the interest of space -- we show that we learn the score in a new $1-\delta$ robust sense, rather than standard $L^2$ to circumvent barriers in learning in $L^2$, and we make use of a careful net argument to obtain our exponential improvements in the dependencies on the depth and range of the neural network relative to [BMR20].

Section 4

• As above, section 4 shows that learning in $L^2$ requires $\text{poly}(1/\gamma)$ samples to obtain a final Wasserstein error of $\gamma$. We can circumvent this barrier by learning the score in our weaker $1-\delta$ robust sense, to obtain a final sample complexity only logarithmic in $1/\gamma$.

• Minimax optimality: As explained above, prior works show minimax optimality with respect to the class of arbitrary distrbutions, which results in an exponential in $d$ sample complexity. This does not accurately capture the situation in practice -- neural networks are surprisingly powerful approximators of real-world functions and seem to be able to represent $d$ dimensional distributions with much less than $\exp(d)$ parameters/samples. We circumvent this by making the reasonable assumption that the scores of the distribution can be represented with a $P$ parameter neural network.

[HRX24]: Neural Network-Based Score Estimation in Diffusion Models: Optimization and Generalization. Yinbin Han, Meisam Razaviyayn, Renyuan Xu. https://arxiv.org/abs/2401.15604

Thank you for your detailed response. I have raised the score. However, the presentation should still be improved before publication.

Thank you for your comments and questions. As you state, prior work like Chen et al. have shown that "sampling is as easy as learning the score." Thus the main open question is: how easy is learning the score?

We think that question is clearly important enough for top tier machine learning conferences. We address it by providing the first polynomial sample complexity bound, under the assumption that the score can be effectively represented by neural networks. And this assumption is both reasonable (neural networks have proven to be surprisingly powerful function approximators) and necessary (if the distribution's score isn't well approximated by a neural network, diffusion won't work).

Assumption A2:

Thank you for the suggestion, which is the consensus of reviewers. We will include the precise definition of ``sufficiently small" into the main text in the final version of the paper. As stated in C.3., the required $L^2$ error is $\frac{\delta \varepsilon^3}{N^2 \sigma^2 \log \frac{d}{\gamma}}$.

Technical novelties:

Here is a brief summary of our technical contributions:

• We exploit the fact that recent works on sampling only require an error of $\varepsilon^2/\sigma_t^2$ for each $t$ to show a sample complexity bound independent of $\sigma_t$ for score estimation with this error. That is, the accuracy required fortuitously matches the accuracy achievable at each scale $\sigma_t$, with a sample complexity independent of $\sigma_t$.

• This lets us run until $\sigma_t$ is extremely small; this leads to a sample complexity logarithmic in the final Wasserstein error.

• To show our score estimation result for a single time $t$, instead of making the strong assumption that the distribution is bounded as in [BMR20], we observe that with high probability the score at time $t$ is bounded in terms of $\sigma_t$. This observation, combined with a technical argument allows us to remove the $R^3$ dependence in the sample complexity in [BMR20].

• Other contributions that we will refrain from describing in detail here in the interest of space -- we learn the score in a new $1-\delta$ robust sense, rather than standard $L^2$, to circumvent barriers in learning in $L^2$, and we make use of a careful net argument to obtain our exponential improvements in the dependencies on the depth and range of the neural network relative to [BMR20].

VE vs VP:

The VE and VP processes are simply reparameterizations of each other. That is, if $x_t$ denotes a solution to the VE SDE, then $y_t = e^{-t} x_{e^{2t}-1}$ is a solution to the VP SDE. The reverse process and its discretization can be easily adjusted to match this change of variables. So, these processes are pretty much identical, and the same results hold.