Scaling Convex Neural Networks with Burer-Monteiro Factorization
Authors: Arda Sahiner
Tolga Ergen
Batu Ozturkler
John M. Pauly
Morteza Mardani
Mert Pilanci
Venue: ICLR 2024 poster (2024)
Keywords: burer-monteiro
convex optimization
neural networks
stationary points
global optima
relu activation
Primary area: optimization
TL;DR: We apply the Burer-Monteiro factorization to two-layer ReLU (fully-connected, convolutional, self-attention) neural networks by leveraging their implicit convexity, and provide insights into stationary points and local optima of these networks.

Abstract: It has been demonstrated that the training problem for a variety of (non) linear two-layer neural networks (such as two-layer perceptrons, convolutional networks, and self-attention) can be posed as equivalent convex optimization problems, with an induced regularizer which encourages low rank. However, this regularizer becomes prohibitively expensive to compute at moderate scales, impeding training convex neural networks. To this end, we propose applying the Burer-Monteiro factorization to convex neural networks, which for the first time enables a Burer-Monteiro perspective on neural networks with non-linearities. This factorization leads to an equivalent yet computationally tractable non-convex alternative with no spurious local minima. We develop a novel relative optimality bound of stationary points of the Burer-Monteiro factorization, providing verifiable conditions under which any stationary point is a global optimum. Further, for the first time, we show that linear self-attention with sufficiently many heads has no spurious local minima. Our experiments validate the novel relative optimality bound and the utility of the Burer-Monteiro factorization for scaling convex neural networks.

Published as a conference paper at ICLR 2024
SCALING CONVEX NEURAL NETWORKS WITH BURER- MONTEIRO FACTORIZATION
Arda Sahiner
Tolga Ergen LG Al Research
Batu Ozturkler Stanford University
John Pauly Stanford University
Arcus Inc. Stanford University
Morteza Mardani
Mert Pilanci
NVIDIA Corporation
Stanford University
ABSTRACT
It has been demonstrated that the training problem for a variety of (non) linear two-layer neural networks (such as two-layer perceptrons, convolutional networks, and self-attention) can be posed as equivalent convex optimization problems, with an induced regularizer which encourages low rank. However, this regularizer becomes prohibitively expensive compute moderate seales, impeding training convex neural networks. To this end, propose applying the Burer-Monteiro factorization to convex neural networks, which for the first time enables a Burer- Monteiro perspective
neural networks with linearitssss This factorization leads to an equivalent yet computationally tractable non-conven alternative with no spurious local minima. We develop novel relative optimality bound of stationary points of the Burer-Monteiro factorization, providing verifiable conditions under which any stationary point is global optimum. For the first time, we show that linear self-attention with sufficiently many heads has no spurious local minima.
INTRODUCTION
It has been demonstrated that the training problem for (non-linear) two-layer neural networks are equivalent convex programs (Pilanci & Ergen 2020 Frgen & Pilanci| 2020) Sahiner et al.) 202Thb Ergen et al. 2021; Sahiner et al.) 2021 This has been observed for a variety of architectures, in- cluding multi-layer perceptrons (MLPs) (Pilanci Ergen 2020 Sahiner et al. (021b), convolutional neural networks (CNNs) (Ergen & Pilanci, 2020) Sahiner et al. (021c), and sell-attention based transformers (Sahiner et al. 2022). A major benefit of convex training of neural networks is that global optimality guaranteed, which brings transparency to training neural networks.
The convex formulation of neural networks induces biases by regularization of the network weights. For linear activation, the convex model directly imposes nuclear-norm regularization which is well- known to encourage low-rank solutions (Recht et al., 2010). For ReLU activation, however, the convex model induces type of nuclear norm which promotes sparse factorization while the left factor is constrained to an affine space (Sahiner et al. (b22 This constrained nuclear-norm is NP-hard compute. This impedes the utility of convex neural networks for ReLU activation. To address this computational challenge, we seek a method which (i) inherits the per-iteration complexity of non-convex training of neural network, and (ii) inherits the optimality guarantees and transparency of convex training. To find solution, we leverage the well-studied Burer-Monteiro (BM) factorization (Burer & Monteiro 2003).
BM has been applied
efficient solution strategy for problems ranging from matrix factoriza-
tion ZZeng & Lafferty 2016 Park et al. 2017, Ge et al. 2017 Gillis 2017) to rank minimiza- tion MMrdani et al. 2013; Recht et al. 2010 Wang et al. 2017) and matrix completion (Mardani et al. 2015 Ge et al. 2017). BM has also been used for over- simplifee neural networks such as (Kawaguchi 2016 Haeffele & Vidal 2017 Du & Lee 2018), where optimality conditions for local minima are provided. However, no work has deployed BM factorization for practical non non-linear neural networks, and guarantees available about the optimality of stationary points.

---

Published as a conference paper at ICLR 2024
Thus, our focus in this work is to adapt BM for practical two-layer (non-linear) convex neural networks. We consider three common architectures, namely MLPs, CNNs, and self-attention networks. For these scenarios, we develop verifiable relative optimality bounds for all local minima and stationary points, which are easy and interpretable. In light of these conditions, we identify useful insights about the nature of neural networks contributing to optimality. In particular, we observe that for self-attention networks all local minima coincide with the global optima if there are sufficiently many heads. The optimality guarantees also provide useful algorithmic insights, allowing one to verify whether the -iggt- first-order methods such as SGD achieve the global optimum for the non-convex training of neural networks. Our experiments with image classification task indicate that this BM factorization enables layerwise training of convex CNNs, which allows for convex networks for the first time to match the performance of multi-layer end-to-end trained non-convex CNNs.
1.1 CONTRIBUTIONS
All in all, our contributions are summarized as follows:
We propose the BM factorization for efficiently solving convex neural networks with ReLU activation for moderate and large scales. This is the first time BM theory has been applied to the non-linear neural network setting the best of our knowledge.
We derive novel bound the relative optimality of the stationary points of the BM factorization for neural networks.
We identify simple and verifiable conditions which guarantee a stationary point of the non-convex BM formulation achieves the global optimum of the convex neural network.
We provide insights into the fundamental building blocks of neural networks that contribute to optimality; egg. that linear self-attention has no spurious local minima if it has sufficiently many heads.
.
Our experiments verify the proposed relative optimality bound of stationary points from the BM factorization, and uncovers cases where SGD converges to saddle points, even in two-layer neural networks.
1.2 RELATED WORK
Burer-Monteiro factorization. The Burer-Monteiro (BM) factorization was first introduced in (Burer Monteiro, 2003 (005). There has been a long line of work studying the use of this factorization for solving SDPs (Boumal et al.) 2016 Cifuentes & Moitra 2019; Waldspurger & Waters 2020; Erdogdu et al., 2021). In the rectangular matrix case, gradient descent converges to global optimum of the matrix factorization problem with high probability for certain classes of matrices (Zheng & Lafferty| 2016). The BM factorization has been also studied in the rectangular
case in more generic settings (Bach et al.
2008 Haeffele et al. 2014 Haeffele & Vidal 2017).
The ability of nuclear norm regularization to induce low rank has been studied extensively in compressed sensing (Candès & Recht, 2009) Recht et al. 2010 Candès & Tao 2010). BM factorization has been applied to scale up nuclear-norm minimization (Mardani ct al. 2015 2013). It has also been deployed for low-rank matrix factorization (Cabral al.) [013)] Zhu et al., 2017] Park et al.) 22017; 7: Ge et al.) (0177). The results show that all sccond-order critical points of the BM factorization are global optima if certain qualification conditions are met.
Nuclear norm and rank minimization.
SGD for non-convex neural networks. It has been shown that for parammeeeetee two-layer linear networks, all local minima are global minima (Kawaguchi 2016). Accordingly, line of work has attempted to show that gradient descent or its modifications provably find local minima and escape saddle points (Ge et al.) 2015: Lee et al. [0166] Jin et al. 2017 Daneshmand et al.) 2018). However, these works assume Lipschitz gradients and Hessians of the non-convex objective, which is not typically satisfied. Another line of work shows that gradient descent converges to global optima for sufficiently highly over-parameterized neural networks, with either the parameter count being a high-order polynomial of the sample count (Du al. 2018: 2019; Arora al. 2019), the network architecture being simple (Du & Lee 2018). In practice, it has been empirically observed that SGD
can converge local maxima, or get stuck in saddle points (Du et al. 2017 Ziyin et al. 2021). For
2

---

Published as a conference paper at ICLR 2024
unregularized matrix factorization, it has also recently been shown that randomly initialized gradient descent provably converges to global minima Ye & Du (021))
Convex neural networks. There is long history of architecting convex optimization problems that mimic the performance of neural networks Zhang et al. 20166) 2017). It has recently been found that ReLU neural networks have equivalent convex programs for training, such networks with scalar outputs (Pilanci & Ergen 2020), vector-outputs (Sahiner et al. 2022 convolutional networks (Ergen & Pilanci 2020 Sahiner al. 2021c), nolymomial-activation networks (Bartan & Pilanci 2021). batch-norm based networks (Ergen et al. 2021), Wasserstein GANs (Sahiner et al. 2021a)) and self-attention networks (Sahiner al.) 2022). Despite efforts in developing efficient sol solvers. convex networks are only effectively trainable at small scales (Bai et al. [0022 Mishkin et al., 2022). Our novelty to adapt BM factorization as a fast and scalable solution for training convex networks, with simple, verifiable conditions for global optimality.
2 PRELIMINARIES
We denote (+)+ == max } as the ReLU non-linearity. We use superscripts, say A (1rif)), to denote blocks of matrices, and brackets, say A 12]. denote elements of matrices. We let be the vector of ones of appropriate size, +|| |) be the E-p norm with p = H, || 1||F be the Frobenius norm, and ви be the unit H-norm ball, {u: |||l 1).. Unless otherwise stated, let [ be convex, differentiable function. We use to denote the number of samples, and to denote the output dimension of each
network. All proofs are presented in Appendix
2.1 TWO-LAYER NEURAL NETWORKS AS CONVEX PROGRAMS
A line of work has demonstrated that two-layer neural networks are equivalent to convex optimization problems. We consider data matrix E Ruxd and consider two-layer o-activation network with e outputs, m neurons, weight-decay parameter 3 >> 0:
3 (I[11) 12 |[
PMLP = min ((o(XW,)W,)) +
(1)
WI W. W.
2
When is a linear activation and m 2 m² for some m S min c). this problem is equivalent to (Rennie Srebro 2005), Section 2.2)
PLMLP-
min F(XZ) + +
ZE
whereas for a ReLU activation and m 2 for some unknown, moblem-dependee m S nc ( (Sahiner et al. (021b), Thm. 3.1),
PRMLP=
min F DD + Zz l.,,, K, (2D, -,,X
Z
where (D,),- {diag (1 0)) E Rd) enumerates the possible activation patterns
generated from X, and the number of such patterns satisfies P P 2r
(<(n11) where r == rank
(Stanley et al. 2004 Pilanci Ergen 2020). The expression (3) also involves a constrained nuclear norm expression, which is defined as
Z K.K
min / s.t. Z IC, C: conv (Z uv Ku 0, ||lla 1, 12 1). F20
(4)
This norm is quasi-nuclear norm, which differs from the standard nuclear norm in that the factoriza- tion upon which relies imposes constraint on its left factors. In convex ReLU neural networks, this norm enforces the existence of {u,, V&} such that Z Σ. u.v. and D .Z > k. (Xu)+v and penalizes Σx / T Il.. This norm is NP-hard compute Sahiner et al. 2021 A variant of these ReLU activations, called gated ReLU activations, achieves the piecewise linearity of ReLU activations without enforcing the constraints (Fiat et al. 2019). Specifically, the ReLU gates are fixed to some (hy}/=n to form
o(W1))) diag (1 (1/h 0)) (Xw1)) D,Xwij.
3

---

Published as a conference paper at ICLR 2024
With gated ReLU activation, the equivalent convex program is given by ([Sahiner et al. 2022
P
POMLP
min F F D,X,, + Z
7 l+:
(5)
which thereby converts the constrained nuclear norm penalty to standard nuclear norm penalty, improving the complexity of the ReLU network. In addition the multi-layer perceptron (MLP) formulation, two-layer Rel.U-activatiog convolutional neural networks (CNNs) with global average
pooling have been demonstrated he equivalent to convex programs as well (Sahiner et al. 2021b Ergen & Pilanci 2020). The non-convex formulation is given by
8 w1, + | 2 ==1
PRONN = min WiyER* Way ERR
Σ
W2j 1 (X(w1s)+) + Σ ==1
(6)
where samples X; E Rx are represented by patch matrices, which hold a convolutional patch of size h in each of their K rows. It has been shown (Sahiner et (021b) that as long as m 2 m" where m S ne, this is equivalent to a convex program ((Sahiner et al. (0221)) Cor. 5.1)
P
PRENN= = min
Σ F(C
1 L
x +
Z lllk,
Z
1=1
==1
j=1
X1
K₁ = (2D; Ink) ),,x))X X, X =
(7)
X
where (Dy} {diag (Xu 2 0}): E RR) and DJ E RRKK The only exponential
dependence of P h. which typically fixed.
Lastly. we review existing convexity results for self-attention transformers (Sahiner al. 2022). We have the following non-convex objective for single block of multi-head self-attention with m heads, where X; E Ra with s tokens and d features
m 3
PSA
min Σr
Σ o (X X.W.
|| 1
W WiJERdad WayeR
for which variety of objectives F can be posed, including classification (e.g. F incorporates global average pooling followed by softmax-cross cnosscyy with labels) or denoising (e.g. F is squared loss). For linear. gated ReLU, and ReLU activation, this is equivalent to a convex program (see Appendices A. B.3 Here, we show that these architectures are amenable the BM factorization.
2.2 THE PURER-MONTERR FACTORIZATION
First proposed by Burer & Monteiro 2003). the Burer-Monteiro (BM) factorization proposes to solve SDPs over some square matrix Q in terms of rectangular factors R where Q is substituted by RR. T It was first demonstrated that solving over R does not introduce spurious local minima rank-constrained SDPs, given rank(R) 2 rank(Q") for optimal solution to the original SDP Q" (Burer & Monteiro 2005). We seek applications where optimize
non-square matrix Z, i.e.
PCvx min F(Z)
ZeRaxe
for a convex, differentiable function F. One may approach this by factoring UVT where U E Rdxm V E Rexm for some arbitrary choice m. Then, we have an equivalent non-convex
problem over R = for (RR = F(UV
PCvx min (R).. R
(10)
TUT UV VUT VVT one may apply directly the result of
Noting that 9 is convex over RR T
Boumal et al. 2020) conclude that long m > d + c, all local minima of (1)) are global
4

---

Published as a conference paper at ICLR 2024
minima of (9) (see Appendix A.2 Further, work from Bach al. (2008) and Haeffele et al. (2014)
demonstrates that all sank-deficient local minimizers of (0 achieve the global minimum Pcv x
long line of work has analyzed the conditions where known non-convex optimization al- gorithms will converge second-order critical points (local minima) (Ge et al. 2013 Jin et al. 2017] Daneshmand et al. 20118). Under the assumption of bounded f and its Hessian, a sccond-orids
critical point can be found by noisy gradient descent (Ge al.) 2013). or other second-order algorithms (Sun et al., 2015). Even vanilla gradient descent with random initialization has been demonstrated almost surely converge to a local minimum for f with Lipschitz gradient (Lee et al.) 2016). However, if the gradient of f is not Lipschitz-continnont descent will find a second-order critical point (10): one may encounter stationary point which saddle. For example, in the linear regression setting, i.e.
there are no guarantees that gradient
f(R) XUV -YI
(11)
the gradient of f is Lipschitz continuous with respect U when is fixed and versa, but not Lipschitz continuous with respect to R (Mukkamala & Ochs 2019). Thus, one may not directly
apply the results Ge et al. (2133); [Sun et al. (2015): Lee al. (2016) in this case, Instead, we seek
to understand the conditions under which stationary points (1O) correspond to global optima of
One such condition is given in Mardami et al. (2013 2015).
Theorem 2.1 (From (Mardani et al.
2013 ). Stationary points U, V of the optimization problem
3 2 (1U + IV]
min U.V 2
JVT YII
(12)
correspond to global optima z = UVT of the equivalent convex optimization problem
p²
min Z
-
(13)
provided that Y ÜV 22 3
3 BURER-MONTEIRO FACTORIZATION FOR CONVEX NEURAL NETWORKS
3.1 MLPs
We first seek to compare the convex formulations of the MLP training problem 2 3 and 5 to their BM factorizations. We describe how to find the BM factorization for any convex MLP.
Lemma 3.1. For any matrix M gaxde let f(U, V) == F(MUVT) be a differentiable function. For any B > 0 and arbitrary vector norms II- and Illlc. we define the Burer-Monteiro factorization
p²:
min S(U, +
Σ
lulll: + llv,
(14)
UCRdexm VeR
For the matrix norm || || defined
| == max trace(R T Z) s.t. u Rv $1 Bo, Vy E Br.
(15)
R
the problem
14 is equivalent to the convex optimization problem
F(MZ) + | Ze
min
(16)
Remark 3.2. In the case of linear MLP, M - X, d, d, d, and 1l ||. ... using the = 2 and = 2 (Bach al. 2008
definition of ||- || the corresponding BM factorization, For gated ReLU network, the regularizer is still the nuclear norm, and thus the same R c= regularization appears the BM factorization. In the case of the ReLU MLP, the nuclear norm is
Under mild conditions, see Appendix A.3
5

---

Published as a conference paper at ICLR 2024
replaced by ||. Ei= 1 1 [lx,k,, which in the BM factorization amounts to having the constraint
K, K 2 O. We express the BM factorization of convex MIPs below.
B (IUIE [ 2
PLMLP = min F(XUVT)
(17)
UERdxm VeR
PGMLP
min UJERuAM
F
F DJ) 8 2
Σ
(IU IVV
(18)
VyEReXm
j=1
j=1
8
PRMLP-
min USER (2 VJERoXM
F _
D2XU,
Σ +I +IV
(19)
2
j=1
To the best of our knowledge, (19) presents the first application of BM factorization to a non-linear neural network, which is enabled by the convex model 33.
In the linear case, the BM factorization 17 is identical to the original non-convex formulation of a linear MIP with m neurons. In the case of gated ReLU, the BM factorization when m is equivalent to the original non-convex formulation. However, for ReLU activation two-layer networks, the BM factorization even when m = corresponds to a different (i.c. constrained, rather than ReLU activation) model than the non-convex formulation. While the original convex program is NP-hard due to the quasi-nuclear norm (Sahiner et al. 20216). the per-iteration complexity of the BM factorization is much lower than for the convex ReLU MLP.
The BM factorizations of these convex MLPs non-convex, hence finding global min- imum appears intractable. However, the following theorem demonstrates that as long as a rank-deficient local minimum to the BM factorization is obtained, it corresponds to global optimum. Theorem 3.3. If m 2 de + d,, all local minima of the BM factorization (4) are global minima.
Furthermore, if is twice- shice-ifferentiable any refk-deficient local minimum R :
of (14
V
corresponds a global minimizer Z" = TÜV 16
This result demonstrates that these two-layer convex MLPs have no spurious local minima under mild conditions. However, there remains algorithmic challenge: it is not straightforward to obtain a guaranteed local minima when the gradients are not Lipschitz continuous. The following result provides general condition under which stationary points of the (14) are global optima of (16).
Theorem 3.4. For any non-negative objective function F, for stationary (Ú.V) of 14) with corresponding 2 UVT with objective K for 16). the relative optimality gap p-p" satisfies p
- p
|ZZF(M2)15 -
(20)
p"
where || Ilb is the dual norm of ||
This bound can be calculated by taking the gradient the unregularized objective function, evaluated at candidate solution to the convex problem (16). which formed by the stationary point of BM problem (14). In the case of linear MLP with X La. F squared- objective, and 2 (MZ) I| I 3,, our result exactly recovers the result of Theorem 2.1 from Mardani al. (003)). When this condition is not exactly satisfied, (20) provides a novel result in the form of an optimality gap bound. To knowledge. this is the first result that generalizes the optimality conditions for stationary points from any BM factorization of a neural network. This provides an easily computable bound after solving (14) which quantifies how close solution is to the global minimum. In the case of ReLU MI.,, the relative optimality gap is given by
$ p"
S
max
Zz F
Σ ===
Dj. XZ, )ulp
p"
jEiP) uEB,
K, K20

---

Published as a conference paper at ICLR 2024
Computing this quantity requires solving cone-constrained PCA problem (Deshpande et al. 2014). In certain cases, the optimality gap stationary points (20) is always zero.
Theorem 3.5. A stationary point (U, V) of [4) is global minimizer of 16 if R
C and
rank (Ü) = rank (V) = minnde, dr).
(21)
Thus, for linear and gated ReLU MLPs, achieves a stationary point with full rank, is corresponds with the global optimum of the convex program. We now further extend these results CNNs and self-attention architectures.
ensure that if the Burer-Monteiro factorization
3.2 CNNS
Before proceeding to explore the BM factorization in the context of two-layer CNNs, we first provide a new result on an equivalent convex program two-layer ReLU CNNs with arbitrary linear pooling operations, which extends the results of [Sahiner et al. (021bb; Ergen & Pilanci] (2020) on Global Average Pooling CNNs. Define P. E gax* to be a linear pooling matrix which pools the K spatial dimensions to arbitrary size Then, express the non-convex two-layer CNN problem
PONN == min
Σn 1=1
Σ
W2 W2, P Pa (X.Wi)))))
(22)
WIJER" Way EB
Theorem 3.6. For 3 and ReLU activation o(.) = (.)+, if 2 m where m S nac, then equivalent to convex optimization problem, given by
22
trace P.D
x (1)
PONN
min Z ZiERbxec
Z |l,,
(23)
trace(P (P.D. x (6)
X1
Kg = (2Ds Ink)
z E Rhxxaa Vcl [c].
X,
Thus, we provide a novel result which characterizes two-layer CNNs with arbitrary linear pooling operations convex program. Similar results be shown for the linear and gated-ReLU activation case With this established, we present our main results on the BM factorization for CNNs.
Lemma 3.7. The BM factorization of the convex CNN problem with ReLU activation is given as:
m Σ
m
PRONN= min U VaeeReXa
ΣF F
jk J P.D. X,Uj
Σ
(1u,, + [VVx+11))
2
k=1j=1
k=1j=1
(24)
s.t. (2D I)X,u 2 OAj E [m], Ak E [P]
The BM factorization closely resembles the original non-convex formulation (2 Generally, 24
inherits the results of Theorems 3.3 and
3.5 we present one such corollary here.
Corollary 3.8. A stationary point (ūjz Uiys, L Ei
1 of 24 corresponds global minimizer
vec
(V
jk
of
23
provided that
trace P. (.
Vz.
u//2 SB
Vk E [P] Be: (2D
I)X 2 0.
trace P.D x
(25)
We examine linear and gated ReI.U activations for CNNs in Appendix B.2
7

---

Published as a conference paper at ICLR 2024
3.3 MULTI-HEAD SELF-ATTENTION
We now for the first time extend BM factorization theory to self-attention networks. Lemma 3.9. The BM factorization of the convex linear-activatioon welf---ttetionee problem given as:
m Σ
8
PLSA
min
F
XXU,XXXVVV
U,l + +V
(26)
UYERaEd
VjeRdxe
In addition to inheriting all of the results of Theorems 3.3 3.4 and 3.5 noting the equivalence the the first to show conditions under
BM factorization with the original non-convex program 8). which there are no spurious local minima for self attentionnio networks.
Corollary 3.10. The iinearactivation self-attention network (26) has no spurious local minima as long the number of heads satisfies twice wwice-differentaable objective F. if for any local minimum (U,, v of 26). the matrix
mt" where m" df + dc. Furthermore, for any
veo((U1)
vec(Ūm) voc(Vm))
R ==
E [d(d+o)xm
(27)
[voc(V1
is rank-deficien then this local minimum is also global minimum of 104
4 EXPERIMENTAL RESULTS
Fashion-MNISS MIIST Accuracy
92
80
m=
m=1
CIFAR 10 Test Accuracy
m-2
m=2
90
75
88
80
1
2 Stage
Stage
(a) CIFAR-10
(b) Fashion-MNIST
Figure I: BM enables layerwise training of convex ReLU CNNs, which competitive with end-to-end ReLU networks the same depth. CIFAR-10, with m 4 we achieve a test accuracy of 81.5% compared 81.6% end-to-end non-convex training, and Fashion-MNIST. with m we achieve test accuracy of 91.6% compared to 91.2% for non-convex training (Kiliçarslan Celik 2021
Bhatnagar et al. 2017).
4.1 THE RELATIVE OPTIMALITY GAP BOUND
In this section, we illustrate the utility of our proposed relative optimality bound for stationary points in the setting of two-layer fully-connected networks. We also seek to examine how this bound changes with respect to the number of samples n, the regularization parameter /3 (which controls the sparsity of the convex solution), and the number of factors in the BM factorization m. We initialize class-balanced three-class spiral data set with varied number of samples n (see Figure Appendix C.I for examples). For this dataset, factorization (18) with varying number of factors m. We then compare the stationary points of these BM factorizations found by gradient descent (GD) the global optimum, which we compute from (5).
then train the gated ReLU MLP BM
For each stationary point of the BM factorization, we compute the relative optimality gap bound provided in our result in Theorem 3.4 We note that since d 2, 3 in this case, for all j. as long as m > all local minima of the BM factorization are global minima (Burer & Monteiro 2005 Haeffele al. 2014). While Lee al. 2016) demonstrated that gradient descent
We examine gated ReI.U and ReL.U activations for self-attention Appendix
B.3
8

---

Published as a conference paper at ICLR 2024
with random initialization converges to local optimum almost surely for losses whose gradient is Lipschitz continuous, we use squared loss with one-hot-encoder class labels, which is not Lipschitz continuous (Mukkamala & Ochs 2019). Thus, there is guarantee that GD will find the global minimum. We display results over / in Figure 2 Our bound gives a useful proxy for whether the BM factorization converges to the global minimum. For larger values of B, it be- comes much easier for GD to find an optimal solution, but GD almost never finds the global minimum.
We find that GD applied to the BM factorization finds saddle points that are not quite local minima, but close. Interestingly, there is only a minor relationship between the optimality gap and the rank of the BM factorization m. This experiment further validates the need to consider stationary points of the BM factorization, rather than just local minima.
4.2 BM ENABLES LAYERWISE TRAINING OF CONVEX CNNS
We consider the task of leveraging the theory of two-layer convex ReLU neural networks for training deep image classifiers. Following the
10-1
Relative Optimality Gap
approach of (Belilovsky et al., 2201)), we seek to train two-layer convex CNNs greedily to mimick the performance of deep network. In the non-
10-3
10-
10-
T
Bound, m=1
convex setting, the greedy approach proceeds by training single two-layer CNN, then freezing the weights of this CNN, using the latent repre- sentation of this CNN as the features for another two-layer CNN, and repeating this process for specified number of stages. We leverage the result of Theorem 3.6 to convert this non-convex layerwise training procedure to convex one, training stages of convex two-layer gated ReLU CNNs. We apply this procedure the CIFAR-IO Krizhevsky et al.] (009) and Fashion-MNIST (Xiao et al., (207)) datasets, using the architec- ture of (Belilovsky et al. 2019) (see Appendix
10-9
Bound, m=5
20-11
Actual, m-1
Actual. m-2
10
m
10-4
10-2
10-2
10-1
beta
Figure 2: Relative optimality gap of the non- convex BM factorization of a rated-ReLU two- layer MLP for three-class spiral data classification (n = 150, d = 2, c = 3). demonstrate how 3 and m affect relative optimality gap, both in in terms of the proposed bound and the actual gap, where the global minimum is found by convex op- timization.
In a memory-limited setting, layerwise training with the convex model equation 23 is im- possible, because the latent representation to be used as input for the second stage, given by has Pac channels, which for reasonable choices of P = 256, = 4, 10 yields upwards of 10 channels for the input to the second CNN stage. Accordingly, we employ the BM factorization of size m, so the latent representation only consists mP channels.
(DD) X (c ) 1
Figure demonstrates that this BM scheme for layerwise training allows for performance to improve stage the next of the layerwise training procedure, reaching the performance of much deeper networks while enabling a convex optimization procedure. Training five stages of a BM factorized convex two-layer gated ReLU CNN CIFAR- 10 resulted in a final test accuracy of 80.9%, 81.1% 81.1%, and 81.5% for m E [1, 2, 4] respectively. Previously, it has been demonstrated that a six-layer ReLU CNN achieves 81.6% on CIFAR-I0 when trained end-to-end (Kiliçarslan & Celik 2021). Interestingly, we find that increasing m generally improves performance. The three-stage trained BM factorized convex two-layer gated ReLU CNN on Rasion-MNIST achieved a final test accuracy of 91.5%, 91.4%, and 91.6% for m [1. 2, 4) respectively, compared to 91.2% for a four-layer ReLU CNN trained end-to-end (Bhatnagar et al. 2017). On this dataset, we observe that the impact of increasing
performance is less pronounced than for CIFAR-10.
The BM factorization is essential for convex neural networks to match deep ReLU net- works. 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. While inheriting the guarantees of Theorems 3.4, and 3.5] these layerwise trained BM networks match the performance of end-to-end. highly non-convex ReLU deep networks.

---

Published as a conference paper at ICLR 2024
REFERENCES
Erling D Andersen and Knud D Andersen. The mosek interior point optimizer for linear programming: an implementation of the homogeneous algorithm. In High performance optimization, pp. 197-232. Springer, 2000.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. International Conference Machine Learning, pp.
PMLR, 2019.
Francis Bach, Julien Mairal, and Jean Ponce. Convex sparse matrix factorizations. arXiv preprint arXiv:0812.1869, 2008.
Yatong Bai, Tanmay Gautam, and Somayeh Sojoudi. Efficient global optimization of two-layer relu networks: Quadratic-hie algorithms and adversarial training. arXiv preprint arXiv:2201.0196. 2022.
Burak Bartan and Mert Pilanci. Neural spectrahedra and semidefinite lifts: Global convex op- timization of polynomial activation neural networks in fully polynomial-time. arXiv preprint 07Xi: 2021.
Eugene Belilovsky, Michael Eickenberg. and Edouard Oyallon. Greedy layerwise learning can scale to imagenet. In International conference machine learning, pp. 583-593. PMLR, 2019.
Shobhit Bhatnagar, Deepanway Ghosal, and Maheshkumar H Kolekar. Classification of fashion article images using convolutional neural networks. In 2017 Fourth International Conference on Image Information Processing (ICIIP), pp. 1-6. IEEE, 2017.
Nicolas Boumal, Vlad Voroninski, and Afonso Bandeira. The non-convex burer-monteiro approach works smooth semidefinite programs. Advances in Neural Information Processing Systems, 29, 2016.
Nicolas Boumal, Vladislav Voroninski, and Afonso Bandeira. Deterministic guarantees for burer- monteiro factorizations of smooth semidefinite programs. Communications on Pure and Applied Mathematics, 73(3):581-608, 2020.
Samuel Burer and Renato DC Monteiro. A nonlinear programming algorithm for solving semidefinite programs via low-rank factorization. Mathematical Programming, 95(2):329 -37,, 2003.
Samuel Burer and Renato DC Montciro. Local minima and convergence in low-rank semidefinite programming. Mathematical programming, 103(3):427-444, 2005.
Ricardo Cabral, Fernando De la Torre, João Costeira, and Alexandre Bernardino. Unifying nuclear norm and bilinear factorization approaches for low-rank matrix decomposition. In Proceedings of the IEEE international conference on computer vision, pp. 2488 2498, 2013.
Emmanuel Candès and Benjamin Recht. Exact matrix completion via convex optimization. Foun- 2009.
dations of Computational mathematics,
Emmanuel Candès and Terence Tao. The power of convex relaxation: Near-optimal matrix
completion. JEEE Transactions on Information Theory, 56(5):2053-22222 200000055-000,,,0000000000 2010.
Diego Cifuentes and Ankur Moitra. Polynomial time guarantees for the burer-monteiro method. arXiv preprint arXiv: 2019.
Hadi Daneshmand, Jonas Kohler, Aurelien I Lucchi, and Thomas Hofmann. Escaping saddles with stochastic gradients. International Conference on Machine Learning, pp. 1155-1164, PMLR, 2018.
Yash Deshpande, Andrea Montanari, and Emile Richard. Cone-constrained principal component analysis. Advances Neural Information Processing Systems, 27, 2014.
Steven Diamond and Stephen Boyd. Cvxpy: A python-embeddeed modeling language for convex optimization. The Journal of Machine Learning Research, 17(1):2909-29) 2016.
10

---

Published as a conference paper at ICLR 2024
Laurent Dinh, Jascha Sohl-Dicksteins and Samy Bengio. Density estimation using real nvp. arXiv preprint arXiv:1605.0880 08803, 2016.
Simon Du and Jason Lee. On the power of over-paramttrizatioo in neural networks with quadratic activation. International conference on muchine learning, pp. 1329-1338, PMLR, 2018.
Simon Du, Jason Lee, Yuandong Tian, Aarti Singh, and Barnabas Poczos. Gradient descent learns one-hiddennlaye layer enn: Don't be afraid of spurious local minima. In International Conference on Machine Learning, pp. 1339-1344 1348. PMLR, 2018.
Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International conference on machine learning, pp. 1675-1685 PMLR, 2019.
Simon S Du, Chi Jin, Jason D Lee, Michael Jordan, Aarti Singh, and Barnabas Poczos. Gradient descent can take exponential time to escape saddle points. Advances in neural information processing systems, 30, 2017.
Murat A Erdogdu, Asuman Ozdaglar, Pablo A Parrilo, and Nuri Denizean Vanli. Convergence rate of coordiccoe maximization burer- monteiro method for solving large sdps. Mathematical Programming, pp. 1-39, 2021.
Tolga Ergen and Mert Pilanci. Implicit convex regularizers of emn architectures: Convex optimization of two-and three-layer networks in polynomial time. In International Conference Learning Representations, 2020.
Tolga Ergen, Arda Sahiner, Batu Ozturkler, John M Pauly, Morteza Mardani, and Mert Pilanci. Demystifying batch normalization in relu networks: Equivalent convex optimization models and implicit regularization. In International Conference on Leurning Representations, 2021.
Jonathan Fiat, Eran Malach, and Shai Shalev-Shward Decoupling gating from linearity. arXiv
preprint arXiv:1906.05000 2019.
Rong Ge. Furong Huang, Chi Jin, and Yang Yuan. Escaping from saddle points-online stochastic gradient for tensor decomposition. In Conference on learning theory, pp. 797-842. PMLR, 2015.
Rong Ge, Chi Jin, and Yi Zheng. No spurious local minima in nonconvex low rank problems: A unified geometric analysis. In International Conference on Machine Learning, pp. PMLR, 2017.
Nicolas Gillis, Introduction to nonnegative matrix factorization. arXiv preprint arXiv: 2017.
Benjamin Haeffele, Eric Young, and Rene Vidal. Structured low-rank matrix factorization: Optimality, algorithm, and applications to image processing. In International conference on machine learning. pp. 2007-2015. PMLR, 2014.
Benjamin D Haeffele and René Vidal. Global optimality in neural network training. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 7331-7339, 2017.
Chi Jin, Rong Ge, Praneeth Netrapalli, Sham M Kakade, and Michael Jordan. How to escape saddle points efficiently. In International Conference on Machine Learning, pp. 1724-1732. PMLR, 2017.
Kenji Kawaguchi. Deep learning without poor local minima. Advances in neural information processing systems, 29, 2016.
Serhat Kiliçarslan and Mete Celik. Rsigelu: A nonlinear activation function for deep neural networks. Expert Systems with Applications, 174:11480; 14805, 2021.
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. 2009.
Jason D Lee, Max Simchowitz, Michael Jordan, and Benjamin Recht. Gradient descent only converges minimizers. In Conference learning theory, pp. 1246-1257. PMLR, 2016.

---

Published as a conference paper at ICLR 2024
Jan R Magnus and Heinz Neudecker. Matrix differential calculus with applications in statistics and econometrics. John Wiley & Sons, 2019.
Morteza Mardani, Gonzalo Mateos, and Georgios B Giannakis. Decentralized rank minimization: Algorithms and applications. IEEE Transactions on Signal Processing, 61(21): 5374-5388, 2013.
Morteza Mardani, Gonzalo Matcos, and Georgios B Giannakis. Subspace learning and imputation for streaming big data matrices and tensors. IEEE Transactions on Signal Processing, 63(10):
2663-2677, 2015.
Aaron Mishkin, Arda Sahiner, and Mert Pilanci. Fast convex optimization for two-layer relu networks: Equivalent model classes and cone decompositions. In Proceedings of the 39th International Conference Machine Learning, 2022.
Mahesh Chandra Mukkamala and Peter Ochs. Beyond alternating updates for matrix factorization with inertial bregman proximal gradient algorithms. Advances in Neural Information Processing Systems, 32, 2019.
Dohyung Park, Anastasios Kyrillidis, Constantine Carmanis, and Sujay Sanghavi. Non-square matrix sensing without spurious local minima via the burer-monteiro approach. In Artificial Intelligence and Statistics, pp. 65-74. PMLR, 2017.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32: 8026-80377 2019.
Mert Pilanci and Tolga Ergen. Neural networks are convex regularizers: Exact polynomial time convex optimization formulations for two-layer networks. In International Conference on Machine Learning, pp. 7695-7705. PMLR, 2020.
Benjamin Recht, Maryam Fazel, and Pablo Parrilo. Guaranteed minimum-rank solutions of linear matrix equations via nuclear norm minimization. SIAM review, 52(3):471-501, 2010.
Jasson DM Rennie and Nathan Srebro. Fast maximum margin matrix factorization for collaborative prediction. In Proceedings of the 22nd international conference on Machine learning, pp. 713-719, 2005.
Arda Sahiner, Tolga Ergen, Batu Ozturkler, Burak Bartan, John M Pauly, Morteza Mardani, and Mert Pilanci. Hidden convexity wasserstein gans: Interpretable generative models with closed-form solutions. In International Conference Learning Representations, 2021a.
Arda Sahiner, Tolga Ergen, John Pauly, and Mert Pilanci. Vector-output relu neural network problems copositive programs: Convex analysis of two layer networks and polynomial-time algorithms. In ICLR, 2021b.
Arda Sahiner, Morteza Mardani, Batu Ozturkler, Mert Pilanci, and John M Pauly. Convex regulariza- tion behind neural reconstruction. ICLR, 2021c.
Arda Sahiner, Tolga Ergen, Batu Ozturkler, John Pauly, Morteza Mardani, and Mert Pilanci. Un- raveling attention via convex duality: Analysis and interpretations of vision transformers. In Proceedings of the 39th International Conference on Machine Learning, 2022.
Alexander Shapiro. Semi-infinite programming, duality, discretization and optimality conditions.
Optimization, 58(2):333116) 2009.
Richard P Stanley et al. An introduction to hyperplane arrangements. Geometric combinatorics, 13 (389-496):24, 2004.
Ju Sun, Qing Qu, and John Wright. When are nonconvex problems not scary? arXiv preprint 15 2015.
12

---

Published as a conference paper at ICLR 2024
Irene Waldspurger and Alden Waters. Rank optimality for the journal on Optimization, 30())2577-2602, 2020.
factorization. SIAM
Lingxiao Wang, Xiao Zhang, and Quanquan Gu. A unified computational and statistical framework for nonconvex low-rank matrix estimation. In Artificial Intelligence and Statistics, pp. 981-990. PMLR, 2017.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms. arXiv preprint
07100000777, 2017.
Tian Ye and Simon S Du. Global convergence of gradient descent for asymmetric low-rank matrix factorization. Advances in Neural Information Processing Systems, 34, 2021.
Gavin Zhang, Salar Fattahi, and Richard Y Zhang. Preconditioned gradient descent for overparame- terized nonconvex montero factorization with global optimality certification. J. Mach. Learn. Res., 24:163-1, 2023.
Richard Y Zhang. Improved global guarantees for the nonconvex murer-monteiro factorization via rank overparamettrriatton. arXiv preprin! arXiv:220..01789, 2022.
Yuchen Zhang, Jason D Lee, and Michael Jordan. 11-regularized neural networks are improperly learnable polynomial time. International Conference on Machine Learning, pp. 993-1001. PMLR, 2016.
Yuchen Zhang, Percy Liang, and Martin J Wainwright. Convexified convolutional neural networks. In International Conference on Machine Learning, pp. 4044-4053, PMLR, 2017.
Qinqing Zheng and John Lafferty. Convergence analysis for rectangular matrix completion using burer-monteiro factorization and gradient descent. arXiv preprint arXiv:1605.070511 2016.
Zhihui Zhu, Qiuwei Li, Gongguo Tang, and Michael Wakin. The global optimization geometry of low-rank matrix optimization. arXiv preprint arXiv: 01266, 2017.
Liu Ziyin, Botao Li, James B Simon, and Masahito Ueda. Sgd converge local maxima. In International Conference on Learning Representations, 2021.
13

---

Published as a conference paper at ICLR 2024
A PROOFS
A.1
COMPLEXITY AND CONVERGENCE RESULTS OF THE BM FACTORIZATION
In this subsection, provide some additional details on the complexity and convergence of the BM factorization 19 We consider the problem dimensions of X E guxd BM factorization size m, and network output dimension c.
First, we can compare the per-iteration time complexity of a naive implementation of pro- jected gradient descent on (19) to the the per-iteration time complexity of a Frank- Wolfe algorithm designed to solve (3) as proposed Sahiner al. (202Tb). Consider the case of squared loss, e-g. Y = 2 Y YIg. For r == rank(X (X), shown Appendix A.4.2 of Sahiner et al. Ib) that single iteration of Frank-Wolfe has time complexity O(Pr"), where P is the number of hyperplane arrangements.
In comparison, projected gradient descent applied to 19 takes two steps:
1. Gradient descent steps Vy)ff The gradient of the objective of (9) with respect to Uj is given by
(D,X)
Σ Dj XU - Y V + 3U,,
(28)
Noting that the residual R := E XU Y can be computed once in O(P(n'c + nde + dem)) time and re-used in all gradient computations, the total complexity of computing all gradients and taking gradient descent step is given by (P(n2(d c) + ndc dem)).
2. A projection step of U, onto the space (2D, - XU, 3 2 for each j E [P]- this naively can be solved in O(P(n* + dm))m)) time.
Overall, then, using projected gradient descent on O(P(n + n²(d + dm) + dem)). which unless m exponential in r significantly more computationally efficient than solving 3 with Frank- Wolfe in O(Pn") time.
is of time complexity
For the convergence rates of naive gradient descent applied to BM factorizations and (19), we point to Zhang et al. (2023), which suggests that these formulations may converge sub-linearly to stationary point in the case that exceeds the rank of the optimal solution Z However, Zhang et al. (2023) suggest a pre-conditioning method to gradient descent applied to BM formulations which guarantees linear convergence and can be applied here.
A.2 RESULT OF (BURER & MONTEIRO 2005) AND ITS APPLICATIONS TO RECTANGULAR MATRICES
In this subsection, outline the precise theoretical statement of (Burer & Monteiro 2005 and describe exactly how it corresponds to our summary in Section 2.2 and thus the application to the later derivations in our work. We first describe the following result, without proof, from (Burer & Monteiro 2005).
Lemma A.1 (Lemma 2.1 of (Burer Monteiro, 2005 Suppose R. (dd+))xr S E R(dtc)xr
satisfy R. = TS Then, S RQQ.Q.. for some orthogonal Q E Brrx.
Now, we proceed to prove analogs of Theorem 2.3 of (Burer & Monteiro 2005) for general SDPs with rank constraint.
Lemma A.2 (Analog of lemma 2.2 of (Burer Monteiro
2005
Consider the problem
min R
F'RRR.
(29)
R local minimum of 29 if and only if RQ is a local minimum of 29) for all orthogonal Q E Rixr.
14

---

Published as a conference paper at ICLR 2024
Proof. Since Q is orthogonal, (RQ) RQQ same objective value, gradients, and higher order derivatives R. Thus, R a local minimum if and only R/ is local minimum.
RR Thus, R' := RQ attains the
Theorem A.3 (Analog of Theorem 2.3 of (Burer & Monteiro 2005 problems.
Consider the following two
X FX(X),
min
(30)
X==
tank())Sr
min RERI
F'RRR
(31)
Then, for any continuous function F', feasible solution X is local minimizer of 30 if and only if. for X RR R is local minimizer of 31
Proof. We follow the exact same lines (Burer & Monteiro 2005). By continuity of the map
R RR T we know that if X is local minimizer of (30). then R is local minimizer of 31
Now, must prove the other direction, namely that if X RR is not local minimizer of (30 then R a local minimizer of (31
Suppose that X is not local minimum. By continuity of F' then, there must be se- quence of feasible solutions {X*} of (30) converging to X such that F'(X*) < F for all For each k, choose R$ such that X R Since {X*] is bounded, it follows that (R*) is bounded and hence has subsequence (R* KER converging some R such that X RR T. Since
F' F'(R***) Using the fact that X not a local minimum
F* < F(X) = (RR ), we see that R is not a local minimum of (1).
RR RR together with Lemmas A.1 and A.2 conclude that R is 31
With this established, we now describe the setting described in Section 2.2 i.e. the rectangular matrix, non-SDP case.
Lemma A.4. Consider the optimization problems
p
XER( F'(X),
min
(32)
Xx0 rank rank(X)mm*
and
P2 = min F(Z). ZeRaxe
(33)
where convex. Define F²; (dd+c)x - IR such that
X
X,
F(
F(X2)
(34)
X
X4
for X1 E Rdxd X2 R X2 E Bexd X E Roxe Then, any local minimizer X
32
corresponds global minimizer Z X, of
33 for m S d
Proof. This follows from Boumal et al. (2020). Corollary 3.2.
Lemma A.5 (Used in Section 2.2 For any convex, continuous function F, a local minimum (Ù, V)
of
min
F(UV)
(35)
UERdxm V
corresponds global minimum Z" U of 33 long m m² for some m² < d + c.
15

---

Published as a conference paper at ICLR 2024
Proof. Define F', : [(d+))x(dte)
R such that
X1 X3
X2 F(X2)
F((
=-
(36)
X4
for X, Rddd X2 Rdxe, X3 Rexd X, E Rexx Then, let
R(d+))x One can
V
re-write F(UV) as (RRR"), ). Then, we see that 35 can be expressed as 31 and any local
minimizer to 35 is local minimizer
Next, note that from Theorem A.3 that any local minimizer of 31 local minimizer of
(30). Lastly, note from Lemma A.4 that any local minimizer of [30 corresponds to a global
minimizer of (33) for m² d Thus, any local minimizer of 35 is global minimizer of 33 as
long as m 2 m where m < c.
A.3 RESULT OF (HAEFFELE ET AL. 2014) AND ITS APPLICATION TO RECTANGULAR MATRICES
In this subsection, outline the precise theoretical statement of (Haeffele al. 2014) and describe exactly how corresponds to our summary in Section 2.2 and thus the application to the later derivations in our work. We first deseribe the following result, without proof, from (Haeffele et al.) 2014).
Theorem A.6 (Theorem (Haeffele et al. 2014)) Fee : St - R be of the form such that (X (X) + (X),, where G' is convex, twice differentiable with compact level sets, and H' proper convex function such that F' is lower semi-continuous. Then, if R refi-eef local minimizer of
min
F(RR
(37)
X
then X RR is a global minimizer of
min XEa(d+eix X
F'(X).
(38)
We now describe how this theorem applies to the setting described in Section 2.2 i.e. the rectangular matrix, non-SIDP case.
Lemma A.7 (Used in Section 2.2h. Let F Raxe R be of the form such that F(Z) G(Z) HZ),, where G is convex, twice differentiable with compact level sets, and H is a proper convex
Ü
function such that F is lower semi-continuous. Then, if R
is a rank-deficient local minimizer
V
of
min F(UVT), 4e VeRdxm
(39)
then Z" = T is global minimizer of
min F(Z). z
(40)
Proof. Define F', [(d+c)x(d+e) + RR such that
X
F(
= F(X2) X4
(41)
X
for X, E Rdxd X, E _dxe X3 E Rexd X4 E Rxxe Clearly, if F = G H, where G twice-differentiable and H is proper convex, then, F' = G H' where G" is twice- differentiable and H' is proper convex. From the proof of Lemma A.5
know that 33 is
U V
exactly the same as (37 for R =
Furthermore, we know from Lemma A.4 that any
16

---

Published as a conference paper at ICLR 2024
global minimum Z of (40) corresponds to global minimum X of (38) Lastly. we know from
Theorem| any rank-deficient local minimum of (37) corresponds to global minimizer (38)
Putting it all together, we have that rank-deficient local minimizer R
of (39)
rank-deficient local minimizer of 37 which corresponds global optimizer X R.R. of 38
(Theorem A.6 which corresponds global optimizer Z UV of (40) (Lemma A.4
A.4 ADDITIONAL RESULTS ON SIZE OF BM FACTORIZATION
We note the additional result from Zhang 2022). p. In the case that F is twice-differentiable, L-smooth, and μ-Strongly convex, a local minimum (U, V) of (35) corresponds to global minimum of 333 as
long as m (L/H- 1)2m", where m rank (Z)) and Z' the optimal solution (33)
A.5 PROOF OF LEMMA
Proof. We first analyze the solution of the following optimization problem
m
min Vy
L 11 + l1v, I1:
s.t. UV Z.
(42)
can write this as an equivalent problem here (Bach et al. 2008
Pilanci & Ergen
2020
min uyeBo.yy
Vvllnn
s.t. UVT = Z.
(43)
We form the Lagrangian of this as
m l|vlln
f min uj B.eBe.,, R
max
trace UUVV + trace (R (Z
(44)
By Sion's 's minimax theorem, we can switch the maximum over R and minimum over V and minimize over V to obtain
min max trace (R T Z) s.t./lu' R S 1.
(45)
uEBe
As long rank ()), by Slater's condition, 2009) to obtain
switch the minimum and maximum (Shapiro
f max trace (R Z) s.t./lu Rlli S Vu e Bo.
(46)
R
By the definition of dual norm,
also write this
max trace (R Z) s.t.u Rv < 1 Vu Be, Vv E BR = IZIlo.
(47)
R
Thus, with this result, we have
m
p min S(U
Σ
1 + liv
(48)
fe VeRd.xm.
2
equivalently as
min F(MUVT)
Σ L 112 + 1iv,
(49)
Ue Km Ve Xi
2
or also as
min Z:rankZZ)
F(MZ)
+ min
|
+ +11v,
(50)
de
2
4em
UVT-2 =
17

---

Published as a conference paper at ICLR 2024
where we now apply our previous result to obtain
min Crank(2))M
F(MZ) + b
(51)
which if rank(Z*) 2 m is equivalent to
p¹
min F(MZ) + Z
(52)
2
A.6 PROOF OF THEOREM 3.3
Proof. We simply note that 14] is the Burer-Montciro factorization of (16). Thus, from Lemma A.5 as long as m 2 de + dy. all local minima of (14) are global minima. Furthermore, note that (66 is composed of two components, proper convex function. Thus, by Lemma A.7 all rank-deficient local minima are global minima.
of which is wwice-differentiabl function, and the other is
A.7 PROOF OF THEOREM 3.4
Proof. Stationary points of satisfy
O Vuf(ú, v) +
0 V + +V
where we define
| Iú,lee
and the same for [|V] nt By the definition of the subgradient, this stationarity condition can be written as
U s.t. trace - TT) = |U]e. | s 0 - Vof(ū, V) + O/
3V' s.t. trace( V V')= = R. | - 1, 0 Vv + + V V
By the chain rule, we have that
0 V2F(M²)V + +
0 VzF(M2) Ū ||
We now right-multiply the top equation by ÚT and the bottom equation by VT to obtain
0 VZF(M2)VÚ' T + U VzF(M2) UV V
(53)
(54)
Taking the trace, we have
trace (VZF(M2) T2) 2)
|| lotrace (U'TT)= = r[ (V'VT).
(55)
Noting the definitions of U and V we have
trace (VZF(M2) T 2) | |[V
(56)
Furthermore, since clearly F(MŽ) = f(Ù, V). we have that
trace (V2F(M2) 2) = 1111 |i
1 (1111 + + -1/2/
(57)
Now, we examine the optimality conditions for (6)). In particular, we have
n² min F MZ)) + Z
(58)
Z
18

---

Published as a conference paper at ICLR 2024
which by definition of the dual norm, is equivalent to
min max F(MZ) Strace Z'). Z'EBp.
(59)
Now suppose we have an approximate saddle point (Z, Z') with objective value ş such that V
32' = 0, trace(Z _ 2') = [|Z] p. and ž' E (1 + B) for some 2 Let 2 T+c Z'. By strong
duality and non-negativity F, we have
p² max min F(MZ) + Straco Z z')
(60)
Z'E Z
2 min F(MZ) + Btrace (Z 2)
(61)
2
3
= min F(MZ) +
1+c trace Z 2')
(62)
2
2
1
min (F(MZ) + Btraco ( 21))
(63)
1 + E
z
1
(64)
1+
Rearranging, we have that
δ p" p"
Sc.
(65)
For the assumptions of this inequality to hold, for any candidate solution Ż. one must sat- isfy VzF(M2) + 3ž' = 0 and trace Z') |ا|| Solving the former equality for
V2)( we have by 57) that the latter equality is satisfied at any stationary point of Z' = - - the BM factorization. Lastly, for 65 to hold for particular c, one must have
Ż' V2F(MZ) E (1 + c)D.r.
(66)
B
60 [V 11
1) i.e.
p
|VzF(M2)) 11
-
(67)
1"
3
A.8
PROOF OF THEOREM 3.5
Proof. From the stationary point condition, we have
0 V2F(M2) BÙ
(68)
0 VZF(MZ) Ú + +.
(69)
From 69
we can obtain
0 Ú'V2((M2)= -V..
(70)
Substituting this into 68 we have
0 -V V Ú -
(71)
(72)
Thus, let = rank(V) rank(ü)= min(do; drJ. We can write the compact SVD of the stationary point as
Ú = LUAR
V. LvAR
19

---

Published as a conference paper at ICLR 2024
where Lu Rdexr, Lv Rdrxr Assume without loss of generality that de dr, so d,
min(de: d We have
0 V2F(MZ)V + +
(73)
0 VZ + ALAAR
(74)
LL = VzF(MM)
(75)
I dr, Ly square and therefore unitary, so
|VZF(M2)|2 = Vz LvM
(76)
- SLul2
(77)
3.
(78)
Thus, we satisfy 20 with equality.
In general, note that when r < (di, d, }. we have
V F(M2
| Zz (22)
2 FVF(MM)Lvv
(79)
= || BLul|
(80)
=
(81)
so 20 is lower bound, which depends on how Vz (2) behaves when operating on vectors in null
A.9 PROOF OF THEOREM 3.6
Proof. We begin with the non-convex objective
22
m
8 m
PONN = min
W 2/
(X,,,)))
Pa
+
w Wiy
12 || 1B.
(82)
WayER Waj e
j=1
can re-write this as (Bach et al.
2008
Pilanci & Ergen 2020)
PONN
min Wiy WiyEB, Way WReek
Σ F
Σ W P (X,1s)))
+ B 2
W W2llFF
(83)
We can also re-write this as
n Fr.)
m Σ
m W W s.t. Σ
PCNN
min
PX =
(84)
EiyEBa
Way ERexa 1=1
1=1
Forming the Lagrangian, we have
m Σ F(r)) + Σ
n
m
PONN= min WiyEB, ERoe r;
max
2
Σv v
Σ
W2 Pa(,Wa) r
V.
j=l
(85)
By Sion's minimax theorem, swap the minimum minimize over W2j, r; to obtain
W 2jj r, and maximum v,, and
PCNN min max- ueB, V;
Σ FF)
s.t.
P (x 5
(86)
1
20

---

Published as a conference paper at ICLR 2024
where F* is the Fenchel conjugate of F. Now, long B and m m where m < nac,
can switch the order of max and min by Slater's condition (Shapiro 2009) Sahiner et al. 2021b to obtain
PENN max
(v:):)
1=1
s.l. max ueB2
Σ
Pa (x 1F S
(87)
Enumerating over the hyperplane arrangements [D&] we can further write this as
max- PNNN
Fv,
V;
s.t.
max PE[
Σ P.D (1 P 1l S B.
(88)
u
1=1
(2D (6) -
Now, noting that vec(ABC) = (C @
\eee B
Magnus & Neudecker 2019), this is equivalent
PCNN
max Σ (v.v.)
i=1
s.t.
max PE
Σ 1=1
v. P P.D X,uull S
(89)
uEB
(2D
-1)X:220
This may also be written further as
PNNN
Σ F*(vi)
v.
s.t.
max kE[P
Σ
v. & P.D 7 X, S 3,
(90)
uEB,
i=1
(2D u20
REB,
and thereby as
PONN- = max Σ F*(vi)
1=1
s.t.
max
trace Σ
(v. &
P.D. X.)ugT SB.
(91)
uEBa
1=1
(2D
(1)
u1)
gEB,
Now, we let Z = ug
to obtain
PCNN
Σ (viv))
v.
max [c[P]
trace
Σ (v. (T SB.
(92)
Z=ng nEB2
(2D (6)
-1) X220
ge
21

---

Published as a conference paper at ICLR 2024
We let = conv {ug"
I)X,u 2 0, u B2, ge = and note that since our objective is linear take the convex hull of the constraints without changing the objective, obtain
(2Dg
PONN
max
F (v:)
s.1. max trace
Σ 1=1
(v. 8 P.D x,)Z
s 3.
(93)
kE[P] ZEC.
Note that the constraint Z E Ck equivalent to stating that || S for the constrained nuclear norm definition with K2 (2D; I)X. Then, we have
PCNN = max Σ (v.).)
2m1
s.t. max
trace Σ v, 2 P
P.D.
10 x,)
BB.
(94)
KE[P]
|Z||-.Kg $1
Now, form the Lagrangian, given by
PCNN
max
min
min
F (v.)
A
B
vec (z:)
Tvec vec (v. @ (P.D x,)
V. |Zel..K, SI 20
(95)
By Sion's minimax theorem, we are permitted change the order of the maxima and minima, to obtain
PNNN min 2x20
min
Σ Fd
AR
B Σ
voc(Zz) vec (v. vec
(P.D X,
Zu 1. $1
(96)
Now, defining K as the (a, 1) commutation matrix we have the following identity from (Magnus & Neudecker, 2019):
vec v. (P.D. (.DD) = (Ie @ (Ka,, @ [ vec(v.).
Using this identity and maximizing over v;. obtain
PCNN min
min IZ.l-.K. SI Au
F
Σ L
De (K1,a (n))) vec(Zs)
+1
Σ AR-
k=1
(97)
Rescaling such that Zi
A Zs. we obtain
P
PONN = min ZuERAxec
Σ
Σ
I &
(vec(X. (K1,a @ (b)))) veo(Zx)
+B
Z
l.
k=1
(98)
Simplifying further, we can write this as
(race((..D) x,Z
PONN
min ERERex
F
Z 1
(99)
trace P.D (,z(
A=
where ((d) gaxa
A.10
PROOF LEMMA 3.7
Proof.
We start with the convex formulation
ttrcce (P.D (0) x, (1)
P
PRENN
min
F
Σ
+ 3
Σ
(100)
Zy ERA
tracee P..D xiZ
22

---

Published as a conference paper at ICLR 2024
In order to compute the Burer- Moner-Monteiro factorization, factor Z₁ V where U₃ R U) 2 0. Then, with v/(²) Raxm Then for each k,
Va E R and (2D
Trace P. ( x Z (1)
trace P. P.D. 10 X
ttrace P.DDE x,Zz (1)
Lrace(Pa.D E X, U (
trace(V P P.D. U²U
trace( (0)
P.D U,)
-=1 jk
P X
m
P PaD
k X
j=
jk
j=1
where
(1) T
jk
cx
(101)
jb
(
jk
The equivalent Burer-Monteiro formulation thus given by
P
PRCNN=
min (T ((V)ER<xa)m. )- (2D -10
Σ
F
V jk P
D.
(1
1=1
k=1j=1
(102)
A.11 PROOF OF COROLLARY 3.8
Proof. We simply apply the result of Theorem 3.4 noting that stationary points correspond global
minima if the norm of the gradient is less than B. Thus, this condition is equivalent to
tracee P..D. X (1)
Z.
F
u1I2 S B. Vk [P], Vu B2
(2D
(8) k
IX 0.
trace P.D X
(103)
-
A.12 PROOF OF LEMMA 3.9
Proof. Courtesy of Theorem 3.1 of (Sahiner al. 2022 we first present the equivalent convex mod- els for linear activation self-ate ooff-attention. Specifically, as long as m 2 m", where m < min(d? de}. the objective in equation equivalent
d ΣΣ k=1t=1
PLSA
min ZER dde
F
G.lk, Q11 + Z|
(104)
where G. := X. X, G.Ik, E R, and (Z().) E Rdxc] are block matrices which form Z.
23

---

Published as a conference paper at ICLR 2024
To prove Lemma 3.9
Now, seek find the Burer-Monteiro factorization. We let Z UV where U R xm and V E Rdexm Let vec vec-((u) E Rdxd be the result of taking chunks of d-length vectors from uj for j E [m] and stacking them in columns. Similarly, let (vj) E Rexd be the result of taking chunks of c-length vectors from Vj and stacking them in columns. Furthermore, we will let vec "(uz): be the kth column of vec (uj). Then, recognize that
begin from the convex formulation 104
(k, =
Σ vec j
Thus,
d d
d
d
G [k x (k.)
G Xx,vv vec " ((v)
k= t=1
d
X
vec
Gr (vv))
1=1
X
vec -1
X, vec vve (uy)X, X,vec (v;))
j=1
Now, overloading notation, let vec Burer-Monteiro factorization of (104) is given by
(uj) U, and vee (vy) -y.- We have clearly that the
3 + Σ XXX,XXXVVV
PLSA = min
F
U + VV
(105)
U ViERexc
2
1=1
A.13 PROOF OF COROLLARY 3.10
Proof. We simply need to apply the result of Theorem 3.3 to this setting. In this case, the non-convex linear self-attention network equivalent to the Burer-Monteiro factorization the convex form.
To obtain this Burer-Monteiro factorization, we factorize convex weights Z E of xdc so letting observe that long the number of heads exceeds m from Lemma A.5
mn < d2 + de,
Ü
all local optima global. Further, form R
and as long as this is rank-deficient
V
local minimum, it also corresponds to global minimum when F is twice-differentiable by Lemma
B
ADDITIONAL THEORETICAL RESULTS
B.
MLPS
The following theorem demonstrates that can extend the results of (Sahiner et al. simply weight-decay regularization, and to arbitrary regularization.
202 beyond
Theorem B.1. The non-convex ReLU training objective
3
p" = min F
Xw1 W25)
+
w + 11
(106)
Waj
24

---

Published as a conference paper at ICLR 2024
is equivalent the convex training objective
P +.3 Σ ==1
p min Zj 1=1
D2
|2,||o,
(107)
as long as and m 2 m where m I nc, where
Z D == max trace(R R Z) s.t. u Rv S 1 E Bc: (2D; I Xu 2 0, Vv E Bn-
(108)
Proof. We start by re-stating the convex objective Bach et al. 2008
Pilanci & Ergen 2020
min F
+ B
(109)
WiyEBe EBC
Was
Then, re-write this in constrained form
min EBC
F(R) +
Σ
Was a s.t.
Σ
(w11))
W25 R.
(110)
==1
==1
and then the Lagrangian
Wwi - R )) (111)
m
p" min
max F(R) + 3 Σ
way |R trace
V
Σ
WieEBe
Waj
R
By Sion's minimax theorem,
swap the order of the maximization over V and minimization over W2j and R. Then, minimizing over these two, we have
p min max (V) s.t. V (Xu)+ S. uEBe
(112)
By Slater's condition, which holds when /3 0 and m S m² where m S ne (Shapiro Sahiner et al. 2021c), we can switch the order of minimum and maximum to obtain
2009
max F* (V) s.t. max uEBe
(Xu)++ll B.
(113)
Introducing hyperplane arrangements, we have
p" = max - - (V) s.t.
max PI
IV D,Xu i S B.
(114)
V
uEBe (2D, - Dix220
By the concept of dual norm, this is equivalent to
p" max - (V) s.t.
max
trace DD SB.
(115)
ZE P
(2D, Iuu20
gEB,
Define
/ ; max f s.t. Z E (ug T u E Be, (2D, I)Xu 2 0, g B,).
(116)
Then, we can write our problem as
p' max F*(V (V) s.t. max trace D, < 3.
(117)
JEP
|2
Now, observe that 1 = Z D,. particular, let us examine 108). which we can re-write as
||||,, = max trace R T () s.t. max trace(Z (Z R) <
(118)
R
11
max trace(R T Z) s.t. |R <
(119)
= =|[[|
(120)
25

---

Published as a conference paper at ICLR 2024
where the simplifications made noting the definition of the dual norm. Now, write objective as
p max - F (V) s.t.
max PE | $1
trace
(V DzXZ) S B.
(121)
Forming the Largrangian, we have
p" max
min 1Z S1
(V) + Σ N3 (B trace (V (VT,xZZ)))
(122)
A 20
By Sion's minimax theorem, can switch max and min and solve over V obtain
min z ll, 1
F
D,D XZ,
+
Σ A
(123)
-=1
R 20
Lastly, combine Ž, into variable to obtain
P + +
p- min F Σ XZ,)
Σ
| ||D,
(124)
Z
j=l
j=1
as desired.
B.2 CNNS
Lemma B.2. The Burer-Monteiro factorization of the convex CNN problem with linear and gated ReLU activation are given as follows.
3
PLCNN
min (uyeR")
F
VjPoX.us
(1
11.)
(125)
2
1=1
(VGERexOim
13
8
PGONN
min H(ueERR* mm 1P
F
M
Xjujk
+
|| Vy
2
k=1/=1
k=1
(V
(126)
Proof. The proofs follow almost identically from the proof of Lemma 3.7 In the linear case, the convex objective is given by
tracc P. (.X,((11)
PiCNN-
min x
F
+ Zll..
(127)
ttrace(P Pance(P,,..0) ((,X)
26

---

Published as a conference paper at ICLR 2024
In order to compute the Burtr-Monteiro factorization, we factor Z UVT, where U Rhxm
V E Racxm Then, with v(c) E -ax= Then,
Trace P)
Irace(Pa (ox,Uva))
trace (oxoceP
(tace((..XUUv(i))
trace (((1) P P
trace(V(c) P
Σj (1) P.X.u
(6) m 1== j
P.X.u
PaX.u,
where
(1)
y
Rxxa
(128)
The equivalent Burer-Monteiro formulation thus is given by
PLCNN
min [uyeR" Dy 1=1
Σ Σ V VyPaxxu
X,u,
+ Σ
(1|5 + I[V
(129)
j=1
(VyeRoxa3m mm
In the gated ReLU case, the convex program given by
trace P.D X
11)
n F
PGONN
min Z.ERhxae
+ 3 Σ
Izll..
(130)
[trace X 10)
k=1
In order to compute the Burer-Monteiro factorization, we factor Z₁ U V where Ug E Raxmm
VA E Racxm Then, with v(d')
Raxm Then for each k;
k
trace P D ) x 2Z (1)
trace P P.D. (6) X
trace P.D (
X.U.V (
trace (P.D.
trace( (1) ("()) P.D. X₂U))
Irace(
P.D X,U))
k
(1) T
P X,uj
jk
P b X,usk
1= ik
X,ukk
jk
1=1
27

---

Published as a conference paper at ICLR 2024
where
(1)
fkb
R
(131)
jk
jk
The equivalent Burer-Monteiro formulation thus is given by
P
P
PGONN =
mim [[ug L- 1P ([V,EEReX)ja Jan
ΣF F
V
P
D
(1 + (
i=1
k=1j=1
k=
(132)
B.3 SELF-ATTENTION
Lemma B.3. The Burer-Monteiro factorization of the convex self aefffattention problem with gated ReLU and ReLU activations are given as follows.
PGSA
min Vi
F
diag -a--''' "x. X.V
(133)
9
U jkl F + V
2
k=
j=1
m
(daa limp'''D o x.V
PRSA
min
F
(134)
Ugh
Vak
8
U 1 1 s.t. (2diag Baaiigg'(DD) (DD 11 (X,U,xX))) 2 0,
2
k=1j=1
where diag -1 (((()) E xx takes elements along the diagonal of x and places them in matrix form.
Proof. Courtesy of (Sahiner al. 2022), we first present the equivalent convex models for gated ReLU and ReLU activation self attention. First, define
X1 X1
X =
X, X
(D,} (diag (1 (Xu 2 01))
then, have
P
P |Z Il+,
PaSA
min
Σ c
(k.0 x,00 (t:0) Y;
+
(135)
2,
t=1
j=l km 8=1
J=l
P d
P + Σ
PASA =
min
Σ c
G i.
(k.0 (k.2)
(k.0
x,Z Yi
|Z l1K,
(136)
Z,
CERPXX xd
1=1
1= 1=1
1=1
where
:=y (T) T (X
for (k.6) bg
E *x* and 2h e Rdxe
28

---

Published as a conference paper at ICLR 2024
We then proceed to take the Burer-Monteiro factorization of these models. Here, will show the Burer-Monteiro factorization of the ReLU model, noting that the proof is the same for the Gated ReLU model sans the constraints.
We let Z, U. where U, Rm and Vj E Rddxm where K,U, 2 0. Let vec (uzx) E Raxd be the result of taking chunks of d-length vectors from 42 for j E [m]] and stacking them in columns. Similarly, let vec ((zxx) E gexd be the result taking chunks of c-length vectors from jaz and stacking them columns. Furthermore, we will let vee (Ujx), & be the kth column of vec (ujx). Then, recognize that
Zy (k.1)
Σ vec v vexx))ve--
Thus, for each j,
G x
(k.6) bg
X,vec X, (uux) (Vix)(
(4,0)
(x, & I.) ( (X, & II)
X,vec x (aae)vee--(vs
[L ( (1 "(ux)avvc-
X,y,, k|x, X,vec (
10 2 (X.U,xX,)) V,V 52 diag (DD))
where the constraint that (2D IXXU, 2 0 can also be re-written as (2diag Dodiag-(DD) (D! 11 1 (X.U;+XT) 2 0 for all E [m.. Thus, we have proven the statement.
C EXPERIMENTAL DETAILS
C.1 THE RELATIVE OPTIMALITY BOUND
class
02
0.0
0.0
-0.4
-0.2
-0.
-04
-0.2
2.0
04
-24
+0.2
2.0
02
24
(a) 45
(b)
150
Figure 3: Example of three-class spiral dataset, with different numbers of samples n.
See Figure for dataset evaluated. In all cases, solve the BM factorization with GD using Pytorch (Paszke et all 2019) on a CPU with a momentum parameter of 0.9, a learning rate of 1.0 which decays by a factor of 0.. whenever the training loss plateaus, and train for 20000 epochs such that
29

---

Published as a conference paper at ICLR 2024
20
3074
32.1
1p-1
Boars: mi
22"1
2ound. mM
$ouns. m1
Esurs,
22-11
Accad 1p11
to
Adkad mm3
50-11
Acoad me]
2219
Actall
Amall
bota
Deta
beta
(b) 45
75
Figure 4: Supplement to Figure 2
Relative optimality gap of the non-convex BM factorization of gated-ReLU two-layer MLP for three-class spiral data classification (d = 2, c 3). For E [3, 45, 75], we demonstrate how 3 and m affect relative optimality gap, both in in terms of the proposed bound and the actual gap, where the global minimum determined by convex optimization.
GD always converges. For convex optimization, to determine the global optimum of each problem. use the MOSEK interior point solver (Andersen & Andersen 2000) with CVXPY (Diamond & Boyd 2016). The parameters we evaluate are n E [3, 45, 75, 150,, B E [10-1, 10-3 10-2 10-1 and m E [1, 2, 5] (see Figure 4 for = 3, 45, 75). We use randomly subsampled set of f = 100 hyperplane arrangements. We perform this experiment over three random seeds, which are used to generate the hyperplane arrangements well as the random Gaussian initializations of the weights.
C.2 BM ENABLES LAYERWISE TRAINING OF CONVEX CNNs
CIFAR-10 Train Accuracy
90
m-2
m=1
Fashion -NNIST Train Accuracy
100.0
m=
97.5
m
80
95.0
75
92.5
65
87.5
60
85.0
1
2
Stage
Stage
(a) CIFAR-10
(b) Fashion-MNii -NNIST
Figure 5: Supplement to Figure gated-ReLU CNNs.
Training accuracies of layerwise trained BM factorized convex
Our layerwise training procedure was trained on a single NVIDIA 1080 Ti GPU using the Pytorch deep learning library (Paszke et al. 2019). In particular, follow the implementation of (Belilovsky et al., 2019), who proposed greedily, sequentially training two-layer CNNs. At each stage, a two-layer CNN (convolutional layer average pooling + fully connected layer) is trained, and then the weights are frozen, the fully connected layer and average pooling are discarded, and the trained convolutional layer used
Seature-generator for the following stage. At certain stages, before the CNN is applied, an invertible downsampling operation (Dinh et al. 2016) is used to reduce the spatial dimensions of the image. In (Belilovsky al. 2019) ReLU activations are used well as an average pooling operation to spatial dimensions of 2 2 (1.c. 4) is used, followed by a flattening operation and a fully connected layer. They also use a softmax cross-entropy loss, a batch size of 128, weight decay parameter of = 5e 4, along with stochastic gradient descent (SGD) with momentum fixed to 0.9, 50 epochs per stage, and learning rate decay by a factor of 0.2 every 15 epochs.
In our experiments, keep all network and optimization parameters the same, aside from replacing the non-convex CNN at each stage with our convex CNN objective (23). We then apply the Burer-Monteiro factorization with 21, to this architecture make it tractable for layerwise
30

---

Published as a conference paper at ICLR 2024
learning as described in the main paper. At each stage, randomly subsample p 256 hyperplane arrangements. We further use gated ReLU rather than ReLU activations for simplicity, which can work as well ReLU in practice (Fiat et al.) 2019). These techniques have been used effectively for convex learning to exceed the performance of two-layer non-convex neural networks (Pilanci & Ergen 2020
Ergen & Pilanci
2020 Ergenn et al. 202T).
For the CIFAR-I0 experiment, we use 5 stages (following (Belilovsky et al. 2019)) ), whereas for the Fashion-MNIST experiment, we use 3 stages, since the training accuracy saturates after 3 stages. We choose learning rates per stage from {10-1, 10-2 10 100*} per stage based on training accuracy for CIFAR-10. The chosen learning rates were [10-|.10-3 , 100 CIFAR-10. For Fashion-MNISS we empirically observed the training loss was better optimized with slightly higher learning rates, so we used [2 x 10-1, 5 x 2.- x 10-3). In Figure we also provide the training accuracies at each stage for the results provided in Figure T Ultimately, our CIFAR-10 network with 5 stages took 9163 seconds to train, and the Mushion-MNISS network with 3 stages took 4931 seconds to train. In (Kiliçarslan & Celik 2021), it is shown that an end-to-end 6-layer CNN with ReLU activations takes 640 seconds to train CIFAR-10 and 285 seconds to train on FasIOn-MNIST. We note that the purpose of this experiment is not to advocate for the use of layerwise BM networks over end-to-end trained networks, but simply to demonstrate the utility of the BM network in enabling convex neural networks to scale to the performance of end-to-ane deep networks by using layerwise learning.
1003 ,, for
10
C.3 ADDITIONAL EXPERIMENTAL DETAILS
250
BM, m=1
200
BM, m=4
GPU Memory (GIB)
Time per Epoch (s)
150
BM, m=1
BM, m-2
BM, mm4
100
Convex
50
2 Stage
Stage
(a) Time per epoch.
(b) Peak GPU memory.
Figure 6: Supplement to Figure Time per epoch and peak GPU memory consumption of different formulations of gated ReLU CNNs trained in a layerwise fashion CIFAR-10. For "Convex", training causes memory consumption to exceed NVIDIA 1080 Ti GPU's 12 GB memory capacity after the first stage, preventing ability to train convex CNN 23
in layerwise fashion.
In this subsection, we show additional details on the experiments in Figure I to show the scaling issues in using the convex formulation of the gated ReLU CNN described in (23) for layerwise training. We further compare the per-iteration complexity of the convex formulation to that of the BM factorization (26)). In Figure 6 we measure the training time per epoch and GPU memory consumption across various stages of the training CIFAR-10, using the same experimental setup as described in Appendix training the BM formulation to training the convex formulation 23 in this setting. Due to memory constraints, we are not able to train more than single stage of the convex formulation, illustrating that convex neural networks encounter major scaling issues that prevents their application to layerwise learning. We observe that the BM factorization leads to significantly reduced memory consumption and time per iteration compared to the original convex formulation, demonstrating the utility of the B.M factorization for scaling convex neural networks to the layerwise training setting.
faatorizd CNN in a layerwise fashion on
C.2 We also compare
31


Reviews:
[{"review_id": "I7jBZaiVM0", "reviewer": "Reviewer_Mur9", "summary": "The paper adds to the growing literature of viewing 2-layer nonlinear neural networks (MLP, CNN, self-attention) as convex programs (sometimes with large parameters) to study stationary points of these programs via tools from convex and SDP optimization. In particular, using Burer-Monteiro (BM) factorization, the authors give a different parameterizations of these convex programs and demonstrate empirically that the new parameterization allows for faster computability. They also provide theoretical guarantees of local minima of BM programs being global minima of the original program and bound relative optimality gap of BM programs.", "questions": "1. How feasible is it to generalize this framework to more than 1 hidden layer?\n2. 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. \n3. Is there a feasible framework to analyze saddle points of the BM formulation and not just local minima?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I believe the paper is highly interesting, novelly applying existing techniques to a new problem and give ample theoretical and empirical justification for the adoption. \n1. BM theory is shown to be flexible for many different architectures, allowing the authors to show no spurious local minima for linear self attention layer. Although I am not an expert in the theory of transformers, I believe this claim is significant if correct.\n2. BM factorization provides computable objectives for convex 2-layer ReLU networks and is shown empirically to scale under widely used, simple datasets.\n3. If a local minima is obtained, low-rank-ness is sufficient for global optimality.", "weaknesses": "1. The proposed BM factorization for nonlinear MLP is nonconvex and it is not certain how to optimize them to obtain local minima (although the authors do provide methods to check optimality and bound the optimality gap to the original formulation)\n2. It appears that BM factorization corresponds to the \"conventional wisdom\" in practical machine learning of \"adding another layer\" to the model and enjoy nicer optimization landscape. The resulting program, then, is nonconvex (as also pointed out by the authors). In light of this, although the authors provide a small paragraph in page 5 comparing BM MLP to original nonconvex problems, it seems that the formulation in the paper is much closer to this original nonconvex problem (and thus needs further comparison) than to the convex formulation in which the framework is developed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bV1b1RteQW", "reviewer": "Reviewer_8nq3", "summary": "This paper proposes to solve convex formulations of nonlinear two-layer neural networks with Burer-Monteiro factorization, which is known to be computationally tractable for solving convex programs with nuclear norm constraints. It provides theoretical optimality guarantees for two-layer MLPs, CNNs, and self-attention networks. Experiments on FashionMNIST and CIFAR-10 are provided to validate the feasibility and effectiveness of the proposed method.", "questions": "N/A", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The topic studied, i.e., scaling recently proposed convex learning methods is interesting and technically challenging.\n\n2. The proposed method is applicable to different convex neural network structures.\n\n3. The proposed method seems theoretically sound.", "weaknesses": "1. 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¨om 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.\n\n[1] Yuchen Zhang, Jason D. Lee, Michael I. Jordan. ℓ1-regularized Neural Networks are Improperly Learnable in Polynomial Time. ICML 2016.\n[2] Yuchen Zhang, Percy Liang, Martin J. Wainwright. Convexified Convolutional Neural Networks. ICML 2017.\n\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4y2pjhhTMe", "reviewer": "Reviewer_4XdN", "summary": "This paper proposed a Burer-Monteiro (BM) formulation for convex 2-layer neural network training problems. In particular, the BM formulation for MLP (with linear, ReLU and gate ReLU activations), CNN (with ReLU activations) and self-attention networks (with linear activations) are included.\n\nThe convex neural network training problem for linear and gate ReLU activated MLP can be directly solved using polynomial time algorithms such as interior-point methods. However, the problem becomes NP-hard for ReLU activated MLP, CNN and self-attention networks due to the  use of quasi-nuclear norm regularizer, as it is computationally intractable in practice. \n\nThe proposed nonconvex BM formulation turns the quasi-nuclear norm regularizer in convex training problems into a constrained Forbinuous norm, which is computationally tractable. Despite nonconvexity, the authors developed an optimality bound for stationary points in the BM formulation, which upper bounds the optimality gap between the original convex problem and its BM formulation, and, in some cases, is able to provide certification on the global optimality of the given stationary point.", "questions": "1. 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.\n2. 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$)?\n3. Could the authors provide experimental results on demonstrating the scalability of BM formulation? For example, the time and memory complexity (per iteration) v.s. the number of neurons.\n4. Could the authors provide more insight on why it is important to find the global optimum in neural network training? Though the global optimum gives the best training error, it does not guarantee to give the best generalization error. It is entirely possible that the local minimum would have better generalization error than the global minimum. In addition, finding the local minimum is a lot cheaper than solving the BM formulation proposed in this paper.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. The proposed BM formulation provides a computationally tractable way to handle convex neural network training problems that involve quasi-nuclear norm regularizers.\n2. This paper is the first to combine BM with convex neural network training, which is an interesting direction.", "weaknesses": "1. 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.\n2. 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.\n3. Continuing from 1 and 2. The paper 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.\n4. The novelty of this paper is really limited. In particular, this paper is just applying Burer-Monteior to the existing 2-layer convex neural networks, which would be too incremental to be published on ICLR.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UGrw189i8h", "reviewer": "Reviewer_xgT4", "summary": "The paper considers two-layer neural networks formulated as convex programs, involving minimizing a loss term regularized by some (quasi-)nuclear norm term.\n\nThe paper applies the Burer-Monteiro factorization to the neural convex networks, and derives several theorems characterizing local minimizes (Theorem 3.3), duality gap of stationary points (Theorem 3.4), global optimality of stationary points (Theorem 3.5), and similar results for CNNs (Section 3.2) and Multi-head Self-Attention (Section 3.3).", "questions": "See \"Weakness\".", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Omitted.", "weaknesses": "- **First of all, 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, I found no experimental evidence showing that applying BM factorization to convex neural networks leads to more scalable implementations. In particular, I think the title would be only meaningful if the paper compares the baseline that optimizes the two-layer networks by gradient descent and presents improvements over efficiency. This mismatch also comes with several claims that seem inaccurate to me or lack sufficient justifications. For example, \n   - *The BM factorization is essential for convex neural networks to match deep ReLU networks.* The paper is limited to 2 layer networks, I am not sure whether it is suitable to generalize to deep networks. Even in the 2-layer case, the BM factorization is not shown to be essential or promising compared to the above baseline.\n   - *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.* However, What remains to be justified is this: Why is layerwise training necessary, and why cannot one just train all layers together (as the above baseline)?\n\n\n- **Second, the theorems seem to be straightforward extensions of prior works (e.g., as the paper mentions, Bach et al. (2008) and Haeffele et al. (2014))** Please compare prior works at the technical level. What's the novelty of the paper at the technical level? Please justify. Maybe it makes sense to make a table for comparison.\n\nOverall, I found the paper is focused too much on writing and rewriting the problems in different ways, without making a special effort to actually solve the problems. Having presented many reformulations, the paper leaves the readers wondering whether BM factorization works or not. In particular, the theorems are weaker than the convex counterparts (the latter guarantees global optimality). If the paper is unable to show BM factorization leads to improvements in scalability or efficiency, the values of the proposed approach would be greatly limited.\n\n\nMinor:\n- In Eq. (1), maybe it would be great to have parentheses for the summation term.\n- Sentence below Eq. (5): two \"thereby\", thereby not reading well.\n- The paragraph after Eq. (19): \"*While the original convex program is NP-hard*\". It is unclear to me why convex programs are NP-hard.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "M8YdRmAINl", "reviewer": "Reviewer_TUoh", "summary": "This paper revisits recent developments which reformulate the training of a two layers neural network into a convex optimization program but are however often intractable computationally. This paper proposes to apply the Burer-Monteiro factorization to such formulations, in order to make them as tractable as the original convex formulation. In some cases, this retrieves the original non-convex formulation, and in some cases (such as for ReLU MLPs) it gives a novel formulation. The analysis includes linear but also, for the first time in the literature, non-linear ReLU networks, and it tackles MLPs, ConvNets, and self-attention networks. The overall result is that such Burer-Monteiro factorization allows to obtain algorithms for training 2 layers networks which are as efficient as their original non-convex formulations, but which can have guaranteed theoretical properties such as a bound on the relative sub-optimality gap. Finally, these developments are plugged into a (tractable)  layer-wise training procedure of (deeper than 2 layers) CNNs, which allows them to inherit from the theoretical guarantees developed in the paper (and empirically, such CNNs are comparable to state of the art training of CNNs in term of test error).", "questions": "It is still a bit unclear to me how the training of the BM factorization for ReLU networks (19) is done in practice: since this is a constrained optimization problem, I guess its training should depart from a simple vanilla, unconstrained GD ? From the experimental section and also after briefly looking at the corresponding code, unless I am mistaken, it seems that training of pure ReLU networks has not been investigated (I think only Gated ReLU and Linear networks were trained in practice ?). Since analyzing (non-gated) ReLU networks are one of the main result of the paper, I think it could have been interesting  (perhaps in Appendix) to just compare in more details, even just on the toy spiral dataset and for 2 layers ReLU MLPs, the training of the BM factorization vs. the original MLP training, since from eq. (19), the two formulations differ, and since BM has some theoretical guarantees which the  original non-convex formulation does not have.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "### Originality\n\nI believe this paper is original, since up to my knowledge it is the first one which considers the Burer-Monteiro factorization of the convex formulation of non-linear ReLU two layers neural networks: additionally, in such case the BM factorization differs from the original non-convex formulation, which makes it an original object to study in itself.\n\n### Quality\n\nI believe the work is of quality, with theorems and their assumptions clearly stated, and their detailed proofs provided in appendix. Additionally, the literature review seems exhaustive, and some care was given to the experiments, with the code provided in the supplemental.\n\n### Clarity\n\nI think the paper is clear, with the motivations and main results clearly highlighted.\n\n### Significance\n\nI believe this work is very significant, since ReLU networks (MLPs, CNNs, and self-attention), contrary to linear networks, are actually widely used in the machine learning community. Therefore, I believe the BM formulation obtained in the paper, as well as the theoretical guarantees that follow from it will be of interest to the community. Additionally, the empirical result at the end regarding the training of deeper than 2 layers CNNs is encouraging and hints at the applicability of the results in that paper even to deeper than 2 layers neural networks. Additionally I think that the discussions around the convergence of (S)GD, stationary points, global optima, and the new theoretical developments in the paper related to those issues, will be interesting to the community.", "weaknesses": "I just have a question regarding the paper (see below).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)