Paper: Analyzing Neural Network Based Generative Diffusion Models via Convexification (ICLR-2024)

Reviewer comment (Theory):
The paper considers the regularized score matching loss rather than the standard one.

Author response:
Firstly, our convex program for the DSM objective holds for any positive weight decay parameter which can be taken to be very small and zero in the limit. Our convex program for the SM objective requires weight decay parameter to be greater than certain small threshold, which is usually necessary for the SM objective to have finite optimal value. We illustrate this via a small example below.  Consider for example we have only one data point and one hidden neuron, then objective function for NN with ReLU activation and no skip connection would be (see the general form in Appendix A.3)
\begin{equation*}\begin{aligned}\frac{1}{2}((xw+b)_+\alpha+b_0)^2+w\alpha\mathbb{1}(xw+b\geq 0)+\frac{\beta}{2}(w^2+\alpha^2).\end{aligned}\end{equation*}
WLOG consider data $x=1,$ then when weight decay $0\leq \beta<1,$ set $b=-w+\sqrt{1-\beta}$ and $b_0=0$ above, we get
\begin{equation*}
\begin{aligned}
\frac{1-\beta}{2}\alpha^2+w\alpha+\frac{\beta}{2}(w^2+\alpha^2)
\end{aligned}
\end{equation*}
Then we can set $\alpha=-w$ and the above expression becomes $(\beta-1)w^2/2$. Thus the objective goes to $-\infty$ when $w$ goes to $\infty$. Note in Theorem 4.1, we requires $\beta>1$, and we consider this reasonable.

Reviewer comment (Theory):
The paper studies the convergence of Langevine dynamics rather than the more widely used backward process in DDPM papers.

Author response:
One of our purpose is to show the convexification bypasses some conventional computation difficulties in optimizing SM objective, we consider Langevin dynamics since there is no denoising score matching. We discuss and use the annealed Langevin sampling in our experiments. The annealed Langevin strategy is employed together with denoising score matching for multiple noise levels in Algorithm 3 and then used in our simulation result, see Figure 4 for example.

Reviewer comment (Experiments):
The Gaussian data distribution in experiments look too simple compared to the ones in practice.

Author response:
The simulation in Figure 3 is used to verify our theoretic results in Section 4.1.1. We explain that no matter what the ground truth data distribution is, when weight decay parameter $\beta$ lies in certain  range, the predicted score function  would always be the same and can be parameterized by sample mean and sample variance. Here we simply pick the ground truth data distribution to be Gaussian. We have added more simulation results with more complex data distributions in the doc and note our theory holds through all these cases. For simulation in Figure 4, the middle plot shows sampling result with non-convex NN solution as score predictor. The sample distribution is nowhere near to standard Gaussian while the sample result with convex score predictor on the right is close to standard Gaussian. We feel in this case, simple Gaussian distribution is enough to reveal the superiority of our convex score predictor compared to its non-convex counterpart.

Reviewer comment (Theory):
I am curious how the convexity arises. Is it because of the regularization term in the SM and DSM objectives?

Author response:
This is due to the hidden convexity in ReLU neural networks. Essentially, there exists an embedding of the non-convex problem to a higher-dimensional convex problem. Moreover, the solutions can be mapped between them. However, our results were not known before since the specific form of SM and DSM objectives are quite different from standard supervised learning objectives. 

The convexification is not due to the regularization term since the regularization can be taken zero in the limit. However, the weight regularization takes the form of a convex regularizer ($\ell_1$) in the convexified problem.

Reviewer comment (Theory):
In many convergence theory, we hope to have a score estimator close to the ground truth score function in $L^2$ sense. Although the optimization objective is derived, it is not clear whether the optimal solution can have good generalization ability in population loss.

Author response:
In this paper, we mostly focus on the optimization aspects of the SM and DSM objectives and constructing a more transparent convex model equivalent to those. The generalization error can be analyzed based on our results.

Reviewer comment (Evaluation):
For Gaussian data distribution, the score function is actually linear. I am surprised in Figure 4 that the non-convex neural estimator trained by DSM cannot recover the ground truth score function. Could you provide some intuition?

Author response:
In Figure 4, we do DSM with annealed Langevin sampling, so the neural network is trained to estimate the score function of the noisy data distribution corresponding to multiple noise level. The poor sample result likely indicates that two-layer NN of the form we consider is not good at predicting some of the noisy data distribution and the annealed Langevin sampling with the predicted score function then fails to produce standard Gaussian samples.

Reviewer comment (Writing):
The writing can be significantly improved. I suggest adding more interpretations after each theoretical result.

Author response:
Thanks for the suggestions, we will make further improvements regarding to the writing.

Reviewer comment (Theory):
It seems unnatural to add a regularization term along with the SM or DSM objective function. By adding the regularization term, the minimizer of the SM/DSM objective function (population version) will not be the intended score function, so the generated samples from such diffusion models will not follow the distribution of the training data. Indeed, Theorem 6.1 and 6.2 show that the obtained sample will follow Gaussian/Gaussian-Laplace distribution, no matter which distribution the training data is from. This defeats the original purpose of diffusion generative modeling.

Author response:
Firstly, our convex program for the DSM objective holds for any positive weight decay parameter which can be taken to be very small and zero in the limit. Our convex program for the SM objective requires weight decay parameter to be greater than certain small threshold, which is usually necessary for the SM objective to have finite optimal value. We illustrate this via a small example. Consider for example we have only one data point and one hidden neuron, then objective function for NN with relu activation and no skip coNNection would be (see the general form in Appendix A.3)
\begin{equation*}\begin{aligned}\frac{1}{2}((xw+b)_+\alpha+b_0)^2+w\alpha\mathbb{1}(xw+b\geq 0)+\frac{\beta}{2}(w^2+\alpha^2).\end{aligned}\end{equation*}
WLOG consider data $x=1,$ then when weight decay $0\leq \beta<1,$ set $b=-w+\sqrt{1-\beta}$ and $b_0=0$ above, we get
\begin{equation*}
\begin{aligned}
\frac{1-\beta}{2}\alpha^2+w\alpha+\frac{\beta}{2}(w^2+\alpha^2)
\end{aligned}
\end{equation*}
Then we can set $\alpha=-w$ and the above expression becomes $(\beta-1)w^2/2$. Thus the objective goes to $-\infty$ when $w$ goes to $\infty$. Note in Theorem 4.1, we requires $\beta>1$, and we consider this reasonable.
For DSM objective, we require regularization $\beta>0$ for the strong duality to hold and $\beta$ can be set very small to mimic the zero weight decay case. 

Theorem 6.1 and 6.2 discuss some large weight decay regime where we are able to solve the convex program analytically and discover Gaussian/Laplacian mixture. For general weight decay,  once an optimal solution $y^\star$ to the corresponding convex program is derived, the predicted score function is just $\sum y^\star_i\sigma(x-x_i)+b^\star$ where $b^\star$ can be computed from $y^\star$. We update our doc with the following general sampling result. For ReLU activaton, similar result holds by replacing sign function with indicator function.

For NN with absolute value activation, note once the optimal $y^\star$ to the convex program is known, the optimal score function is simply 
\begin{equation*}
\begin{aligned}
s(x)=\sum_{i=1}^n (y_{i}^\star+y_{n+i}^\star) |x-x_i|+b^\star
\end{aligned}
\end{equation*}
Integrating the above score function, we obtain 
\begin{equation*}
\begin{aligned}
\int s(x)dx=\frac{1}{2}\sum_{i=1}^n (y_{i}^\star+y_{n+i}^\star)(x-x_i)^2\text{sign}(x-x_i)
+b^\star x
\end{aligned}
\end{equation*}
The corresponding probability density is proportional to
\begin{equation*}
\begin{aligned}
\exp(\int s(x)dx)=\exp(\frac{1}{2}\sum_{i=1}^n (y_{i}^\star+y_{n+i}^\star)(x-x_i)^2\text{sign}(x-x_i)+b^\star x).
\end{aligned}
\end{equation*}

Note that the log-density is piecewise quadratic with breakpoints only at a subset of training points.

Reviewer comment (Theory):
The convex program in 1 dimension is efficiently solvable (the regularized objective). However, in high dimensions, the convex objective function involves combinatorial sums, which is not efficiently computable.

Author response:
The convex objective function says that we enumerating all sign patterns, denote the number of sign patterns as $P$, and solve a convex program with $P$ variables of dimension $d$ by $d$. Here $D_i$ and $X$ is fixed and we are only optimizing over $W_i$, which contains no combinatorial sums. Note $P\in O(r(n/r)^r)$ where $r$ is the rank of data matrix. In order to guarantee global optimality, the program  takes $O((r(n/r)^{r}d^2)^3)$.  In practice, when $P$ grows large, one can simply sample a subset of sign patterns and solve the convex program with these sub-sampled sign patterns. See Theorem 2.1, Theorem 3.3 and Algorithm 1 in [1].

[1] Aaron Mishkin, Arda Sahiner, and Mert Pilanci. Fast convex optimization for two-layer
relu networks: Equivalent model classes and cone decompositions, 2022.

Reviewer comment (Theory):
In the statement of Theorem 4.2, $D_j = D_j$ when ... and $D_j = 2 D_j - 1$ when ... is not a serious statement.

Author response:
Thanks for pointing out the ambiguity, here we are saying the coefficient matrix $D_j$ in Theorem 4.2 is equal to sign pattern $D_i$ explained at the beginning of Section 4.2 and $D_j=2D_i-I$ for absolute value activation. We've updated the doc to reflect this.

Reviewer comment (Presentation):
More generally, the presentation is difficult to follow, and things are not well motivated. For instance, theorem 4.1 is stated, but no intuition is provided for what it means. The matrices/vectors in the theorem are described in subsequent sections, but again, not much motivation/intuition is provided. It seems like a lot of this should just be pushed to the appendix.

Author response:
Thanks for this comment. We will improve the readability in the revision. Since we study the convex program of different neural network types including ReLU and absolute value activation, and with or without skip connections, the convex programs are of similar quadratic form with slightly different coefficient matrices/vectors, so we express the general form in Theorem 4.1 and  specify coefficient matrices/vectors for different neural network types in the following paragraph. We feel that the specific coefficient matrices/vectors are useful for readers to see how the derived convex program is related to the input data so we put them in the main text. Again, our work aims at both  bypassing difficulties in training NN with vanilla SM objective and understanding NN approximation error in the finite training data regime, which is crucial for NN-based score-based generative models (please see our reply to your comment 1).

Reviewer comment (Experiments):
The Langevin algorithm is studied for sampling, while in practice, annealed Langevin is used after learning the score at multiple noise levels. None of this is explored in this paper. Furthermore, the sampling result itself seems to be a consequence of classical work, since the distribution learned is always very simple in the weight decay regime studied.

Author response:
In fact, we discuss and use the annealed Langevin sampling in our experiments. The annealed Langevin strategy is employed together with denoising score matching for multiple noise levels in Algorithm 3 and then used in our simulation result, see Figure 4 for example. Our sampling results are enabled by the convex reparameterizations followed by classic convergence result of log concave sampling (see https://chewisinho.github.io/main.pdf). However, they are specialized to the explicit score function obtained by solving the convex programs. Ultimately, our convergence results apply for NN-based score function learned with finite data and we are unaware of any prior work on guarantees for any NN-based learned score function.

Reviewer comment (Novelty):
What is the motivation for studying this problem?

Author response:
Our project is indeed well-motivated. Existing theory works mainly establish convergence of diffusion process when the learned score approximates score of unknown data distribution well, but in reality only empirical approximation is viable due to finite training samples. Current literature falls short in understanding the role of NN approximation error for score-based generative models and it's difficult
to characterize the distribution from which these models sample in practice. For example, ([1],[2],[3]) show NN-based score-based generative models given finite training data generalizes due to approximation errors introduced by limited NN model capacity and also optimization errors. Our work contributes in understanding what score function NN learns in finite training data regime (see reply to your comment 2 below). Simply speaking, we characterize what two-layer NN exactly learns in finite data regime, there is no optimization error since we are able to solve the convex program exactly.

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

[2] TaeHo Yoon, Joo Young Choi, Sehyun Kwon, and Ernest K. Ryu. Diffusion probabilistic models
generalize when they fail to memorize. In ICML 2023 Workshop on Structured Probabilistic
Inference & Generative Modeling, 2023. 

[3] Mingyang Yi, Jiacheng Sun, and Zhenguo Li. On the generalization of diffusion model.

Reviewer comment (Experiments):
Is there anything you can say when there is no weight decay?

Author response:
Please see our reply to your comment 2. To summarize, we show that weight decay is crucial for vanilla SM objective for optimal value to be finite. For DSM objective, our convex program holds for any positive weight decay and can be taken to be zero in the limit. We list some unexpected Gaussian/Gaussian-Laplace results in specific weight decay regimes since we can solve the convex program analytically. But our convex program is general and once the optimal value $y^\star$ and $b^\star$ is derived, the predicted score function is just $\sum y^\star_i\sigma(x-x_i)+b^*$, which is a piecewise linear function with breakpoints at a subset of data points.  Note that the $\ell_1$ regularization encourages the convex program to have a small subset of data points as the breakpoins of the score function.

Reviewer comment (Theory):
Can you provide more detail in the section studying the denoising score matching objective?

Author response:
Our main result for DSM is the convex program derived in Theorem 5.1. We spent a large portion of the paper on the SM objective since one of our purpose is to show the convexification bypasses some conventional computation difficulties in optimizing the SM objective. For DSM convex program, note it holds for any positive weight decay and it's a simple Lasso problem. Our theoretical and simulation results make the comparison between SM and DSM fully transparent.

Reviewer comment (Theory):
The general promise of the paper seems too ambitious, as of course, there is no way (yet) to obtain general convex optimisation formulation. The paper focuses on specific cases and somehow it has to be clarified for full transparency that the results apply to specific cases.

Author response:
We state clearly that we present a theoretical framework to analyze two-layer NN-based diffusion models in the abstract, this is realized via Theorem 3.1 for univariate data and Theorem 3.2 for multivariate data for score matching objective, and via Theorem 4.1/4.2 for denoising score matching objective. The neural network model we study is presented at the beginning of Section 4, which includes both ReLU and absolute value activation and with or without skip connection. We consider this general enough for two-layer neural network architecture.  In our analysis for predicted score function, we focus on specific weight decay regime where we can solve the convex program analytically and we discover Gaussian/Laplacian mixture. For general weight decay,  once an optimal solution $y^\star$ to the corresponding convex program is derived, the predicted score function is just $\sum y^\star_i\sigma(x-x_i)+b^\star$ where $b^\star$ can be computed from $y^\star$. We update our doc with the following general sampling result. For ReLU activaton, similar result holds by replacing sign function with indicator function.

For NN with absolute value activation, note once the optimal $y^\star$ to the convex program is known, the optimal score function is simply 
\begin{equation*}\begin{aligned}s(x)=\sum_{i=1}^n (y_{i}^\star+y_{n+i}^\star) |x-x_i|+b^\star\end{aligned}\end{equation*}
Integrating the above score function, we obtain 
\begin{equation*}\begin{aligned}\int s(x)dx=\frac{1}{2}\sum_{i=1}^n (y_{i}^\star+y_{n+i}^\star)(x-x_i)^2\text{sign}(x-x_i)
+b^\star x\end{aligned}\end{equation*}
The corresponding probability density is proportional to
\begin{equation*}\begin{aligned}\exp(\int s(x)dx)=\exp(\frac{1}{2}\sum_{i=1}^n (y_{i}^\star+y_{n+i}^\star)(x-x_i)^2\text{sign}(x-x_i)+b^\star x).\end{aligned}\end{equation*}

Note that the log-density is piecewise quadratic with breakpoints only at a subset of training points.

Reviewer comment (Theory):
The results of convexification relies on specific structure of the networks that makes sense. However, it is quite obvious that, also, these networks are not what's used in practice. It would be appropriate to update the text to reflect the limitations.

Author response:
Please see our reply for  your comment 1. Our work aims at both revealing the ability of  bypassing conventional difficulties in training NN with vanilla SM objective via convexification and understanding NN approximation error in finite training data regime, which is crucial for NN-based score-based generative models.  Existing theory works mainly establish convergence of diffusion process when the learned score approximates score of unknown data distribution well, but in reality only empirical approximation is viable due to finite training samples. Current literature falls short in understanding the role of NN approximation error for score-based generative models and it's difficult
to characterize the distribution from which these models sample in practice. For example, ([1],[2],[3]) show NN-based score-based generative models given finite training data generalizes due to approximation errors introduced by limited NN model capacity and also optimization errors. We establish convergence results apply for NN-based score function learned with finite data and we are unaware of any prior work on guarantees for any NN-based learned score function. 

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

[2] TaeHo Yoon, Joo Young Choi, Sehyun Kwon, and Ernest K. Ryu. Diffusion probabilistic models
generalize when they fail to memorize. In ICML 2023 Workshop on Structured Probabilistic
Inference & Generative Modeling, 2023. 

[3] Mingyang Yi, Jiacheng Sun, and Zhenguo Li. On the generalization of diffusion model.

Reviewer comment (Theory):
The authors also didn't discuss the general difficulty of extending these results to more realistic networks. A discussion on this would be helpful. In fact, ReLU can be replaced with other basic units, would these results hold or easily extended to those cases?

Author response:
The current convex reformulation can be extended to any piecewise linear activation, it is not obvious how to extend it to smooth activations such as SiLU/GeLU trivially. However, since our reformulation resolves the differentiability issue of the ReLU activation, and hence resorting to smooth activations might be unnecessary. 

It's possible to extend the convex formulations to deeper networks based on recursive hyperplane arrangements (see [4] and [5]). However, our work uncovers many complexities even in two-layer ReLU score models that were not known before.

[4] T. Ergen, M. Pilanci, Path Regularization: A Convexity and Sparsity Inducing Regularization for Parallel ReLU Networks. NeurIPS 2023.

[5] T. Ergen, M. Pilanci, Global Optimality Beyond Two Layers: Training Deep ReLU Networks via Convex Programs.
ICML 2021.

Reviewer comment (Experiments):
While authors mentioned the computational difficulty of score matching, it would be appropriate to discuss the numerical cost of this procedure and have a comparison with score matching in terms of runtimes for similar problems. Is it much more efficient?

Author response:
Our convex program in Theorem 3.1 is a standard cone program with multiple of $n$ variables and no constraints  when we have $n$ training points. Therefore, the complexity of solving it with standard interior point methods is $O(n^3)$. Our convex program in Theorem 3.2 is a standard cone program with $Pd^2$ variables where $d$ is the dimension of the training samples and $P$ is the number of samples for the hyperplane arrangement. The complexity is $O((Pd^2)^3)$. In order to guarantee global optimality, the program in Theorem 3.1 takes $O(n^3)$ time and the program in Theorem 3.2 takes $O((r(n/r)^{r}d^2)^3)$ where $r$ is the rank of the training data matrix.

In contrast, applying a gradient method such as (stochastic) gradient descent has no guarantees of convergence or local optimality. At best, one could aim to prove convergence to a stationary point, however, to the best of our knowledge, this is not known for the SM or DSM objectives.

Reviewer comment (Writing):
The word "Langevine" is used in many parts of the paper instead of "Langevin".

Author response:
Thanks for pointing out the typos, we have made update according to the suggestions.