Sample-Efficient Training for Score-Based Diffusion

We appreciate your positive reception of the concept of the $1 - \delta$ error we introduced. As you suggested, we have uploaded a revised manuscript with a clearer presentation. In particular, we appreciate and have addressed your corrections and suggestions such as defining $m_2$ before using it, and correcting several typos.

Explanation of $\text{poly}(\log(m_2/\gamma))$: This is motivated by the discussion in the previous section, where we note that prior works improve a $\text{poly}(m_2/\gamma)$ to a $\text{poly}(\log(m_2/\gamma))$ for sampling. The natural question is to see if this improvement is possible in training, which is what we investigate.

Domain size: You make a fair point, $\mathcal{H}$ doesn't necessarily depend on the domain size but it typically will. For neural networks, for example, this induces a logarithmic dependence on the domain size. (See Theorem 1.5).

Appendix D of Lee et al: The two examples are essentially the opposite of each other. They give an example of distributions with small $L^2$ error and large TV error; we give distributions with large $L^2$ error and small TV error. They show non-annealed Langevin dynamics fails even with small $L^2$ error; we show annealed Langevin dynamics succeeds even for large $L^2$ error.

Note also that the conclusion of Lee et al explicitly asks the question we address in this work:

The assumption that we have a score estimate that is $O(1)$-accurate in $L^2$, although weaker than the usual assumptions for theoretical analysis, is in fact still a strong condition in practice that seems unlikely to be satisfied (and difficult to check) when learning complex distributions such as distributions of images. What would a reasonable weaker condition be, and in what sense can we still obtain reasonable samples?

Our distance measure $D_p^\delta$ is a weaker condition than $L^2$ that is easier to learn in but still yields reasonable samples.

Bound on score: This is implied by Tweedie's formula: $s(x) = \mathbb{E}_{z|x}[-z/\sigma^2]$, and the numerator is $O(R)$.

Difference between $\hat{s}$ and $\tilde{s}$: $\tilde{s}$ refers to some candidate score function in $\mathcal{H}$ that we are trying to analyze the loss function evaluated for, while $\hat{s}$ is the empirical loss minimizer among these.

Thank you for your review. Below we have addressed the issues that you brought up, and hope that this helps clarify and explain the theory.

Scope: "While the improvement in terms of accuracy in Wasserstein distance is remarkable, the dependency on the accuracy in TV remains polynomial."

Learning a distribution within TV error $\varepsilon$ requires $1 / \varepsilon^2$ samples information theoretically, so we cannot hope for going beyond polynomial for the dependency on the accuracy in TV.

"What's the point of improving drastically the accuracy in W2, if the accuracy in TV remains bad?"

The W2 and TV error are incommensurate. The W2 error is "how blurry is the result", while the TV error is "what fraction of the distribution are completely missed". These are on different scales; if you want very sharp images, you want low W2 error, and we show that this can be achieved.

The authors should track how the bound scales in d and compare to existing work. The new version presents this. For sampling, our bound is identical to the best existing analysis of the SDE (so, $O(d)$). For training, estimating a single score (Theorem 1.1) takes sample complexity linear in $d$; estimating all the scores gives a $d^2$ dependence for Corollary 1.3. For comparison, using the Block et al. bound would need $d^3$ for Corollary 1.3. [They give a bound in terms of the Lipschitzness, which can be $\Theta(d)$ in the regime they consider.]

It would also add value to the paper to track the dependency on something more explicit than the cardinality of the hypothesis class $\mathcal{H}$..

What would you propose? The cardinality is pretty explicit, and leads to good bounds for neural networks (see Theorem 1.5). One might try the Rademacher complexity of $\mathcal{H}$, but this doesn't work well: the Rademacher complexity scales with the magnitude of the outputs, which (e.g. in the neural network example) is large and will yield a $\text{poly}(1/\gamma)$ dependence.

Along the same lines, for the result on neural networks (Theorem 3.1), one needs to assume the existence of a network that approximates the score well enough.

How could it be otherwise? If you can't represent the score, you aren't going to sample well with score-based diffusion.

Proof novelty: The idea of the proof is not actually to exclude a region of low probability. The Block et al analysis shows for each possible score function $s$ that the score matching objective approximates its $L^2$ error, with enough samples. The problem is that getting this many samples requires $\text{poly}(\frac{1}{\gamma})$ samples, because for some score functions $s$ the value in the score matching objective has very large variance.

The idea in our proof is to essentially argue that those score functions are very bad, so it's OK to have a poor approximation to just how bad they are. We show that the difference between the score matching objective for the optimal $s^*$ and any candidate $s$ has variance related to its expectation, such that every sufficiently-far $s$ will be rejected with high probability.

The proof involving excluding regions of low probability is just to bridge the gap between existing theory and our new proposed error metric, which we justify the use of in section 4. Proving convergence in our new error metric is the main contribution.

Thank you for your review. We have made several changes to the manuscript with your comments in mind. In particular, we have added formal statements to the main text to make it clearer what our results are.

We clarify some questions you brought up below:

The setting looks artificial compared to true score-based model.

We have no idea what you mean here. Our setting is exactly that of DDPM, as deployed in settings like Stable Diffusion. You train a neural network with score matching; you sample by discretizing the SDE. The question is: how many training samples do you need for accurate sampling?

Example. You are correct that on this example, the actual sampling procedure will have no problem. But the existing theory does: the score functions are not learned in $L^2$, so the existing theory cannot say the sampling will be accurate.

Your argument is essentially the intuition for our paper: it's OK to have large $L^2$ error on average, as long as the squared error is "generally" small (see our distance measure (8)).

Thank you for your review. We would like to first clarify that our main contribution is not so much about unbounded distributions, as about higher accuracy: to get accuracy $\gamma m_2$ in Wasserstein, our approach proves a $\text{poly}(\log \frac{1}{\gamma})$ sample complexity while prior work needed $\text{poly}(\frac{1}{\gamma})$. And getting this improved dependence requires switching from $L^2$ approximation to a different measure.

Addressing weaknesses (1, 2): We have rewritten the paper significantly, and included the formal theorem statements in the introduction.

Addressing weakness (3): First of all, our bound for finite $\mathcal{H}$ is strong enough to extend to continuous neural networks as a corollary by essentially rounding the weights to finite precision (see Theorem 1.5). So " finite" vs "continuous" is not what makes the result challenging.

Generally speaking, convergence of the empirical mean of a distribution depends on the variance of that distribution. The score matching objective involves terms involving the norm of the output of a neural network; for the wrong-weight networks, this norm can be huge, making the convergence very slow; and to get accuracy $\gamma$, such an approach gives a $\text{poly}(\frac{1}{\gamma})$ rather than the $\text{poly}(\log \frac{1}{\gamma})$ that we are able to achieve.

Answering specific questions:

1, 2, 3. Thank you, we have rewritten the introduction more clearly.

4. Again, we have written this more clearly; sampling requires sufficiently accurate scores, which Theorem 1.1 provides.

5. We now include the dependence. It is linear in $d$ and the complexity class $\log \mathcal{H}$. For neural networks (Theorem 1.5), the dependency is linear in $d$, the number of parameters $P$, and the neural network depth $D$.