Paper: Sample-Efficient Training for Score-Based Diffusion (ICLR-2024)

Reviewer comment (Writing):
Also on page 1, it writes 'logarithmic in $m_2/\epsilon$', where neither $m_2$ or $\epsilon$ is defined before. Actually $m_2$ is defined in Theorem 1.2 for the first time.

Author response:
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.

Reviewer comment (Writing):
Before equation (6), we see $m_2$ again, still undefined.

Author response:
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.

Reviewer comment (Theory):
In the part named Our results, what's the intuition of 'one would like to show ... with poly($log \frac{m_2}{\gamma}$)'? Where does this poly($log \frac{m_2}{\gamma}$) come from?

Author response:
**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.

Reviewer comment (Theory):
In the remark after Corollary 1.1, it states that 'Corollary 1.1 doesn’t depend on the domain size at all'. But H contains sufficiently accurate approximations to each score. Does it mean that $|H|$ should be related to the domain size?

Author response:
**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).

Reviewer comment (Novelty):
How is the discussion in section 4 related to that in Appendix D of [1]?

Author response:
**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.

Reviewer comment (Theory):
On page 5, it states that 'If x is bounded by some value $R$, then the score would be bounded by $R/\sigma^2$. Can you prove it or provide a reference?

Author response:
**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)$.

Reviewer comment (Theory):
In Theorem A.1, are $\hat{s}_r$ and $\tilde{s}_r$ different? It seems that in the proof you assume they are equal (but actually they are not I think).

Author response:
**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.

Reviewer comment (Evaluation):
The scope of the paper is quite limited. While the improvement in terms of accuracy in Wasserstein distance is remarkable, the dependency on the accuracy in TV remains polynomial. What's the point of improving drastically the accuracy in W2, if the accuracy in TV remains bad? For this reason, while I like the idea of using the quantile measure, the benefit of doing so in the context of score-based diffusion remains unclear.

Author response:
**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.

Reviewer comment (Theory):
The complexity is also polynomial in $d$. Recent papers analysing diffusion models (see e.g. Table 1 in "Linear convergence bounds for diffusion models via stochastic localization") show that this dependency is linear or quadratic in $d$. I appreciate that the setting of this paper is different (most works assume access to a good L2 score, which is shown to be impossible if one sticks with a logarithmic dependency in $1/\gamma$). Nonetheless, the authors should track how the bound scales in $d$ and compare to existing work.

Author response:
**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.]

Reviewer comment (Novelty):
The novelty in terms of proof is not high. Basically the idea is to exclude a region with low probability and then use existing analyses (mostly, Benton et al., 2023). This is admittedly a minor point (simple ideas can be very useful!). The two key weaknesses above are the main reason of my score.

Author response:
**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.

Reviewer comment (Presentation):
I feel the presentation is not clear enough for readers to fully understand the merit of this paper. For more details see my question list.

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

Reviewer comment (Theory):
Corollary 1.1 does not seem very surprising as $H$ is finite. The authors contain proof in their paper. But maybe it would be helpful to explain intuitively why it is interesting to establish Corollary 1.1, and why it is not a straightforward consequence of the uniform law of large numbers.

Author response:
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.

Reviewer comment (Writing):
I think $m_2$ appears a lot before the authors define it as the second moment.

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

Reviewer comment (Theory):
What is the relation between the last sentence on page 3 and Theorem 1.2? Theorem 1.2 seems to be independent of training, why can we see from this theorem that our outlier-robust approximation suffices for sampling?

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

Reviewer comment (Theory):
Maybe the authors can comment a little bit on the polynomial dependency? Like what is the order of the polynomial. Block et al 2020 has an explicit polynomial dependency, and it is not clear if the results presented here is indeed better if the form of polynomial is not presented.

Author response:
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$.