Paper: Neural Network-Based Score Estimation in Diffusion Models: Optimization and Generalization (ICLR-2024)

Reviewer comment (Novelty):
The analysis strategy and framework presented in this paper do not entirely convince me. The primary contribution lies in establishing a connection between stochastic optimization (as defined in equation 5) and deterministic optimization (as defined in equation 8). Once this connection is made, the convergence and generation results appear as corollaries drawn from existing literature. Additionally, this connection is also not new, as it was proposed in [1].

Author response:
We respectfully disagree with the assessment that our results are merely incremental. There are major challenges that prevent the existing results (or their simple modifications) to be applied in our setting. To clarify this point and the misunderstandings, let us do a quick review of our results:

Our theoretical results cover three essential aspects of a learning problem: **approximation, optimization, and generalization**. To achieve our goal, we first derived a simple and easy-to-parameterize equivalence to the original loss function. The primary challenge is to find the **right error decomposition and control each component**. This was accomplished by analyzing four key elements: coupling, label mismatch, early stopping, and approximation. In particular, we established the coupling arguments between the neural network training and kernel regression with NTK-based analysis.  **An additional layer of complexity** arises from the non-standard nature of our problem, in contrast to the standard regression problems. These complexities include: a) unbounded input, b) vector-valued output, and c) the incorporation of a time variable. These challenges are nontrivial to address and  distinguish our NTK-based results from the existing ones in the deep learning literature. 

We sincerely hope that this explanation in addition to our general response (regarding our other general contributions) clarifies the misunderstanding.

Reviewer comment (Theory):
However, if one were to directly analyze or train the stochastic optimization in equation 5, the results would likely differ significantly, even when employing NTK-type analysis. This is because, in this case, the NTK would encapsulate randomness arising from the Brownian motion.

Author response:
We respectfully disagree with the assessment that our results are merely incremental. There are major challenges that prevent the existing results (or their simple modifications) to be applied in our setting. To clarify this point and the misunderstandings, let us do a quick review of our results:

Our theoretical results cover three essential aspects of a learning problem: **approximation, optimization, and generalization**. To achieve our goal, we first derived a simple and easy-to-parameterize equivalence to the original loss function. The primary challenge is to find the **right error decomposition and control each component**. This was accomplished by analyzing four key elements: coupling, label mismatch, early stopping, and approximation. In particular, we established the coupling arguments between the neural network training and kernel regression with NTK-based analysis.  **An additional layer of complexity** arises from the non-standard nature of our problem, in contrast to the standard regression problems. These complexities include: a) unbounded input, b) vector-valued output, and c) the incorporation of a time variable. These challenges are nontrivial to address and  distinguish our NTK-based results from the existing ones in the deep learning literature. 

We sincerely hope that this explanation in addition to our general response (regarding our other general contributions) clarifies the misunderstanding.

Reviewer comment (Experiments):
It's important to note that this work is purely theoretical and lacks empirical experimentation to validate its assumptions, such as Assumption 3.2, Assumption 3.4, Assumption 3.5, Assumption 3.7, Assumption 3.8, and Assumption 3.11.

Author response:
This is an important concern/question, thank you for raising it. Let us start our answer by first recalling that Assumption 3.2 simply says that the support of the target distribution is bounded. **This assumption is satisfied in almost all applications.**  For instance, for image data, the domain is bounded since the pixel values range  from 0 to 255.  Assumption 3.4 is **technically necessary** to make $\alpha(t)$ well defined. Notice that this boundedness assumption  is a clear generalization of  assuming $g$ is a constant. Furthermore, in our revision, **we verified Assumptions 3.5, 3.7 and 3.8** in Appendix G. In particular, we provide an **upper bound** of $\beta_x$ in Lemma G.1 to verify Assumption 3.5. Moreover, **explicit dependency** of $\delta_1(\Delta, R)$ and $\delta_2(d, N)$ in Assumptions 3.7 and 3.8 are provided in Lemma G.2 and G.3, respectively. We also include a reference in the revision to **validate Assumption 3.11** (see the bottom of page 8).  We would like to thank you for raising this concern which made us clarify these points in the revision and improve the quality of our manuscript.

Reviewer comment (Experiments):
The authors categorize errors into four parts: coupling, label mismatch, early stopping, and approximation. Without conducting numerical experiments, it becomes challenging to determine which error contributes the most. As a result, this work does not provide substantial practical insights.

Author response:
From a theoretical perspective,  we have analyzed the  **dependency** of each term with respect to model parameters  rigorously. Therefore, our result can  provide  partial insights on how large each term could possibly be. The approximation error and label mismatch terms are polynomial in $R$ and $R_\mathcal{H}$. Also, the coupling error is $1/\sqrt{m}$ dependent on the width of neural networks. 

We believe that establishing a theoretical foundation to fill the gap in the literature is essential for the community to develop a comprehensive understanding of diffusion models. We hope that this, in turn, will lead to practical insights on algorithm and architecture design. That being said, we do agree with the reviewer  that conducting extensive  numerical experiments is an important and interesting future direction to explore and would complement our work.

Reviewer comment (Experiments):
Can this analysis be extended to other architectures, let’s say, transformers?

Author response:
We are the **first** to address the open question of developing **algorithm-dependent** score matching guarantee in the literature. To achieve this, we establish the **convergence** of gradient descent algorithm and an over-parameterized two-layer neural network. Although simple, over-parameterized two-layer neural network already enjoys the universal approximation property when the width is sufficiently large. It has been widely used in other deep learning theory settings such as regression problem and reinforcement learning (https://arxiv.org/abs/1810.02054 and https://arxiv.org/abs/1909.01150). When the ReLU activation is applied, it results in a non-linear, non-convex, and non-smooth objective, leading to significant optimization challenges. We agree with the reviewer that studying how the analysis can be extended to other neural network architectures and other optimization algorithms is certainly an interesting next step.

Reviewer comment (Theory):
The curse of dimensionality not discussed. In particular, it is interesting to know for this problem exact dependency of all constants on dimension and discuss this in limitations of the work if exponential dependency is present.

Author response:
Great point! We agree that quantifying the explicit dependency of all the parameters with respect to the dimension is an interesting and important question. In our revised manuscript, we discussed **the issue of the curse of dimensionality** (see page 10) and left it as an important future research direction. In particular, the curse of dimensionality has been observed in the literature in the context of diffusion models and regression problems. See https://arxiv.org/pdf/2309.11420.pdf, https://arxiv.org/pdf/2302.07194.pdf, https://arxiv.org/pdf/2303.01861.pdf, and https://arxiv.org/pdf/2306.14859.pdf. In addition, if we assume that the data distribution has a **low intrinsic dimension**, the curse of dimensionality can be addressed by utilizing similar techniques as in https://arxiv.org/pdf/2302.07194.pdf.

Reviewer comment (Theory):
In Lemma 3.3 bound depends exponentially on dimension d, which makes me wonder -- do we have the curse of dimensionality in those bounds? I guess, R can be varied to improve this dependency but what is final dependency in bounds on dimension? If this is exponential, this will somewhat limit applicability of the results, at least, make them good for low dimension setting but for high-res diffusion models dimension is enormous and, hence, might not be something that explains performance of diffusion models.

Author response:
Thanks for raising this question. **One can choose the radius $R$ to improve the dependency.** Also, figuring out the final dependency of the excess risk on dimensionality is an interesting future direction, as we explained in our earlier response. 

For a general non-parametric regression, one can not get rid of the curse of dimensionality if we do not impose additional structure on the regression function and/or the data distribution, e.g., see https://arxiv.org/pdf/2306.14859.pdf. However, we can overcome the curse of dimensionality by assuming a low-dimensional structure of the data distribution. This idea has been explored in recent works on diffusion models, e.g.,  https://arxiv.org/pdf/2303.01861.pdf and https://arxiv.org/pdf/2302.07194.pdf. Our results fill a gap in the literature on diffusion models; please see the general response for a summary of our contributions.

Reviewer comment (Presentation):
The paper's presentation is a bit dense and can be improved.

Author response:
Thank you for your constructive feedback. Based on your suggestion, we have worked on our presentation and included more discussions and motivation (some major changes are marked in blue). We agree that we did not do a proper job in justifying the assumptions in our earlier version. We included the justifications of the major assumptions with verifications and more tractable conditions in our rebuttal version of the paper.

Reviewer comment (Theory):
The upper bounds presented in the paper seem far from tight. It is not clear whether these bounds are useful for getting guarantee in specific contexts.

Author response:
Thank you for your constructive feedback. Based on your suggestion, we have worked on our presentation and included more discussions and motivation (some major changes are marked in blue). We agree that we did not do a proper job in justifying the assumptions in our earlier version. We included the justifications of the major assumptions with verifications and more tractable conditions in our rebuttal version of the paper.

Reviewer comment (Novelty):
The parametric form this paper proposes does not seem to be novel.

Author response:
Thank you for your constructive feedback. Based on your suggestion, we have worked on our presentation and included more discussions and motivation (some major changes are marked in blue). We agree that we did not do a proper job in justifying the assumptions in our earlier version. We included the justifications of the major assumptions with verifications and more tractable conditions in our rebuttal version of the paper.

Reviewer comment (Theory):
Could the authors elaborate a little bit more on why is it reasonable to fix \(a\) throught training and only update \(W\)? In a random feature model \(a\) is trained while \(W\) is generated randomly. In the NTK regime it is also assumed that \(W\) does not change too much during training, hence the problem reduces to fitting a linear model with \(a\) representing the coefficients.

Author response:
Fixing $a$ and only updating $W$ is a **standard set-up** used in the literature of (over-parameterized) deep learning theory and reinforcement learning, e.g., see https://arxiv.org/abs/1810.02054 and https://arxiv.org/abs/1909.01150, as it already enjoys universal approximation property. However, when the ReLU activation is applied, it still leads to a **non-linear, non-convex, and non-smooth objective** that is challenging to optimize. Note that the random feature model is essentially a *linear model*, which is much simpler compared to our non-linear model. In the NTK regime, we need to fit a non-linear model parametrized by $W$ (with $a$ fixed). *Although NTK is similar to a random feature model, they are **not equivalent**.* See the discussion on page 4.

Reviewer comment (Experiments):
Why the authors propose to uniformly sample the time? In practice usually a non-uniform weight function is employed. How does the choice of weight function affect the result.

Author response:
Although we sample the time uniformly over an interval $[T_0, T]$, we introduce a (non-uniform) **weight function** $\lambda(t)$ in our score matching objective (4). In particular, we propose to choose $\lambda(t) = h^2(t)/\alpha^2(t)$, which leads to a simplified regression objective (12). In practice, we expect a similar performance if we sample the time according to the density function $\lambda(t)$.

Reviewer comment (Novelty):
I feel the most general form of Lemma 3.1 has already been established in many past works. See for example, section 5.1 of https://arxiv.org/pdf/2306.09251.pdf and the intro section of https://arxiv.org/pdf/2309.11420.pdf. I think the authors should at least cite these papers and discuss the relation.

Author response:
Thank you for bringing these works to our attention. We have cited these references and discussed them in the revised manuscript. A similar form of Lemma 3.1 has been proved in https://arxiv.org/pdf/2302.07194.pdf for data with linear structure,  in https://arxiv.org/pdf/2306.09251.pdf for the discrete-time setting and in a concurrent work https://arxiv.org/pdf/2309.11420.pdf. We have included a discussion on this on page 5.

Reviewer comment (Reproducibility):
How is \(\gamma\) initialized?

Author response:
We initialize $\gamma(0) = H^{-1}u(0)$, where
$$ H \coloneqq \begin{pmatrix}
            H_{11} & \cdots &  H_{1N} \\\\ \vdots & \ddots & \vdots \\\\ H_{N1} & \cdots & H_{NN}
        \end{pmatrix}, \quad H^{ik}_{j\ell} \coloneqq z^{\top}\_{j}z\_{\ell}\mathbb{E}\left[a\_{1}^{i}a\_{1}^{k}\mathbb{I}\left\lbrace z\_{j}^{\top}w\_{1}(0) \geq 0, z\_{\ell}^{\top}w\_{1}(0) \geq 0\right\rbrace\right],$$

and $u(0) = (u\_1(0)^\top, \dots, u\_N(0)^{\top}) \in \mathbb{R}^{dN}$ with $u\_j(0)^{\top} = (u\_j^i(0))\_{i = 1}^{d} =  (f\_{{\bf W}(0)}^i(X\_{t\_j}, t\_j))_{i = 1}^{d}$. We have clarified this in the revised manuscript. See the bottom of page 6.

Reviewer comment (Writing):
The notation \(\beta_x\) is a bit confusiong. I assume it should not depend on \(x\). Maybe the authors can state what does it depend on?

Author response:
You are correct that the Lipschitz constant is  *independent of $x$*. In particular, the constant $\beta_x$ only depends on $f_*$ and $T_0$ and does not rely on $x$, $t$ and $T$. We have the index $x$ to show that the Lipschitz constant is w.r.t. the variable $x$ (and not any other variable). This is a standard practice in the optimization literature. We have clarified this in the revised manuscript. See the middle of page 7.

Reviewer comment (Theory):
I think Assumption 3.5 is a fact instead of an assumption when the target is bounded, at least when \(g\) is a positive constant. This is because taking the gradient of the conditional expectation gives the conditional covaraince, which has bounded operator norm when data is bounded.

Author response:
Great point! Based on your suggestion, we added Lemma G.1 to establish the **upper bound** $\beta_x = \mathcal{O}(D/h(T_0))$ in the revised manuscript. See the middle of page 7.

Reviewer comment (Theory):
In Theorem 3.6, should I interpret \(c_1\) as a universal constant? If not, what does it depend on?

Author response:
The constant $c_{1}$ is equal to the constant $C(d+1, 0)$ in [Proposition 6, Bach, 2017]. We have clarified this in the revised manuscript. See the footnote on page 7.

Reviewer comment (Writing):
In Assumption 3.6 you mean \(1 - \delta (\Delta, R)\)?

Author response:
Correct. An explicit dependency is provided in Lemma G.2 in the appendix (see the middle of page 8). In particular, we show that $\delta(\Delta, R)$ is linear in $\Delta$ and exponentially dependent on $R$.  We have clarified this point in our revision.

Reviewer comment (Theory):
There are two delta functions in Assumption 3.7 and 3.8, what are their relation?

Author response:
Thank you for pointing out the issue. We have changed the notations to $\delta_1$ in Assumption 3.7 and $\delta_2$ in Assumption 3.8. See the bottom of page 7 and the top of page 8.

Reviewer comment (Theory):
I feel the second term in Theorem 3.9 upper bound is huge. If it is not, maybe the authors can comment on it a little bit.

Author response:
When both $R$ and $\Delta$ are fixed, choosing $m = {\rm poly}(d, N, C_{\max}, C_{\min}, \delta, \lambda_0)$ ensures that the second term becomes **arbitrarily small**. Consequently, the coupling error in Theorem 3.9 can be arbitrarily small. We added a comment in the revised manuscript accordingly. See the middle of page 8.

Reviewer comment (Theory):
Could the authors elaborate more on why "Assumption 3.11 can be satisfied by an extension of classical early stopping rules for scalar-valued kernel regression." Maybe giving an example in which this assumption is satisfied would be helpful for readers to digest.

Author response:
Assumption 3.11 requires an **early stopping rule for a vector-valued kernel regression problem**. In the statistical learning literature, early stopping rules for scalar-valued kernel regression are well-established, e.g., Raskutti et al., 2014. We believe a *generalization* of the scalar-valued analysis satisfies Assumption 3.11. See the middle of page 9 and Lemma G.3 in the appendix.