Paper: Scaling Convex Neural Networks with Burer-Monteiro Factorization (ICLR-2024)

Reviewer comment (Experiments):
There appears to be a mismatch between the title and the actual contents of the paper. The title is "Scaling convex neural networks with BM factorization". However, there is no experimental evidence showing that applying BM factorization to convex neural networks leads to more scalable implementations. The title would only be meaningful if the paper compares the baseline that optimizes the two-layer networks by gradient descent and presents improvements over efficiency. This mismatch is accompanied by several claims that seem inaccurate or lack sufficient justifications.

Author response:
1. **Showing that applying BM factorization to convex neural networks leads to more scalable implementations.**

To help show more results on the fact that the BM factorization leads to more scalable implementations, we have added an additional section in Appendix A.1 which demonstrates the superior time complexity of the BM formulation compared to the original non-convex objective.

In addition, the reviewer has some questions around the need for layerwise training and how it relates to this paper. Below we can explain in more depth the rationale behind these experiments.

While training deep neural networks directly is possible using SGD heuristically applied to poorly understood non-convex objectives, we seek to find approaches that can optimize equivalent architectures to deep neural networks while having convex guarantees. While most of the theory around convex neural networks allows us to re-state two-layer neural networks as convex programs, we wish to extend the capabilities of this approach to deeper networks. Thus, layerwise learning is a natural solution -- each two-layer sub-component of a deeper neural network can be trained to global optimum using convex optimization. Layerwise learning has already shown great success in the non-convex training domain as shown in Belilovsky et al., 2019. However, naively using convex optimization to train layerwise networks does not scale computationally for the reasons discussed in Section 4.2.

Therefore, we show how the BM factorization applied to these convex networks (1) scales in a way that the original convex program doesn't, allowing us to train these architectures in a layerwise fashion in an approach that matches the performance of end-to-end trained deep neural networks; and (2) preserves convex guarantees around optimality and has an easily computable relative optimality gap bound. We hope this makes the motivations of the layerwise training experiments more clear to the reviewer.

Reviewer comment (Theory):
The claim that "Without the BM factorization, the induced regularizer of convex CNNs is intractable to compute, and the latent representation used for layerwise learning is prohibitively large" remains insufficiently justified. Specifically, it is unclear why layerwise training is necessary, and why one cannot simply train all layers together as in the baseline method.

Author response:
In addition, the reviewer has some questions around the need for layerwise training and how it relates to this paper. Below we can explain in more depth the rationale behind these experiments.

While training deep neural networks directly is possible using SGD heuristically applied to poorly understood non-convex objectives, we seek to find approaches that can optimize equivalent architectures to deep neural networks while having convex guarantees. While most of the theory around convex neural networks allows us to re-state two-layer neural networks as convex programs, we wish to extend the capabilities of this approach to deeper networks. Thus, layerwise learning is a natural solution -- each two-layer sub-component of a deeper neural network can be trained to global optimum using convex optimization. Layerwise learning has already shown great success in the non-convex training domain as shown in Belilovsky et al., 2019. However, naively using convex optimization to train layerwise networks does not scale computationally for the reasons discussed in Section 4.2.

Therefore, we show how the BM factorization applied to these convex networks (1) scales in a way that the original convex program doesn't, allowing us to train these architectures in a layerwise fashion in an approach that matches the performance of end-to-end trained deep neural networks; and (2) preserves convex guarantees around optimality and has an easily computable relative optimality gap bound. We hope this makes the motivations of the layerwise training experiments more clear to the reviewer.

Reviewer comment (Novelty):
The theorems seem to be straightforward extensions of prior works (e.g., Bach et al. (2008) and Haeffele et al. (2014)). A comparison with prior works at the technical level is needed, with clear justification of the novelty of the paper’s contributions. A comparison table might be appropriate.

Author response:
2. **What's the novelty of the paper at the technical level?**

The primary **technical** novelty against other works such as Bach et al. (2008) and Haeffele et al. (2014) is in our **characterization of stationary points**. In Theorems 3.3 and 3.4, we characterize the relative optimality gap for stationary points of the BM factorization, which is not shown in any other work. Most other work around the BM factorization focus on local minima only. This is especially significant because for most NN objectives, there is no guarantee that BM will actually find local minima, as we discuss in Section 2.2.
We note that there are other aspects to novelty in our work such as results in neural networks, though we believe the reviewer does not dispute this and is more curious about the technical innovations in how we study the BM factorization.

Reviewer comment (Presentation):
In Eq. (1), it might be beneficial to add parentheses for the summation term.

Author response:
Thank you for the minor comments as well -- these have been revised in the updated version (see changes in the revised PDF in blue).

Reviewer comment (Writing):
In the sentence below Eq. (5), there are two occurrences of "thereby", resulting in awkward phrasing.

Author response:
Thank you for the minor comments as well -- these have been revised in the updated version (see changes in the revised PDF in blue).

Reviewer comment (Theory):
In the paragraph after Eq. (19), the statement "*While the original convex program is NP-hard*" is unclear, as it is not explained why convex programs would be NP-hard.

Author response:
Thank you for the minor comments as well -- these have been revised in the updated version (see changes in the revised PDF in blue).

Reviewer comment (Novelty):
The proposed BM formulation doesn’t seem to provide any advantages other than dealing with quasi-nuclear norms. As BM requires \( m \geq d + c \) to ensure no spurious minima, it increases the number of variables from \( dc \) to \( m(d+c) \) (for MLP). For linear and gate ReLU activated MLP, it seems like it would be a lot more efficient to directly solve for the original convex problem.

Author response:
We have provided additional explanations of the per-iteration complexity of the Burer-Monteiro factorization compared to the original convex problem.

While the Burer-Monteiro (BM) factorization adds additional variables (the original convex problem has $Pdc$ variables, while the BM factorized optimization problem has $Pm(d+c)$ variables), the per-iteration complexity is greatly reduced in the BM case. In the ReLU activation case,  we can compare applying the Frank-Wolfe algorithm presented in (Sahiner et al., 2021b) to the convex problem (3) to projected gradient descent applied to (19). 

In (Sahiner et al., 2021b) Appendix A.4.2, it is shown that a single Frank-Wolfe iterate on (3) is $\mathcal{O}(Pn^r)$, where $r := \mathrm{rank}(X)$. In contrast, a single iteration of projected gradient descent on (19) naively takes $\mathcal{O}(P(n^3 + n^2(d+c + dm) + dcm))$ time. It's clear that unless $m$ is chosen to be exponential in $n$ (which would never occur in practice), the per-iteration complexity of solving (3) is order of magnitude larger than that of (19) due to the exponential dependence on $r$. Therefore using the BM factorization greatly reduces the per-iteration complexity of solving the neural network optimization problem

We have added an additional section in the Appendix (Appendix A.1) to walk through the details of this result.

Reviewer comment (Theory):
The proposed BM formulation for CNN is still not practical as the memory requirement is too high. In fact, the authors adopt layerwise training for solving the BM formulation, which has no guarantee to converge to global optimum.

Author response:
The proposed BM formulation for two-layer CNNs **is** practical -- this is why we use it in Section 4.2. The purpose of doing layerwise training is to expand these results to carry to deeper neural networks, rather than to overcome shortcomings of the computational complexity of the BM factorization for CNNs.

Reviewer comment (Theory):
Continuing from Points 1 and 2, the paper is titled “Scaling Convex Neural Network” but it seems like the proposed BM formulation only addresses the intractability of quasi-nuclear norm. It does not reduce the time complexity and the memory requirement for solving the convex neural network training problem. It is questionable that the proposed BM formulation would be scalable to large 2-layer neural networks due to the above reasons.

Author response:
See above and Appendices A.1 and C.3 for an explanation of how we reduce the time complexity and memory requirements for solving the convex neural network problem.

Reviewer comment (Novelty):
The novelty of this paper is really limited. In particular, this paper is just applying Burer-Monteiro to the existing 2-layer convex neural networks, which would be too incremental to be published on ICLR.

Author response:
Aside from the BM factorization proposed in this work being the first time the BM factorization has been applied to neural network optimization, we present novel results in a few ways:
1. We derive a novel bound on the relative optimality of the stationary points of the BM factorization for neural networks. This is the first time stationary points of BM factorization have been showed in any work -- other works have solely focused on local minima, which are not guaranteed to be found with e.g. SGD.
2. Accordingly, we are the first to develop conditions to check if stationary points of BM factorizations have achieved their global optimum.
3. We also use this BM theory to uncover insights about novel neural network architectures, e.g. that linear self-attention networks with sufficiently many attention heads have no spurious local minima.

Reviewer comment (Reproducibility):
Could the authors elaborate the architecture of the CNN network used in Section 4.2? Specifically, exactly which “architecture of (Belilovsky et al., 2019)” is used.

Author response:
See Appendix C.2 for a description of the experimental setup we use in this experiment based on Belilovsky et al., 2019. Namely, we train two-layer CNNs (consisting at convolutional layer with 256 filters + gated ReLU + average pooling to size $2 \times 2$ + fully connected layer) for multiple stages. After training each stage, the fully connected layer is discarded, and the convolutional filters of the current stage are frozen and used to generate the input for the next stage.

Reviewer comment (Theory):
What is the rate of convergence for solving the proposed BM formulation using GD? Does GD become sublinear convergence when it is close to the optimal due to the ill-conditioning on \( U \) and \( V \) (using MLP for example, \( U \) and \( V \) must be rank deficient at optimality due to \( m \geq d + c \))?

Author response:
For this, we can refer to some literature [1] showing the linear convergence of preconditioned GD applied to BM factorized problems with convex objectives, though, as pointed out by the reviewer, naive GD may indeed suffer from sublinear convergence if $m$ exceeds the dimension of the optimal solution, i.e. $U_j$ and $V_j$ are rank deficient. We have noted this in Appendix A.1.


[1] Gavin Zhang, Salar Fattahi, and Richard Y Zhang. Preconditioned gradient descent for overparameterized nonconvex burer-monteiro factorization with global optimality certification., JMLR 2023

Reviewer comment (Experiments):
Could the authors provide experimental results demonstrating the scalability of the BM formulation? For example, the time and memory complexity (per iteration) vs. the number of neurons.

Author response:
We have provided an additional experiment of this nature in Appendix C.3. We demonstrate that while the convex neural network is much more inefficient compared to the BM factorization when training a single two-layer CNN, this only becomes exacerbated by the layerwise learning setting.

Reviewer comment (Novelty):
I would bring to the author's attention that Zhang et. al. proposed a convex version of MLP [1] and of CNN [2]. In particular, for [2], the framework also imposes a low-rank nuclear norm constraint. In the arxiv version of [2] (https://arxiv.org/pdf/1609.01000.pdf), they also provided experiments to the scale of CIFAR-10 dataset, where they also applied low-rank kernel matrix factorization techniques like Nyström approximation, random feature approximation, Hadamard transform, etc. The discussion on how the proposed method is compared to [1] and [2] should be included at least theoretically, if not experimentally.

Author response:
We thank the reviewer for notes on [1] and [2] -- these are interesting works that propose convex alternatives to neural networks. We should note that these are slightly different than the line of work we focus on in our paper, because they do not exactly match the objectives and architectures of typical non-convex neural networks, whereas in this paper we focus on exact convex formulations of an identical neural network objective (i.e. a solution to (3) maps one-to-one with a solution to (1)). We have included a brief mention of [1] and [2] in our related work section.

Reviewer comment (Theory):
The authors could include a big O notation analysis on the computational cost to better illustrate exactly how much the BM factorization helps with scaling the method.

Author response:
We have now included this big O notation analysis to illustrate the per-iteration complexity of solving the BM factorization compared to the original convex formulation in Appendix A.1.

Reviewer comment (Theory):
How feasible is it to generalize this framework to more than 1 hidden layer?

Author response:
Unfortunately, there have not been many satisfactory convex formulations of deeper neural networks that are amenable to this kind of analysis. Hence our emphasis on layerwise training in Section 4.2.

Reviewer comment (Writing):
It is hard to keep track of different norm notations at times, so the authors may consider defining all of them at once at least in the appendix. For instance in Lemma 3.1, when R and C are introduced, it's not clear that those are just L^p norm with p = R and C or something else. It's also not clear that ||.||_F corresponds to Frobenius norm or some special norm that comes from the objective function F.

Author response:
We have updated our "notations" section in Section 2 to make these more clear.

Reviewer comment (Theory):
Is there a feasible framework to analyze saddle points of the BM formulation and not just local minima?

Author response:
Theorems 3.3 and 3.4 do exactly this -- they show results around the stationary points (which include saddle points) of the BM formulation, rather than just local minima which has been the focus of much of the existing BM literature.