CRONOS: Enhancing Deep Learning with Scalable GPU Accelerated Convex Neural Networks
Authors: Miria Feng
Zachary Frangella
Mert Pilanci
Venue: NeurIPS 2024 poster (2024)
Keywords: convex neural networks
deep learning
admm
alternating minimization
gpu acceleration
Primary area: optimization

Abstract: We introduce the CRONOS algorithm for convex optimization of two-layer neural networks. 
CRONOS is the first algorithm capable of scaling to high-dimensional datasets such as ImageNet, which are ubiquitous in modern deep learning. 
This significantly improves upon prior work, which has been restricted to downsampled versions of MNIST and CIFAR-10.
Taking CRONOS as a primitive, we then develop a new algorithm called CRONOS-AM, which combines CRONOS with alternating minimization, to obtain an algorithm capable of training multi-layer networks with arbitrary architectures.
Our theoretical analysis proves that CRONOS converges to the global minimum of the convex reformulation under mild assumptions. 
In addition, we validate the efficacy of CRONOS and CRONOS-AM through extensive large-scale numerical experiments with GPU acceleration in JAX.
Our results show that CRONOS-AM can obtain comparable or better validation accuracy than predominant tuned deep learning optimizers on vision and language tasks with benchmark datasets such as ImageNet and IMDb.
To the best of our knowledge, CRONOS is the first algorithm which utilizes the convex reformulation to enhance performance on large-scale learning tasks.

CRONOS: Enhancing Deep Learning with Scalable GPU Accelerated Convex Neural Networks
Miria Feng Electrical Engineering Stanford University piria0@stanford edu
Zachary Frangella
Mert Pilanci Electrical Engineering
Management Science & Engineering
Stanford University
Stanford University pi pallacioonn sanncisttanfode edu
zfran&stanf ord. edu
Abstract
We introduce the CRONOS algorithm for convex optimization of two-layer neural networks. CRONOS is the first algorithm capable of scaling to high-dimensional datasets such as ImageNet, which are ubiquitous in modern deep learning. This significantly improves upon prior work, which has been restricted downsam- pled versions of MNIST and CIFAR-10. Taking CRONOS as a primitive, we algorithm called CRONOS- AM, which combines CRONOS
then develop with alternating minimization, to obtain an algorithm capable of training multi- layer networks with arbitrary architectures. Our theoretical analysis proves that CRONOS converges to the global minimum of the convex reformulation under mild assumptions. In addition, we validate the efficacy of CRONOS and CRONOS-AM through extensive large-scale numerical experiments with GPU acceleration in JAX. Our results show that CRONOS-AM accuracy than predominant tuned deep learning optimizers on vision and language tasks with benchmark datasets such as ImageNet and IMDb. To the best of our knowledge, CRONOS is the first algorithm which utilizes the convex reformulation enhance performance on large-scale learning tasks.
obtain comparable or better validation
1
Introduction
The non-convex landscape of deep neural networks (DNN) poses significant challenges for modern deep learning, especially since non-convexity implies it is NP-hard to train a neural network to optimality Blum and Rivest [1988]. Common stochastic first-order optimizers, such stochastic gradient descent (SGD), offer no guarantees of producing more than an approximate stationary point Arjevani al. [003),. which may be considerably suboptimal Ge al. (2013). The effectiveness of methods such as SGD and Adam are therefore reliant on heavily conditioned training environments. As result, most models require extensive hyperparameter tuning of the optimizer train successfully, This leads to expensive iterations in high compute settings with variable performance depending on optimizer selection and problem domain Yao et al. (00TT) Scaling Rosenffll [20TT] also indicate this regime will yield an increasingly larger number of hyperparameters, with a disproportionate dependence on computational resources and data cost.
Our objective is to achieve more efficient targeted optimization deep learning tasks by leveraging the connection between neural networks and convex optimization. This approach provides clearer insight to the underlying optimization problem deep learning is trying to solve, despite the complexity of its non-convex landscape. The authors of Pilanci and Ergen (2020) have recently proven that the training of shallow neural networks can be equivalently formulated as a convex optimization program. The strategy leverages semi-infinite duality theory develop algorithms which converge the global minimum in polynomial time. However, the resulting convex program is constrained high-dimensional linear model that is intensely difficult to solve scale. Other more recent works of Mishkin et al. [2022] and Bai et al. 2023 have made progress in solving this problem with
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

small downsampled versions of MNIST Noever and Noever [2021] and CIFAR- 10 Krizhevsky and Hinton (2010), yet were unable to scale realistic high-dimensional datasets that are ubiquitous in modern deep learning. This inability to handle large real-world data significantly limits the practical deployment of convex neural networks, despite their strong theoretical guarantees.
In this work, we propose CRONOS: the Convex Reformulated Neural Network Operator Splitting algorithm. CRONOS is fast, efficient and nearly hyperparameter- method for training two- layer convex neural networks. We then augment our algorithm in CRONOS-AM: Cronos with Alternating Minimization, which extends applicability to neural networks of arbitrary architectures beyond just two-layer networks. We implement all experiments and algorithms in JAX Roth et al. (2024) with the RTX-4090 GPU. This allows CRONOS fully leverage GPU acceleration and eliminate memory bottlenecks. thus successfully tackling problems with large data. In order to sustain strong convergence guarantees, parallelization, and robustness to hyperparameter tuning we utilize ADMM BByy et al. [2011] as a core solver in our method. This enables exciting new time-efficient strategies for CRONOS to handle scalability in real world problems. We evaluate the performance of our algorithms in binary and multi-class classification tasks across wide domain of large scale datasets (both image and language), on three model architectures (MLP, CNN, GPT-2). Additionally, our theoretical analysis proves the convergence of CRONOS to the global minimum of the convex reformulation under mild assumptions.
To the best of our knowledge, this is the first time convex reformulated neural networks have been successfully applied to large data such ImageNet Recht al. (2019) and large language modeling tasks with GPT-2 BBdzianowski and Vulié 2019] architecture. Our main contributions can be summarized as follows
We develop a practical algorithm via convex optimization and the Alternating Directions Method of Multipliers (ADMM) Boyd al. [2011] train two-layer ReLU neural networks with global convergence guarantees.
Using this primitive, extend our algorithm effectively train luyti-layer networks of arbitrary architecture by combining CRONOS with alternating minimization.
We demonstrate the efficient practical applications of our method real-world large scale image and language datasets, including ImageNet and IMDb.
Our experiments are implemented in JAX functional programming paradigm for sig- nificant GPU acceleration. Results demonstrate performance speedups and successfully overcome the problem of memory bottlenecks tackle large data.
Our theoretical analysis proves the convergence guarantees to global minimum under mild assumptions of proposed algorithms.
2 Related Work
CRONOS builds upon key ideas from previous literature convex neural networks, ADMM, randomized numerical linear algebra and synergy with JAX implementation. Please see Appendix A for detailed discussions of related work according each of these three areas of interest.
3 Background
This section introduces the convex reformulation of two-layer neural networks. We formalize this in two steps: 3.1) foundational defimition of convex ReLU neural networks and 3.2) its equivalence as linearly constrained Generalized Linear Model (GLM) with group lasso regularization.
3.1 Convex ReLU Neural Networks
Given a dataset X E Raxd a two-layer ReLU multilayer perceptron (ReLU-MLP) with weights W(!) Rmxd (2) E Rm outputs the prediction:
fwa), 20 (X)
Σ (XW1))+
(I)
'Our codebase is available at htttp:///gtaatbcc/////aaaaaaaa//

---

Here (x)+ = m 0} denotes the ReLU activation function.
Given targets y IR", the network in ① is typically trained by minimizing the following non-convex loss function:
m
min (fw, wr (X),9) + W1,202
[ [1, 2 (w2y))
(2)
where : C:RR R" + R is the loss function, and 3 2 is the regularization strength. It is typically challenging to solve ( since the optimizer often needs meticulous tuning of hyperparameters to ensure successful training. Such tuning is expensive, since it requires many iterations of running the optimizer across multiple hyperparameter configurations in a grid search obtain good performance. This dramatically contrasts with the convex optimization framework, where algorithms come with strong convergence guarantees and involve minimal hyperparameters. Fortunately, possible to maintain the expressive capabilities of ReLU neural networks while still enjoying the computational advantages of convex optimization.
Pilanci and Ergen 2020] have shown 2 admits a convex reformulation, which elides the difficulties inherent in solving the deep learning non-convex landscape. This reformulation has the same optimal
value as the original non-convex problem, provided m 2 m"" for some m S n + 1. Therefore, 2 convex program does result in loss of information generality.
reformulating
Pilanci and Ergen 2020 's convex reformulation of 2 also enumerates the actions of all possible ReLU activation patterns on the data matrix X. These activation patterns act separating hyperplanes, which essentially multiply the rows X by 0 or 1, and can be represented by diagonal matrices. For fixed X, the set of all possible ReLU activation patterns may be expressed as
D { D diag (1 2 0)) v E IR4)
The cardinality of Dx grows as Given D; E Dxx the set of vectors for which (Xe)+ = D;Xu, is given by the following convex cone:
[Dx O (r(n/r)"), where r rank Pilanci and Ergen [2020].
K, = fo R (2D, x)X 2 0}.
Learning directly based on the enumeration of Dx impractical due to the exponential size of Dx [Mishkin et al.) 2022]. Instead, we motivated to work with the following convex program, based on sampling activation patterns from Dx:
min e (Exxx(e (wra0r)
Di (0luu-w.,,y)) +
w.).y
||ulle + ||
1=1
s.t. 0i, Wy E K.;
Vi E
Although 33) only works with subsampled version of the convex reformulation of Pilanci and Ergen 22220] it can be shown under reasonable conditions that (3) still has the same optimal solution as Mishkin et al. [2022]. Therefore we can confidently work with tractable convex program 3
3.2 Convex ReLU networks as linearly constrained GLMs with group lasso regularization
Prior derivations Bai et al. 2023 have shown that may be reformulated as linearly constrained composite convex program:
Proposition 3.1. Define the matrices F, D,X and G; (2D; T)X, where i [P]. Then by introducing the constraints u; = Ug 2; = w,, where E [P]. and appropriate slack variables can be reformulated as:
$1...., 8p, tp.t,....,tp
min u.u.a)
((Fu. v) + llllla11 + 1(s 2)
12PP
(4)
s.t.
u
= 0.
s
where
U1 E 2dPP F E Rdd G R2nPx2dP
The reformulation is essentially a very large constrained generalized linear model (GLM) with group lasso penalty Yuan and Lin 2006). The data matrix F associated with 4 has dimensions

---

x 2dP, and the constraint matrix has dimensions 2(n + ddP x 2dP. Although the sizes of F and the constraint matrix seem intractable, they are highly structured. Each F; and G; consists of the data matrix X multiplied by diagonal matrix of 0's and I's. Therefore, F and G do not need to be instantiated, and we can apply matrix-vector products efficiently by exploiting this structure on GPU-accelerated frameworks. Additionally, if X is approximately low-rank (a common phenomenon in machine learning), its singular values must decay fast. F; and G, then also inherit this approximate low-rank structure in (4), which CRONOS exploits and solves via fast matrix-vector products on GPU acceleration.
4 CRONOS: Convex Neural Networks via Operator Splitting
This section introduces the CRONOS algorithm for solving 4
CRONOS is a scalable, convex optimization-ban learning method capable of handling large data and utilizes GPU acceleration for enhanced performance.
4.1
ADMM for robustness and decomposability
Algorithm 1 ADMM for Convex ReLU Networks
Require: penalty parameter P repeat ++1
argmin, 1 y/1 2 Ellu v² + +1 + 2llGu sh + 111
argmin o 0110012,11 1(s 0) 1lk e 1
Primal update
s gh+1
x+1
A + ta (w h+ - ))
Dual À update
zk+1 v
(Gu - s 1
Dual update
until convergence
Equation 4 is in a natural form to apply the Alternating Directions Method of Multipliers (ADMM). the seminal first-order optimization algorithm for solving convex optimization problems Boyd et all [20TT] ADMM is an ideal choice due to its strong convergence guarantees, robustness to hyperparameter variations, and ability to effectively leverage modern computing architectures through parallelism. However, directly applying ADMM to ( requires solving the u-subproblem at each iteration:
2 h+1 = argmin, 11 - y/1²
P ?1
v + + LIGu - s + +
(5)
2
This solution then requires solving the following linear system:
(H + T) b².
11 1 FFFF+
(6)
8 FT F w N GT (s* v²).
Since is 2dP matrix and Gis a 2nP 2dP matrix, the cost of solving (6) via direct method is p4 + p3). This is prohibitively expensive since n and dP' are typically large. A natural method is solve (6) inexactly via the Conjugate Gradient (CG) algorithm, which only requires matrix-vector products (matvecs) with F, FT, G and GT. Under these conditions ADMM will still converge, since the sequences of subproblem errors are summable Eckstein However the number of iterations required by CG to achieve ò- 6-accuracy is (K(H + 1) (1/8)). where =( (H + 7) is the condition number H I. This results slow convergence for machine learning problems, since F and G' are closely related to the data matrix X, which is approximately
summablelEckstein and Bertsekas [192].
low-rank (dell and Townsend (2019). Given that the linear system (6) must be solved at each iteration, CG's slow convergence renders it infeasible for real-world application.

---

4.2
Nyström preconditioning for fast convergence
Algorithm 2 CRONOS
Require:
penalty parameter p. forcing sequence 1k=1 rank parameter
[U, AJ RandNystrómApprox
r,
Compute using Algorithm
repeat
Use Nyström PCG (Algorithm 4h to find un *+1 that solves ( within tolerance 8t
x++1 prox 111-112
(u + )²)
P Primal v update
sk+1 (Gu + +u4
Slack s update
A Al + ta (u 1 1
Dual À update
wk+1 f- ye + 7a (Gu h+1 - 8 sl+l)
Dual v update
until convergence
We exploit the approximate low-rank structure in our problem matrices speed up convergence by applying the NysADMM algorithm from Zhao et al. [2022]. which is an ADMM-based method targeted at solving large machine learning tasks. Notably, the subproblem for e*1 and s has closed-form solution that may be computed in O ((n + P) time. NysADMM employs the Nyström Preconditioned Conjugate Gradient (NysPCG) from Frangella al.
[2023a solve (6). NysPCG is a linear system solver specializing in solving linear systems with large, approximately low-rank matrices.
NysPCG first constructs a low-rank preconditioner P for the matrix H + I in 6 from a randomized Nyström approximation of H. When P constructed with the appropriate rank, it can be shown that Nyström PCG solves the linear system in (6) to 5-accuracy within O (log(1/6)) iterations (Proposition 6.3 The dependence be solved quickly at each iteration. Details of the NysPCG algorithm and construction of the preconditioner are presented in Appendix B.1
the condition number is therefore eliminated, and (5) can
We refer to our algorithm in solving 4 as CRONOS (Convex ReLU Optimized Neural networks via Operator Splitting). The CRONOS algorithm is presented in Algorithm 2
4.3
Scale and speed with JAX and Just-In- Time compilation
Scaling convex neural networks to realistic high-dimensional data is critical for machine learning problems. Therefore we implement our methods in JAX [Bradbury et al. (018)) a high-performance numerical computing library designed to accelerate machine learning research. This framework provides an efficient way perform array operations, automatic differentiation, and optimization of numerical processes. Leveraging Just-In-Time (JIT) compilation capabilities through XLA (Ac- celerated Linear Algebra) allow to execute optimized machine code with extremely accelerated and scalable performance on GPU. Currently, the importance of GPUs in deep learning cannot be overstated. Therefore we note that any practically competitive algorithms need fully utilize parallel processing capabilities. The combination JAX and JIT compilation effectively enables us to scale to high-dimensional datasets such Food (267 267 3), ImageNet (512 512 3), and the IMDb language dataset. Our experiments are summarized Section
CRONOS-AM: Deep Networks via Alternating Minimization
In this section, we introduce the CRONO-AM algorithm for training neural network with arbitrarily many layers. Consider training an L-layer neural network with ReLU activations and the least-squares loss. Training the network involves solving:
I 2
minimize 0
[]](0XX- - 1/12
W 3 + (IIW'-1 12 || 11).
2
where (vee (1), veo(W)),... Wr.) E RP is the vector of the network weights, with W; denoting the weights for the äth layer, while X is the data matrix, and are the labels. We can rewrite
5

---

this objective as:
L-2 (FLL--((01--2) 20 11 - - + W. W.ll+ (1W +l
1|2
minimize
Oii W. 1,10
where 01.L-2 (v)c(W1), ve vec(WW F FI:L-1 consists of the first 2 layers. Let us write ( = -1 X), which may be viewed transformed data matrix. We shall sometimes write X for brevity. Hence, we can write the output of the network as:
F(w; X) (x(0 (X(01.12)WZ
WL.
Thus, the training problem may be rewritten as:
a T-2 | + (1wwz-ll2 ||,rx1/)).
(XwLL-))WW (
minimize 016--2/6611:02
WL -
2
When f1.-22 is fixed, the preceding problem can be viewed as training a two-layer ReLU neural network with transformed data matrix X. Motivated by this, we replace the part of the optimization involving the last two layers with the convex reformulation to obtain:
minimize (u,v,e),
|F (((WLL=1)
ul12
Σ
L i=1
|W + Bll 2.1 + 1ss21100)
2 0).
W
12PP
s.t.
u
= 0.
Equation 7 decouples the convex weights from the non-convex weights. This puts the objective into a natural form to apply alternating minimization, where we alternate between minimizing with respect to W1.L-2 and (u, v, s). To handle the convex minimization, we can apply CRONOS. For the non- convex portion, we propose utilizing DAdapted- Adamm (Defazio and Mishchenko 2023 as it does not require setting the learning rate. call this algorithm CRONOS-AM (CRONOS MOOONOS-AAtemnnn Alternatin Minimization). Pseudocode for CRONOS-AM is given in Algorithm
Theoretical Analysis of CRONOS
In this section, establish the rate of convergence for CRONOS in solving (4) and provide overall computational complexity estimate. Additionally, we show that the linear system for the a -updaae may be solved by PCG rate independent of the condition number.
6.1 F's and Goss inherit approximate low-rank structure of X
We begin by showing the spectrum of x upper bounds the spectrum of F? F; and G. Proposition 6.1 (F; and G; are approximately low-rank if X is). Let E [P), and j E (d]. Then the eigenvalues of FT F; and G? G; satisfy:
max (AAAEG S xx)xx).
Proposition 6. shows that if X is an approximately low-rank matrix, so are all the Ft's 's and G, 's. Most data matrices in machine learning exhibit fast polynomial or exponential spectral decay and so are well-approximmted by low-rank matrices Wainwright (2019). Derezinski et al. [2020). As the matrix H that defines the linear system arising from (5) is built from the F/s and the G.s, it will also be approximately low-rank, which motivates the following assumption.
Assumption 6.2 (Approximate low-rank structure). Let (HH) denote the jth eigenvalue of H. Then the eigenvalues values of H satisfy:
(, - 0 B 1/2.
Assumption 6.2 posits that the eigenvalues of H decay at a polynomial rate, which is reasonable given the preceding discussion. Under Assumption 6.2 we will obtain complexity guarantees that align with CRONOS practical performance.

---

6.2 Fast -update
The next proposition shows that NysPCG solves the linear system in each iteration in (5) fast.
Proposition 6.3 (Fast solution of u-subproblem). Assume Assumption 6.2 holds. Suppose the randomized Nyström preconditioner is constructed with rank r (log())) Then after t
O (l))())) iterations, Nyström PCG outputs a point ++1 satisfying
[[ + T)u h+1 - 1 S
with probability al least 1 c.
For fixed > O, Proposition 6.3 shows that with rank of r = (1), NysPCG solves the u- subproblem within (log( 3z iterations. Thus, PCG solves the linear system quickly. Proposi- tion 6.3 is consistent with practice, NysPCG with rank of 20 and 10-40 PCG iterations work well across all problem instances. Therefore solving ( is not bottleneck for CRONOS.
6.3 Convergence of CRONOS
The following theorem shows CRONOS converges to the global minimum (4)
Theorem 6.4 (Convergence and Computational Complexity of CRONOS). Suppose Assumption 6.2 holds. Fix Ç (0, 1) and denote the minimum of 4 by p* Construct the Nyström preconditioner in Algorithm 2 with rank O (log(t)). ) and at the kth CRONOS iteration run Nyström PCG with tolerance 8 convergence. Then, with probability at least Ç the following statements hold:
1. After K iterations, CRONOS' output satisfies
{(Fa K 3) + 31158
112,1 + 1(8 2 - =0(/))
12aP ūk
Tak
= -11KK)
z
2. The total complexity of CRONOS to produce an c-suboptimal point of (4) is
CeRONOS = ndd2
Theorem 6. shows CRONOS converges ergodically at an O(1/K)-rate. When the minimum of and (2) coincide Theorem 6.4 guarantees CRONOS is within of the global minimum of ( after (1/ee) iterations. By comparison, the worst-case SGD on (2) can only be guaranteed to output an e-approximate stationary point after (L1/e*) iterations Arjevani et al. [2023]. SGD's s output may
also fail to be close to local minimum, and can be highly suboptimal Ge al. 2015]. Therefore,
CRONOS offers much stronger convergence guarantees than SGD.
Theorem 6.4 the first realistic convergence guarantee for ADMM on (4). Previously the convergence analysis of Bai et al. (2023) assumes the ADMM subproblems are solved exactly at each iteration, which is unlikely for large-scale data. Consider a precision > then the cost of the state-of-the-ard ADMM approach from Bai et al. [2023]. is at least O(nd² p4 + d p3) since it solves the linear system derived from Eq. (5) exactly. contrast, the dependence CRONOS upon n, d, and P offers significant improvement, since the cost grows linearly with and d. This enables our scalability to very high- -ige--immnsional datasets both vision and language.
Theorem 6.4 as presented is pessimistic about the overall convergence speed of CRONOS. For sparsity-pred prriity-poooiing..o convex regularizers such as the group lasso penalty, ADMM known to reach the manifold containing the support in finite time, after which it converges linearly Liang et al. [2077). Yuan al. (2020). Therefore, CRONOS convergence is practically much faster than what Theorem 6.4 suggests.
For detailed discussion when this occurs, please Appendix
7

---

DUD 1aagpo
56D
Suapp
Veu
Time
(a) CIFAR-I0
(b) ImageNet
Figure 1: CRONOS-AM vs. competitors on Deep ReLU MLP
7 Experiments
In this section, empirically evaluate the efficacy of CRONOS and CRONOS-AM classification tasks. We first experiment with vision data in multi-layer perceptron (MLP) and convolutional neural network (CNN) architectures" We find that without the necessity of tuning hyperparameterss CRONOS performs as well or better than prevalent standard optimizers. To the best of our knowledge, CRONOS is the first convex reformulated neural network capable of handling large scale data tasks such as ImageNet classification. Secondly, we evaluate the performance of CRONOS on natural language sentiment classification with GPT2 architecture. All experiments were run in JAX v0.4.28 and FLAX v0.8.2. Appendix [ presents detailed discussion of the experimental setup and additional numerical results. Details on each dataset used can be found in Appendix F33
7.1 Training deep Multi-Layer Perceptron
We evaluate the performance of CRONOS- for binary classification with a 4-layer multilayer perception whose architecture is provided in the Appendix [.5 We train the model on three different datasets CIFAR-10, Food, and ImageNer ! TRONOS-AM. is benchmarked against several of the most popular optimizers in deep learning: SGD with Polyak momentum (SGD) (Sebbouh al. 202T Adam [Kingma and Ba 2014). AdamW (Loshchilov and Hutter (017). Shampoo Cuppt et al.
2018 Yogi [Zaheer al. 20188) and D-adapted Adam (DAdam) [Defazio and Mishchenko [2033]
For each competing method we consider 5 learning rates, selected randomly from a logarithmic grid
with range [10-5.5 10 1051.
Fig. plots the median trajectory across different learning rates each competing method along with the 5th and 95th quantiles. Fig. [ shows CRONOS-AM either achieves the best comparable performance relative to its competitors. These plots and the tables show competing methods exhibit an extremely high degree of variance, and poor learning rate selections can yield non-convergent behavior. In contrast, CRONOS-AM does not exhibit these weaknesses and performs comparably to the best-tuned competitor.
Table 1: Results for CIFAR-10 and ImageNet Datasets
CIFAR-10
ImageNat
Optimizer
Peak Validaionn Range
Best Leaminng Rase
Peak Validation Runge
Best Leurning Rate
CRONOS-AM
90.55 5%
88.47%
NA NA
DAdim
15%
NA
87.96%
(9..4.89.68))
10
17
10 6
AdamW
[50.04, 90.25)00 %%
1.56 x 10
[50.87, [70.87.87.46]
4.07
10
Yogi
[60.04, 90.84
[60.87, 88.47]%
6.65
SGI
[50.04, 87.75) %
5.71 x 10-3
[50.87, 87.466 96
4.94 10-3
Shampoo
151.5, 89.15
5.07 x 10-3
69.
1.70 10-4
'The results for the convolutional experiments may he found Appendix F66 Resulls for Food and other downsampled variants of ImageNet may be found in Appendix F.5

---

Table 2: Optimizer Runtimes (s) on CIFAR- 10 and ImageNet
Dateset
CRONOS-AM
D-Adagted Adam
Shangoo
Yogi
CIFAR-10
ImageNet
3.10
1.84
3.14
2.94
2448
5.19
1.98
Volaate
17
Araed
Chovee
Thonrd
Paser//Oater Teeestion
(a) CRONOS vs. tuned AdamW on GPT2-NFT
(b) CRONOS vs. Adam trajectories on GPT2-NFT.
CRUINOS
(c) CRONOS vs. tuned AdamW on GPT2-FT. (d) CRONOS vs. AdamW trajectories on GPT2-FT.
Figure 2: CRONOS AdamW GPT2 configurations for IMDb
Table I presents the range in peak validation across the grid (except for CRONOS-AN and DAdam, which do not require learning rate parameter). CRONOS-AM outperforms DAdam on both tasks and performs comparably to the best-tuned first-order optimizer. On ImageNet, Shampoo does best by fair margin. We attribute this to Shampoo being approximate second-order optimizer. Properly tuned, Shampoo may yield better performance than purely first-order optimizers like CRONOS-AM and Adam for certain tasks.
Table shows the total runtime in seconds for each optimizer on CIFAR-I0 and ImageNet. Despite doing more work than competing optimizers, CRONOS-AM's runtime is comparable with standard optimizers such as Adam, AdamW, and Shampoo.
7.2 Natural Language Classification with CRONOS
Table 3: Results for Different GPT2 Architectures
GPT2-NFT
GPT2-FT
Optimizer
Peak Validation Range
Best Learning Rate
Peak Validation Range
Best Leaming Rate
CRONOS
77.66%
NA 1.56 x 10-3
93.91% [18.30, 93.69)%
NA 1.32 10-4
AdamW
[40.29, 73.69]%
Our experiments explore sentiment classification with the IMDb dataset. For all experiments, use the pretrained GPT2 architecture with 12 transformer blocks and embedding dimension of 768. We use the same GPT2Tokenizer across all language model tasks to ensure consistent evaluation.

---

and all dataset examples are padded to the same length of 60 words. The input to CRONOS is thus shaped into batches of size (1500 x 60 x 768) for each of the labels, and all experiments implement multiple trials across varying data batches. Our baseline benchmark model is the GPT2 pretrained model passed through one dense linear layer for classification accuracy. Numerical results are summarized in Table [3] and Fig. compares CRONOS to tuned Adam on time. CRONOS is seen to reach best validation faster than AdamW. particularly in Fig. 2a Fig. 2b and Fig. 2d plot several AdamW trajectories along with CRONOS against number of epochs for AdamW and ADMM iterations for CRONOS, The more translucent the curve for AdamW indicates larger deviation from median trajectory. Both plots show AdamW is extremely sensitive to the learning rate selection, Appendix F.4 provides further details of the three CRONOS integrated GPT2 experiments.
IMDB-NFT. We extract the GPT2 pretrained checkpoint and then immediately follow up with CRONOS for sentiment classification. Results are shown in Fig. 2a and Fig. 2b) Notably, this setting does not involve any training loops with standard optimizers (such as AdamW), and the foundation GPT2 model does not see any labeled sentiment data from the IMDb dataset. We limit the amount of data seen by CRONOS to only two batches to evaluate the efficacy of our method in the low data regime. Table shows CRONOS significantly outperforms tuned AdamW. reaches higher validation accuracy Taster than AdamW and has the benefit of not requiring hyperparameter grid search. In contrast, Fig. 2 shows AdamW's performance may be quite poor if the learning rate is not judiciously selected.
IMDb-FT. Our fine-tuned experiment setting utilizes the GPT2 pretrained checkpoint followed by one epoch only of training with the AdamW optimizer on default parameters with the standard dense linear layer head followed by CRONOS. Results are shown in Fig. [c] and Fig. 2d] Although the authors of BERT Devlin [2018) recommend 2-4 epochs for training in this setting, we aim to evaluate the performance of CRONOS with limited overhead to extract features for maximum efficiency. The summary of results in Table 3 are particularly promising CRONOS reaches 93.91 % validation accuracy, about 0.3% better than tuned A, which is widely regarded as the most effective method for training language models.
IMDb-DA. We examine the setting of training for one epoch of AdamW on unlabeled IMDb data initialized from the GPT2 checkpoint. The resulting features are then extracted and passed into CRONOS for classification. The motivation for the experiment is to examine the potential of CRONOS on unsupervised domain adaptation settings in future work, which aims to reduce the distribution gap between source and unlabeled target domains. Our goal is to leverage the pre-trained high-level semantic knowledge in GPT2 and use the unlabeled IMDB data to help align this into our sentiment classification setting. Comprehensive experimental results are further summarized in and show promising directions for future work.
Appendix F
8 Conclusion
We introduce CRONOS, the first algorithm to successfully apply convex neural networks to high- dimensional datasets with competitive performance. CRONOS leverages ADMM for robustness and parallelization, Nyström preconditioning for fast convergence, and JAX for GPU acceleration with large data. We extend this framework with CRONOS-AM: alternating minimization to adapt convex reformulation for deeper networks of arbitrary architecture. CRONOS comes with strong theoretical guarantees that match its practical performance, a rarity for optimization algorithms in deep learning. Experiments on large benchmark datasets in image and language tasks validate the efficacy of our algorithms. CRONOS performs better or comparably against other prevalent optimizers, but with virtually zero tuning of hyperparameters. We achieve validation accuracy of 88.47% on ImageNet binary classification, and 93.91 % on IMBd sentiment classification. Our goal is to advance an innovative, alternative paradigm in deep learning through convex networks, thus enhancing both efficiency and interpretability.
novel algorithm which utilizes
Our results raise several important questions for future work and present many exciting avenues for continued research. Can we provide a convergence guarantee for CRONOS-AM that shows an advantage over stochastic first-order methods? Additionally, the efficacy of CRONOS on NI.P tasks suggest that investigating the performance of CRONOS on harder language datasets is an interesting direction, with potential for scalability on multi-GPU or TPU settings with other modalities.
10

---

Acknowledgments. This work was supported in part by National Science Foundation (NSF) under Grant DMS-2133248; part by the NSF CAREER Award under Grant CCF-2236829; in part by the U.S. Army Research Office Early Career Award under Grant W911NF-21-10022;; in part by the Office of Naval Research under Grant N00014-24-1-2164 We thank Lucy Woof and Peter Hawkins for support throughout and many insightful discussions.
References
Yossi Arjevani, Yair Carmon, John C Duchi, Dylan J Foster, Nathan Srebro, and Blake Woodworth. Lower bounds for non-convex stochastic optimization. Mathematical Programming, 199(1-2):
165-214, 2023.
Haim Avron, Petar Maymounkov, and Sivan Toledo. Blendenpik: Supercharging lapack's least- squares solver. SIAM Journal Scientific Computing, 32(3):1217-1236, 2010.
Francis Bach. Sharp analysis of low-rank kernel matrix approximations. In Conference Learning Theory, pages 185-209. PMLR, 2013.
Yatong Bai, Tanmay Gautam, and Somayeh Sojoudi. Efficient global optimization of two-layer relu networks: Quadratie-time algorithms and adversarial training. SIAM Journal Mathematics of Data Science, 5(2):46 2023.
Yoshua Bengio, Nicolas Roux, Pascal Vincent, Olivier Delalleau, and Patrice Marcotte. Convex Neural Networks, Advances in Neural Information Processing Systems, 18, 2005.
Avrim Blum and Ronald Rivest. Training 3-node neural network is np-complete. Advances in Neural Information Processing Systems, 1, 1988.
Stephen Boyd, Neal Parikh, Eric Chu, Borja Peleato, and Jonathan Eckstein. Distributed optimization and statistical learning via the alternating direction method of multipliers. Foundations and Trends in Machine Learning, 3:1-122, 2011.
James Bradbury. Roy Frostig. Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPP programs, 2018. URL
betp:///ithut com google/jax
Pawet Budzianowski and Ivan Vulić. Hello, it's gpt-2-how pretrained language models for task-oriented dialogue systems. arXiv preprint wVXiv:190..05774 2019.
help you? towards the of
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in Neural Information Processing Systems, 32, 2019.
Aaron Defazio and Konstantin Mishchenko. Learning-rate-fres learning by d-adaptation. In Interna- tional Conference on Machine Learning, pages 7449-7479. PMLR, 2023.
Wei Deng and Wotao Yin. On the global and linear convergence of the generalized alternating direction method of multipliers. Journal of Scientific Computing, 66:889-916. 2016.
Michal Derezinski, Feynman Liang, Zhenyu Liao, and Michael W Mahoney. Precise expressions for random projections: Low-rank approximation and randomized newton. Advances in Neural Information Processing Systems, 33:18272-18283, 2020.
Michal Derezinski, Jonathan Lacotte, Mert Pilanci, and Michael W Mahoney. Newton-less: Spar- sification without trade-off's for the sketched newton update. Advances in Neural Information Processing Systems, 34:2835-2847, 2021.
Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.0480) 2018.
Theo Diamandis, Zachary Frangella, Shipu Zhao, Bartolomeo Stellato, and Madeleine Udell. Genios: an (almost) second-order operator-splittngg solver for large-scale convex optimization. arXiv preprint arXiv:2310.0833 2023.

---

Jim Douglas and Henry H Rachford. On the numerical solution of heat conduction problems in two and three space variables. Transactions of the American Mathematical Society, 82(2):421 1956.
Jonathan Eckstein and Dimitri P Bertsekas. On the DouglasRachford splitting method and the proximal point algorithm for maximal monotone operators. Mathematiccal programming, 55: 293-318, 1992.
Zachary Frangella, Joel A Tropp, and Madeleine Udell. Randomized Nyström preconditioning. SIAM Journal on Matrix Analysis and Applications, 44(2):718-755 752, 2023a.
Zachary Frangella, Shipu Zhao, Theo Diamandis, Bartolomeo Stellato, and Madeleine Udell. On the (linear) convergence of generalized newton inexact admm. arXiv preprint arXiv:2302.03860 038663,, 2023b.
Rong Ge, Furong Huang, Chi Jin, and Yang Yuan. Escaping from saddle pointsonline stochastic gradient for tensor decomposition. In Conference on Learning Theory, pages 797-842. PMLR, 2015.
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari, Linearized two-layers neural networks in high dimension. The Annals of Statistics, 49(2):1029-1054, 2021.
Gene H Golub and Charles FF Van Loan. Matrix Computations. Johns Hopkins University Press, 2013.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: tion. In International Conference on Machine Learning, pages 1842-1850. PMLR, 2018.
Preconditioned stochastic tensor optimiza-
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural Tangent Kernel: Convergence and generalization in neural networks. Advances in Neural Information Processing Systems, 31, 2018.
Sungyoon Kim and Mert Pilanci. Convex relaxations of relu neural networks approximate global
optima in polynomial time. In International Conference on Machine Learning, 2024.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv e-prints, pages arXiv-1412, 2014.
Alex Krizhevsky and Geoff Hinton. Convolutional deep belief networks on cifar- 10. Unpublished manuscript, 40(7): 2010.
Jonathan Lacotte and Mert Pilanci. Effective dimension adaptive sketching methods for faster regularized least-squares optimization. Advances in Neural Information Processing Systems, 33: 19377-19387, 2020.
Jingwei Liang, Jalal Fadili, and Gabriel Peyré. Local convergence properties of douglas-rachfoor and alternating direction method of multipliers. Journal of Optimization Theory and Applications, 172:874-913, 2017.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 17100:171, 2017.
Michael W Mahoney al. Randomized algorithms for matrices and data. Foundations and Trends in Machine Learning, 3(2):123-23) 2011.
Per-Gunnar Martinsson and Joel A Tropp. Randomized numerical linear algebra: Foundations and algorithms. Acta Numerica, 29:403-572, 2020.
Xiangrui Meng, Michael A Saunders, and Michael W Mahoney. Lsm: A parallel iterative solver for strongly over-or underdetermined systems. SIAM Journal on Scientific Computing, 36(2): C95-C118, 18, 2014.
Aaron Mishkin, Arda Sahiner, and Mert Pilanci. Fast convex optimization for two-layer relu networks: Equivalent model classes and cone decompositions. In International Conference on Machine Learning, pages 15770-15816. PMLR, 2022.
12

---

Yu Nesterov. Gradient methods for minimizing composite functions. Mathematical Programming. 140 2013.
David Noever and Samantha E Miller Noever. Overhead mnist: A benchmark satellite dataset. arXiv preprint wrXiv:2102.04260 2021.
Hua Ouyang, Niao He, Long Tran, and Alexander Gray. Stochastic alternating direction method of multipliers. In International Conference on Machine Learning, pages 80-88. PMLR, 2013.
Yuyuan Ouyang, Yunmei Chen, Guanghui Lan, and Eduardo Pasiliao. An accelerated linearized alternating direction method of multipliers. SIAM Journal on Imaging Sciences, 8(1 8(1):644-681, 2015.
Brendan Odonoghue and Emmanuel Candes. Adaptive restart for accelerated gradient schemes. Foundations of Computational Mathematics, 15:715-732, 2015.
Brendan Odonoghue, Eric Chu, Neal Parikh, and Stephen Boyd. Conic optimization via operator splitting and homogeneous self-dual embedding. Journal of Optimization Theory and Applications, 109:1042-1068, 2016.
Mert Pilanci and Tolga Ergen. Neural networks are convex regularizers: Exact polynomial-ime convex optimization formulations for two-layer networks. In International Conference on Machine Learning, pages 7695-7705. PMLR, 2020.
Mert Pilanci and Martin J Wainwright. Newton sketch: A near linear-time optimization algorithm with linear-quadratic convergence. SIAM Journal Optimization, 27(1):205- 2017.
Ali Rahimi and Benjamin Recht. Weighted sums of Random Kitchen Sinks: Replacing minimization with randomization in learning. Advances in Neural Information Processing Systems, 21, 2008,
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet elassifiers generalize to imagenet? In International conference on machine learning, pages 5389-5400
PMLR, 2019.
Jonathan S Rosenfeld. Scaling laws for deep learning. arXiv preprint arXiv:2108.0766 2021.
Jakob Roth, Martin Reinecke, and Gordian Edenhofer. Jaxbind: Bind any function jax. arXiv preprint 4rXiv:2403.08847, 2024.
Ernest K Ryu and Wotao Yin. Large-scale convex optimization: algorithms & analyses via monotone operators. Cambridge University Press, 2022.
Othmane Sebbouh, Robert M Gower, and Aaron Defazio. Almost sure convergence rates for stochastic gradient descent and stochastic heavy ball. In Conference on Learning Theory, pages 3935-3971.
PMLR, 2021.
Bartolomeo Stellato, Goran Banjac, Paul Goulart, Alberto Bemporad, and Stephen Boyd. Osqp: An operator splitting solver for quadratic programs. Mathematical Programming Computation, 12(4): 637-672, 2020.
Joel A Tropp, Alp Yurtsever, Madeleine Udell, and Volkan Cevher. Fixed-rank approximation of positive-semideffiite matrix from streaming data. Advances in Neural Information Processing Systems, 30, 2017.
Joel A Tropp, Alp Yurtsever, Madeleine Udell, and Volkan Cevher. Streaming low-rank matrix approximation with an application to scientific simulation. SIAM Journal on Scientific Computing, 41(4):32430-A243,, 2019.
Madeleine Udell and Alex Townsend. Why are big data matrices approximately low rank? SIAM Journal on Mathematics of Data Science, 1(1):144-160, 2019.
Martin Wainwright. High-dimensional Statistics: non-asymptotic viewpoint, volume 48, Cam- bridge University Press, 2019.
13

---

David P Woodruff et al. Sketching as a tool for numerical linear algebra. Foundations and Trends@ in Theoretical Computer Science, 10(1-2):1117, 2014.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXin:1708.07747, 2017.
Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael Mahoney. Adahessian: An adaptive second order optimizer for machine learning. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pages 10665-10673, 2021.
Ming Yuan and Yi Lin. Model selection and estimation in regression with grouped variables. Journal of the Royal Statistical Society Series B: Statistical Methodology, 68(1):49-67, 2006.
Xiaoming Yuan, Shangzhi Zeng, and Jin Zhang. Discerning the linear convergence of admm for structured convex optimization through the lens of variational analysis. The Journal of Machine Learning Research,
2020.
Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive methods for nonconvex optimization. Advances Neural Information Processing Systems, 31, 2018.
Shipu Zhao, Zachary Frangella, and Madeleine Udell. Nysadmm: faster composite convex opti- mization via low-rank approximation. In International Conference on Machine Learning, pages 26840. PMLR, 2022.
14

---

A Related Work
Convex Neural Networks. Given the recent impressive capabilities of neural networks, there has been much work to develop strategies to avoid certain computational difficulties inherent in non- convex optimization. [Bengio et al. [2005] is an early work in this direction; they treat 2-layer neural networks convex models but require the first layer to be fixed. Random features regression is similar approach, where the first layer is randomly initialized, and only the last layer is trained trained|Rahimi Rahimi and Rechi [2008). More recently, it has been shown for certain classes of wide neural networks that training is equivalent to solving a convex kernel regression problem Jacot et al. [2018). Wide neural networks of this form are said to be in the "kernel regime". However although random feature regression and kernel regime neural networks yield convex optimization setups, they are known to exhibit inferior performance standard neural networks Chizat et al. 2019), Ghorbani et al. 2021
This makes them less attractive relative to the convex eeformulation of (3).
Motivated by the strong optimality guarantees associated with the convex reformulation other lines of work directly try and solve (3) using techniques from convex optimization. Two relevant
works and the current state-of fiffthherrt, are Mishkin al. [2022] and Bai et al. [2023]. Mishkin et al.
[022] develops a variant of restarted-FSSTA (r-FISTA) Nesterov 2013). (Odonoghe and Candes
2015 to solve with experiments
MNIST and CIFARTO. Although CRONOS and
et all 2022]> 's r- F-STAA have similar theoretical complexity profiles, we observe that CRONOS is
able to outperform r-FISTA in terms of scalability, as is able to handle large datasets ImageNet and IMDb. In contrast, the r-FISTA solver of Mishkin et al. [2022] was unable to scale to all of
CIFAR-I0 and MNIST. More closely related is the work of Bai al. [223] J, who directly applies
ADMM to ( However this strategy is unsuitable for large problems as it requires forming and
factoring a sizable matrix. CRONOS alleviates this issue by applying NysADMM, which leads to fast subproblem solves that only require matvecs.
ADMM methods. ADMM has long history that goes back to the work of Dfoullls and Rachford 1956] on efficient numerical solution of the heat equation. For a complete historical discussion, see
the survey [Boyd et al. [ the book Ryu and Yin [022]] The popularity of ADMM in machine learning originates with the survey Boyd et al. [0 which demonstrated ADMM could be applied successfully various machine learning problems. Since then, there have been numerous works, 2013 2015), Deng and Yin 2016). Zhao el al. [022]. that have developed variants of ADMM which aim to tackle the large problem instances that are ubiquitous in the era of big data.
Ouyang et al.
Randomized Numerical Linear Algebra. Randomized numerical linear algebra (RandNLA) is a field that combines randomization with fundamental numerical linear algebra primitives to speed up computation dramatically Mahoney et al.] [001 Woodruff et al. 2014]), Martinsson and Tropp [020)_ - RandNLA has been successful in acceleratiog a variety of important computational problems, ranging from low-rank approximation Trrppp et al. [2077] 2019), solving linear systems Avron alg [2010]. Meng et al. [2014), [acotte and Pilanci [2020]. Frangella et al. [2023a]. to convex
optimization Pilanci and Wainwright [2017]. Derezinski et al. [2021], Zhao et al. [2022] The most
relevant work our setting is Frangella al. [2023a], which introduced Nyström PCG, the algorithm
CRONOS uses to efficiently solve the u-subproblem (5). We are the first work to apply Nyström preconditioning the context of convex neural networks. By leveraging this with the structure of (4) and JAX based GPU acceleration, CRONOS is able to be both fast and highly scalable.
B Additional algorithms
B.1 Randomized Nyström approximation and Nyström PCG
In this section, we give the algorithms from Frangella et al. [2023a] for the randomized Nyström approximation and Nyström PCG. Algorithm constructs a low-rank approximation /l = ~U to the input matrix H. A is then used to form the Nyström preconditioner P, which (along with its inverse) given by:
P 1 U(A + pT\UT + + -U-- +
p-1 = (s, + (A + - + (I-UU".
15

---

Algorithm 3 Randomized Nyström Approximation
Require:
psd matrix H f (R), sketch size
Ω = randu(d,
Gaussian test matrix
f qr())
thin QR decomposition
Y HΩ
matvecs with H D compute shift D add shift for stability Cholesky decomposition triangular solve thin SVD
eps(norm(Y:
Y, = Y + us? +
C chol(?"
B Yu/C
U, U,E,, ~ =svd(B, [B, 0)
= max - -1 v remove shift, compute eigs return Nyström approximation A UÀUT
In CRONOS, μ as this corresponds to the value of the regularization parameter in the linear system arising from (5). The most important thing to note about Algorithm is that the time compute a matrix-vector product with H is Taw. the asympotic cost is ( my/) Zhao et al. [2022]
The preconditioner is then deployed with Nyström PCG (Algorithm 4 at each iteration to efficiently solve
Algorithm 4 Nyström PCG
Require: psd matrix H, righthand side r, initial guess To. regularization parameter p. sketch size 8,
tolerance 8 [U, A]
wr RandboiizaaYystiinApppoommmooo
(H pT)
Yo = P-lwo
Po y/0
while ||l|2 ||| > do
(H + pl Po
(w) yo) (PO v)
Zo + apo
= wi ac
y = P Iw
3 (w (wFy)//(nr (wg yo)
Ji < I, uh f w, Po + y + Spo. yo +- y
end while
return approximate solution a
B.1.1 Setting hyperparameters in Nyström PCG
Nyström PCG has two parameters: the rank and the tolerance The tolerance is easy set any summable sequence suffices to ensure CRONOS will converge. Thus, we recommend setting the PCG tolerance εд 1-1 at each iteration. This is consistent with popular ADMM solvers such as SCS, which uses the same sequence in its CG solver [Odonoghue et al.) 2016)
The rank is also not difficult to set. In our experiments, rank of 20 worked well uniformly. So, we recommend this as a default value. However, users wish to make the rank r close to the effective dimension (see Appendix C). they can use the adaptive algorithm proposed in Frangella et al. 2023a to select the rank. This approach iteratively builds the preconditioner and is guaranteed
to terminate when is the order of the effective dimension with high probability. This guarantees that Nyström PCG will converge to tolerance e in O (log(1/e)) 1//=)) iterations (see Appendix C
B.2 CRONOS-AM
We present pseudo-code for CRONOS-AM in Algorithm 5
DAdapted-Adam Addm we run the algorithm
for 1-epoch to perform the approximate minimization.
16

---

Algorithm 5 CRONOS-AM
Require: penalty parameter p, forcing sequence {8
k=1 rank parameter
repeat
Approximately minimize (u of+1, I
with respect to (u, U, s) using CRONOS obtain
Run CRONOS for 5 iterations
k+1 Approximately minimize 7 with respect w1:L-2 using DAdapted-Adam obtain w w1:L-2 > Run DAdapted-Adam for epoch
until convergence
C Proofs of main results
C.1 Proof of Proposition 6.1
Proof. Recall by definition, that F; D,X, and G; (2D; - I)X, where D, E Rax* is a diagonal matrix where Das E {0,1}. Consequently,
D < I and (2D; 1)² =/==1. =I.
Thus, by the conjugation rule, we conclude:
D2D < x" and X" (2D; 3
Recalling A < B implies (A) < (B), and that FT = = DJX and G. = X"D, X, the desired claim follows from the preceding display.
C.2 Proof of fast u-update using Nyström PCG
The key quantity in establishing fast convergence of Nyström PCG is to select the rank to be the level the effective dimension the matrix defining the linear system Frangella et al.) [2023a] Given y > 0 and symmetric positive definite matrix H, the effective dimension given by:
de (H) trace (H p
Roughly speaking, (H) gives smoothed count of the eigenvalues of larger than the parameter > O. We now recall the following result which bounds di (H) when the eigenvalues decay polynomially.
Lemma C.1 (Effective dimension under polynomial decay.). Let 0, 1/2 and let H N be a symmetric posilive semidefintte matrix. Suppose the eigenvalues of H satisfy:
Ay(1)) 00 0 !P]
Then,
d (H) S Cμ 29
where C' > 0 is some constant.
For a proof, see Section C. of Bach 2013) commence with the proof of Proposition 6.3
With these preliminaries out of the way, we now
Proof. Let H denote the randomized Nyström approximation of H output by Algorithm and define
E H A. By Lemma C.1 with y and Assumption 6.2 we have that d/__ ()) = 0(1). As
r (log (t)) - ( H)) + log it follows from Lemma A.7 Zhao et 2022 that
P (| < 2) 1-
We also obtain from Lemma 5.4 of Frangella al. [2023a] that Xx(( 1. Thus, applying Proposition 4.5 of Frangella et al. 2003a gives:
1/2 (H P 1/2) 11 + |E 4.
17

---

Moreover, from the convergence analysis for CG, we know that Golub and Van Loan 2013
VM(P-1/2(H+ + 1) - (M(P-1/2(H + 1)P-1/2) +1)
u k+1 xk+1 1 2
| 1aktt 1++--
(8)
Now,
[I(H 1)u b11 11(11 1)(attt तا VAr(H +))
a++ 11
Hence, by the last display with 8). we reach
[1 7) 1+ bll < VAI(H 1) "0 || (1/3)*,
The desired claim immediately follows from the last display.
C.3 Proof of Theorem 6.4
CRONOS convergence is a consequence of the following theorem, which a simplified version of Theorem of Frangella et al. 2023b) specialized to the setting of Algorithm
Theorem C.2 (Simplified Theorem I. Frangella al. [2023b]). Let p" denote the optimum of (4) For each K 2 1, denote kK++1 * Ck u², ok+1 # Σk vt, and gK+1 = Σk*
where (i*) Jx21 and {=* Jk21 are the iterates produced by Algorithm with forcing sequence h221 with ul = 0, v = 0, =0, A O, and v² = O. Let k--1 denote the exact solution of (5) at iteration Suppose that ||u a++1| S 8th for all E [K]. Then, the suboptimality gap satisfies
((Fa*+1 v) + B|| 2|,11 1( 00 O(1/K).
Furthermore, the feasibility gap satisfies
[[[ [2ar] G
K+1
[K+-1
1 K
g++
=
Consequently, after O(1/e) iterations,
E(Fu K v) +
0| K+1 1a.1 -)) Sc and
Tzar
k+1
k+
z k++
S
Proof. From Theorem
C.2 suffices to show that for all [K] that
| k+1 < 8/.
To this end, observe as O log and Nyström PCG at the keth iteration runs for O (log (+)
it follows from Proposition 6.3 that with probability 1 1- for all & E [K] the output *+1 ++ PCG satisfies:
Nyström
[I(H + 1)u h+1 1 < 8 vk [K].
Consequently, with probability at least
I k+1
[[ + I)(w k+1 5t ++1 II( x++1 < I|u k+1 k |1H41
= [H + 1) k+1 -
( 1 S [[ + +Ju - b² < 8t
Thus, we can invoke Theorem C.2 to conclude the proof of the first statement.
We prove the second statement. Recall that the cost of multiplying vector by H costs O P2). The cost of CRONOS is dominated by two parts: (1) the cost to construct the Nyström preconditioner, and (2) the cost to solve (5) at each iteration. All other operations cost O ((n + d)P) or less. For a rank 0, the cost of forming the Nyström approximation is O(dPP2r) (see Appendix [B.1p. In
Theorem 6.4 O [log
the total cost of constructing the preconditioner is ((ddP)) By Proposition 6.3 with probability at least 1 Ç, the total cost of solving for iterations given by:
0
ndd22 log
1 5h
< KO
ndd2 log
1 OK
)) = ndP22
18

---

where the last equality follows as K = O (1/c). Thus, with probability at least the total complexity of Algorithm is
CCRONOS =6 =6(md)) ((ndP2
c
D
When does the optimal value of the convex reformulation coincide with the optimal value of the original objective?
An important practical question is when does the optimal value of the convex program Eq. coincide with the optimal value the non-convex problem Eq. (2) Pilanci and Ergen 2220 established that the optimal value of the full convex reformulation agrees with that of Eq. but this is computationally intractable. is not apparent apriori that the optimal value the subsampled convex program 3
will agree with the optimal value of Eq. 2
Mishkin al. 2022] has provided an answer to when the solutions of Eq. (3) and Eq. (2) agree. Namely, they show the optimal values agree under the following reasonable conditions: (1) the number of neurons m in (2) must exceed the number of non-zeros in the optimal weights of the convex model, and (2) the sampled ReLU patterns in (3) must contain the ReLU patterns in the
optimal solution to (2.. While this result comforts us that using the subsampled convex model is principled, it is non-quantitative, and its conditions are difficult to verify.
Fortunately, the recent work Kim and Pilanci 2024 provides a more quantitative bound the gap between the optimal values of (3) and (2). Let PMIP sCvx denote the optimal value of (2). and PMLP.CVX denote the optimal value of (3). Theorem 2.1 in Kim and Pilanci 2024] establishes that m = 52 (log(n)) and P m/2, then with high probability:
PMVP S PMLP S C (P
where C is some absolute constant. This result shows that for (((og(n)), PMLP cvx deviates from NCVX by modest logarithmic factor. Moreover, two-layer ReLU networks universal approximators, is often the case that the network is capable of perfectly fitting the training data, in which case PMLP-CVX PMLP.NCVX O.
Combining the preceding result with Theorem 6.4 we see that CRONOS will deliver solution with comparable objective value to PMLP-NCVX in polynomial time. When PMPL-NCV O or the conditions Mishkin et al. [2022] hold, CRONOS can produce solution arbitrarily close to PCLP-NNVx in polynomial time.
E CRONOS for general smooth losses
This section shows how CRONOS can be extended to general smooth losses beyond the least-squares loss.
Beyond least-squares: The NysADMM framework of Zhao et al. [022) applies to losses beyond least-squares loss. Instead of solving (5), NysADMM replaces {(Fu) with a (regularized) second-order Taylor expansion about the current iterate nt, to obtain the subproblem.
argmin, (!(Fu*) (u u F Ve(Fut,u)
+ uu u² T ( +al)(u w²)
2
6 2 lu lu-vi th + 1 LIGu s + 112
2
The solution of 9 may be found by solving the following linear system:
(1 1)u 6².
where H [FT G.G. Note that the matrix defining the linear system in (9) changes from iteration to iteration. Consequently, the preconditioner should be periodically updated to maintain good performance. Other than this, there are no changes to CRONOS. CRONOS for general smooth losses is presented in Algorithm 6 Similar Algorithm the convergence of Algorithm can be established using results from Frangella et al. 2023b
19

---

Algorithm 6 CRONOS for general smooth losses
Require:
penalty parameter p. forcing sequence {8 Jk=1 rank parameter 7, preconditioner update frequency f repeat
if' k is multiple of f then
[U, = RandNysAppx( ( GT GT, & Update Nyström preconditioner
every iterations end
Use Nyström PCG to find th that solves ( within tolerance 8*
x - prox 11l-112 (2 + +*)
b Primal e update
ghtl f (Gu k+1 + +1+
b Slack 8 update
x t X + 2u (u 1l - 1)
Dual À update
whl b is (Gu gltl)
Dual update
until convergence
F
Experimental details and additional results
In this section, include additional details relevant our experimental design. We also provide additional numerical results not present within the main paper.
F1 Hardware
All experiments were performed on an RTX-4090 GPU with 24 GB of memory and 100tFLOPS in JAX functional code. We utilize x86-64 CPU architecture with Ubuntu 22.04 ÓS.
F.2 Hyperparmeters for CRONOS and CRONOS-AM
In our experiments CRONOS (including when used subproblem solver in CRONOS-AM run for 5 ADMM iterations. The number of PCG iterations varies from 5-50 depending upon the task. The rank of the Nyström preconditioner varies from r = 10 to r = 20. The value of p is varied from 0.001 depending upon the task. For CRONOS-AM DAdapted-Adam W is always run for I epoch to get the non-convex weights.
In practice, P can be selected adaptively via residual balancing. Which increases or decreases in such a way that ensures the primal and dual residuals (which measure the quality of the approximate solution) are balanced [Boyd et al., 2011] This is also common strategy employed in practice by ADMM solvers to set p [Stellato al. 2020 Diamandis et al. 2023). Note however, unlike deep learning optimizers which are sensitive to the learning rate, CRONOS is guaranteed to converge for any fixed value of p [Frangella et al., 20236]. So while a judicious choice of p may lead faster convergence, we are not restricted to some small band about the best to ensure convergence. This differs significantly from the non-convex case, where the learning rate not within some small band, non-convergent behavior can arise.
E3 Datasets
Fashion MNIST. The Fashion-MNIS: dataset is collection of 70,000 grayscale images, each of size 28x28 pixels, categorized into classes. includes 60,000 training samples and 10,000 test samples, and is intended as a more challenging replacement for the traditional MNIST dataset of handwritten digits. We use this dataset for the multi-class classification experiment.
Cifar10 dataset. This dataset involves 60,000 images of size 32x32x3 in 10 different classes, with 6,000 images per class. The dataset is divided into 50,000 training images and 10,000 test images, and is commonly used in baseline image classification tasks. Previous work in convex neural network applications have been to only reach binary classification classes on Cifar10, usually downsampled setting.
20

---

Table Validation accuracy achieved by CRONOS on IMDB, across three different settings, with multiple varying batches averaged across seeds.
IMDB-NFT IMDB-FT
IMDB-DA
74.67%
94.14%
76.15%
ImageNet Dataset. ImageNet is a large-scale dataset with approximately 14 million images have been hand-annotated by the project to indicate what objects are pictured. Each image is padded to 312x512x3 for binary classification. We provide downsampled version of size 171x171x3 of ImageNet order to examine scaling abilities of all methods. Our peak validation accuracy with CRONOS is 88.72%
Food-5k Dataset. The Food-5K dataset is a collection of 5,000 images with binary classification label of I1 = food' and 11 non-food', Each image is padded to size 256x256x3 in all experiments. We select this dataset for the intermediate scaling of size between downsampled ImageNet171 and full size ImageNet511 12. Interestingly, we note that CRONOS performs significantly better on larger data and harder tasks with virtually zero tuning of hyperparameters.
IMDb Dataset. The IMDb dataset is a set of 50,000 movie reviews from the Internet Movie Database (IMDb) for binary sentiment classification. It is evenly split into 25,000 reviews for training and 25,000 reviews for testing. Each set consists of an equal number of positive and negative reviews, where the sentiment of each review is determined by the number of stars awarded by the reviewer. consists of only negative and positive reviews, without neutral. We perform all experiments with a batch size of 32, and pad all input sequences length 60 with the PAD token defined as 'EOS', Input to CRONOS is reshaped into (1500x60x768) with 1500 samples and an embedding dimension of 768. Our best accuracy on the GPT2 foundation experiments with CRONOS yields 94.14% 14% validation accuracy.
F.4 Details on Three NLP Experiments with GPT2
We conducted 14 NLP experiments grouped into 3 categories (no train GPT2-NFT, 1 epoch fine tuning GPT2-FT, and unsupervised GPT2-DA). In each experiment, we benchmark against varying seed and learning rate sweeps of AdamW as baseline. We select AdamW as baseline since is predominantly utilized in most NLP tasks. Our results show that CRONOS outperforms tuned AdamW in both settings in terms of speed and accuracy.
IMDB-NFT. We initially extract the GPT2 pretrained checkpoint and then immediately follow with CRONOS for sentiment classification. Notably, this setting does not involve any training loops with standard optimizers (such as AdamW), and the foundation GPT2 model does not see any labeled sentiment data from the IMDB dataset. We limit the amount of data seen by CRONOS to only batches only, in order to evaluate the efficacy of our method on low data regime settings.
Our second language model experiment uses the GPT2 pretrained checkpoint followed by one epoch only of training with the AdamW optimizer on default parameters with the standard dense linear layer head followed by CRONOS. Although the authors of BERT recommend 2-4 epochs for training in this setting, we aim to evaluate the performance of CRONOS with limited overhead to extract features for maximum efficiency. The resulting output of the dense layer is then input to CRONOS for the sentiment classification task, and demonstrates extremely promising results of approximately 94% validation accuracy after utilizing 2 batches of labeled data. This is particularly significant due to the lack of hyperparameter tuning in CRONOS, and its consistency across multiple trials of varying data batches.
IMDB-FT.
IMDC-DA. Finally we examine the setting of training for one epoch of AdamW on unlabeled IMDB data initialized from the GPT2 checkpoint. This experiment was motivated by the idea of subspace alignment. Since pre-trained GPT2 has vast semantic knowledge, even minor alignment to an unlabeled dataset should be beneficial to the eventual downstream task. We are also practically motivated to explore challenging settings with unlabeled data, since this is readily prevalent real world scenarios. We run one epoch of unlabeled IMDB data through pre-trained GPT2 to predict the next token in the language modeling setting. This new checkpoint then used for feature extraction and classification with our convex network and AdamW. This examines the potential of CRONOS on
21

---

4.5
1:0
(a) ImageNet 171
(b) ImageNet-177 trajectories
14B
170
(c) Fond-5k
(d) Food-5k trajectories
Figure 3: Results for ImageNet-177 and Food-5k
unsupervised domain adaptation settings in future work, which aims reduce the distribution gap between source and unlabeled target domains. We recognize that fopaii- tasks generally perform well when the feature distributions of the domains effectiveness of initializing from the broadly pretrained GPT2 model on the unlabeled IMDb dataset. Experimental results show promising directions for future work.
similar. Therefore examine the
F.5
Additional Results for the MLP experiment
Below in Fig we report results for various optimizers on Food-5k.
F.6
Convolutional Neural Network Experiments
Here we test the performance of CRONOS-AM on a 4-layer CNN consisting convolutional layers with average pooling, followed by two densely connected layers. The goal of this experiment to demonstrate the fcasibility of applying CRONOS-AR gives results for CIFAR-10 and ImageNet-1711 171. On CIFAR-10. 10, CRONOS-AM performs somewhat worse than the competition, but delivers excellent performance on ImageNet- 171, where it reaches high accuracy quite quickly. These experiments validate the applicability of CRONOS- A-M beyond MLPs. For future work, it would be interesting to explore the effectiveness of CRONOS- on other popular architectures.
to architectures aside from MLPs. Fig.
F.7 Multiclass Classification Experiment
To see how CRONOS performs in the multi-class setting, we experiment the Fashion MNIST dataset [Xiao et 2017). This yields classification problem with unique classes. We consider two-layer ReLU MLP with 64 neurons. For comparison we also consider Adam, AdamW, SGD+Mo- mentum, Shampoo, and Yogi. fully exploit GPU acceleration, the batchsize these methods set to 10, 000. As in the main paper, we consider a random grid of five learning rates. The results for experiment are provided T The table shows CRONOS delivers best performance by significant margin relative to its non-convex competitors. Morcover, see the non-convex stochastic
22

---

Solet
(a) Training a CNN on CIFAR-10
(b) ImageNet-177
Figure Training a CNN on ImageNet-171
optimizers very sensitive the learning rate. Again, CRONOS does not have this issue, as it does not require a learning rate.
Table 5: Results for Fashion MNIST
Two-layer ReLU MLP with 64 neurons
Optimizer
Validation Range
Learning
CRONOS
85. 8%
NA
Adam
[12.24, 78, 62]%
3.52
AdamW
[12.23, 78 61]%
3.52 10
SGD
[6.63, 75.32|%
1.97
10
Shampoo
[6.87, 78.49]%
2.00
10
Yogi
[19.11, 71.54]
3.52
10
Results for different random seeds
In this subsection we report results for the deep MLP experiments for each dataset for two different random seeds. Looking the figures, is clear the trends observed earlier the main paper are stable across the random seed.
23

---

Time
CIFAR-10
(b) ImageNet- 171
Thae
(c) Food-5k
(d) ImageNet
Figure 5: CRONOS ARONOS-AM&I vs. competitors on Deep ReLU MLP (Seed 2)
CIFAR-10
(b) ImageNet-177
(c) Food-5k
(d) ImageNet
Figure 6: Learning rate trajectories for Deep ReLU MLP (Seed 2)
24

---

CRANCOS per1
Tim
CIFAR-10
(b) ImageNct- 171
(c) Food-5k
(d) ImageNet
Figure 7: CRONOS AM vs. competitors on Deep ReLU MLP (Seed 3)
CIFAR-10
(b) ImageNet-171
(c) Food-5k
(d) ImageNet
Figure 8: Learning rate trajectories for Deep ReLU MLP (Seed 3)
25

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: All claims made in the abstract are supported with theoretical results and numerical experiments within the main paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We mention several directions for improving the algorithms developed here within the conclusion.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
26

---

Answer: [Yes]
Justification: All our results are accompanied by the relevant assumptions, and we provide correct proof for each result.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Yes, the experimental setup in our paper is clear.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. .., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
27

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Yes, have provided the code generate the experimental results presented in the paper in a zip file.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (kttps://nnis..c)
mublic//uudes//oodeuuiiiis
for more details.
While we encourage the release of code and data, understand that this might not be possible, so No is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (httpa
7 //nips. httpps.ccpuublle htce/ublic/guuidess Sotpeeee//ooSSSumissssssee itt.saammmssiooooooooooooo asionnllicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We explain the experimental setup within the paper and provide the code used to generate rthe results in the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Where relevant we have provided shaded quantile plots. We also provide results across different random seeds when it is relevant.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
28

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We mention in the paper all the computational resources we used to run the experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http://oorrrracccooouooclaaaaaaa
Answer: [Yes]
Justification:
We have reviewed the code of ethics and have conformed them.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our paper focuses on developing two new algorithms for training neural networks, therefore it does not have any direct positive negative societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We give attribution to all assets used within the paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
30

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, datasets.paperssti www.ppppptsthhhccc thcode www.coo//a/a..o datasets
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out the asset's s creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
The code used to produce the results in this paper uploaded in zip file.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: No research with crowdsourcing human subjects was performed in this paper.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
31

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
32


Reviews:
[{"review_id": "a4yy1hUhH6", "reviewer": "Reviewer_iBsY", "summary": "In the paper, two neural network optimization methods, CRONOS and CRONOS-AM, are proposed. The authors provide theoretical complexity analysis and convergence proof of the algorithms, and perform experiments on large datasets of image and language tasks to verify the effectiveness of the algorithm. As can be seen from the experimental results, CRONOS provides better or comparable performance compared to optimizers such as SGD, Adam, and AdamW, but hardly any tuning of hyperparameters. This is the first time that convex networks have been applied to large datasets with competitive performance, which has great significance.", "questions": "How is unsupervised training implemented on IMDC-DA?", "limitations": "The authors should conduct experiments on more datasets to test the effectiveness of the proposed algorithms.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Two neural network optimization methods, CRONOS and CRONOS-AM, are proposed. \n2. The authors provide theoretical complexity analysis and convergence proof of the algorithms.\n3. Experimental results indicate that CRONOS provides better or comparable performance compared to optimizers such as SGD, Adam, and AdamW, but hardly any tuning of hyperparameters. \n4. This is the first time that convex networks have been applied to large datasets with competitive performance, which has great significance.", "weaknesses": "1.The symbols W1j, w2 are not consistent with their definition in line 77.\n2. Figure 2 is not referenced in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "** Weaknesses **\n1. Thank you for pointing this out! We have fixed this in the revision.\n\n**Questions**\n\nWe apologize for the typo and have corrected this for IMDB-DA. Thank you for raising the question! The Domain Adaptation experiments were aimed at minimizing the distribution shift between source and target domains. Our goal is to leverage the pre-trained high-level semantic knowledge in GPT2 and use the unlabeled IMDB data to help align this into our sentiment classification setting. \n\nWe run one epoch of unlabeled IMDB data through a pre-trained GPT2 to predict the next token in the language modeling setting. This new checkpoint is then used for feature extraction and classification with our convex network and AdamW. \n\nThis experiment was motivated by the idea of subspace alignment. Since pre-trained GPT2 has vast semantic knowledge, even minor alignment to an unlabeled dataset should be beneficial to the eventual downstream task. We are also practically motivated to explore challenging settings with unlabeled data, since this is readily prevalent in real world scenarios. We have added more numerical and experimental details for the IMDB-DA setting in the revision, and acknowledge this is an exciting area with much more work to be done!", "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": "EJcdGr5Twc", "reviewer": "Reviewer_inX1", "summary": "This work investigated the optimization of convex reformulated neural networks. The authors proposed to solve the sub-sampled convex optimization problem (for a two-layer ReLU network) with operator splitting (using ADMM), and used conjugate gradient method with Nystrom preconditioning to solve one subproblem in ADMM iterations. The proposed Nystrom preconditioning utilizes the low rank nature of data. The authors further proposed to extend the method to multi-layer networks through alternative minimization. The authors provided an efficient GPU implementation of the proposed algorithm using JAX and just-in-time (JIT) compilation.\n\nThe authors theoretically (1) verified that the subproblem in ADMM inherits the low-rank structure of data and thus can enjoy fast convergence through Nystrom preconditioning; (2) proved the fast convergence of the ADMM subproblem with conjugate gradient method; (3) and proved the convergence of the overall algorithm to the optimal solution of the sub-sampled convex optimization. The authors also showed that the proposed algorithm has a lower computation complexity compared to a previous work that directly applies ADMM, and thus is more scalable to larger datasets such as ImageNet.\n\nThe authors conducted extensive empirical experiments and claimed to be the first that successfully scale convex reformulated neural network optimization to large data such as ImageNet and large language modeling tasks with GPT2 architecture.", "questions": "1. How were the experiments on Imagenet actually conducted?\n    * Why were they binary classification tasks as described in Line 599?\n    * Did the authors also use a type of stochastic opitmization with CRONOS-AM so that it can be applied to more than 14M training images?\n\n2. In IMDB experiments, the authors did not compare CRONOS with existing optimizers as they did in Figure 2, while they did use 1 epoch of AdamW training before CRONOS optimization as they described in Line 622-624. Can the authors clarify the experiment settings more clearly and explain how can we decouple the effect of AdamW and CRONOS optimization?\n\n3. Notation $P$ was abused. It was used to denote the number of $D_i$ samplings in Section 3 and later the pre-conditioning matrix in Section 4.\n\n4. In line 91, it should be \"provided $m \\geq m^*$, for some $m^* \\leq n + 1$.", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. As the authors claimed, this work is the first practically useful algorithm for training neural networks in their convex reformulation on large and complex datasets. And the authors verfied this with promising empirical results.\n\n2. The approximate solver with conjugate gradient method plus Nystrom preconditioning looks a practical and promising approach and the authors did provide JAX implementation that will be useful for the following research.", "weaknesses": "1. The description of how the hyperparameter $P$ is selected is missing and the authors did not provide any discussion on this hyperparameter, but I suppose it is important for the following reasons.\n    * Although the authors proved the convergence of CRONOS to the global minimum of (4), the closeness of the global minimum of (4) to that of (3) depends on the effective sampling of $D_X$.\n    * According to Mishkin et al. [2022], the global minima of the two problems coincide if the hidden dimension is large enough AND the optimal $D_i$ are contained in the sampled subset.\n    * Bai et al. [2023] also show that to make the optimal objective of (4) effectively small, we need $P$ to increase linearly as the number of data increases.\n    * The selection of $P$ directly influences the computation cost as the authors analyzed in Section 6.4.\n\nHowever, as far as I observed, the authors did not mention the selection of $P$ in practice, while claiming the proposed CRONOS is almost hyperparameter tuning free. \n\n2. Following the first point, the authors did not mention the actual training time comparison of the proposed method and normal deep learning optimization methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weaknesses**\n\n1.) This is an excellent point! It is important to have guidance on how to select the hyperparameter $P$  in practice.\n Fortunately, large $P$ is not necessary to ensure the optimal values are comparable.\n The recent work [1] (see Theorem 2.1) on convex neural networks has shown that as long as the number of neurons satisfies $\\mathcal O\\left(\n \\log(n)\\right)$, the optimum of the convex reformulation will deviate from the optimum of the original two-layer ReLU problem by a multiplicative factor of at most $\\mathcal O\\left(\\sqrt{\\log(n)}\\right)$.\n This is an extremely modest multiple. \n Therefore from a practical standpoint, only a modest number of neurons is needed to ensure good performance.\n This agrees with practice, as in the experiments in the main paper we have only used $P=10$ across all experiments, but we obtain comparable or better performance relative to non-convex optimizers.\n\n In the revision we will be sure to mention this point. \n We will also include a more formal discussion in the supplement about when the optimal solutions actually coincide, which occurs when the conditions of Mishkin et al. (2022) mentioned by the reviewer hold. \n\n2.) Below is a table showing the runtime of CRONOS and CRONOS-AM to non-convex optimizers.\nNotably, the runtimes are comparable to standard methods, and we can conclude that the proposed methods have reasonable empirical complexity. In addition, we will include a more comprehensive table of runtimes in the revised document, summarizing NLP experiments with larger data sizes. \n## Optimizer Runtimes (s) on CIFAR-10 and ImageNet\n| Dataset | Cronos-AM | Adam | AdamW | D-Adapted Adam| SGD | Shampoo | Yogi |  \n| --------- | :---------:| :---------: | :---------: | :---------: | :---------: | :---------: | :---------: |\n| CIFAR-10 | 3.00 | 3.02 | 3.14 | 3.68 | 2.28 | 6.80 | 2.79|\n| ImageNet | 5.10 | 1.84| 3.14| 2.94|  2.48| 5.19| 1.98\n\n**Questions**\n\n1.)     We randomly selected two classes on ImageNet and performed binary classification. As the reviewer suggests, to be able to handle all of ImageNet would require replacing CRONOS in CRONOS-AM with a stochastic variant. We believe a stochastic variant of CRONOS-AM is the next step in advancing the practical deployment of convex NNs. It is an exciting direction for future work that we will mention in the revision.\n\n2.) Thank you for raising this concern! We have thoroughly added more experimental results, numerical tables, and setting details to the revised document for clarification. We conducted 14 NLP experiments grouped into 3 categories (no train GPT2-NFT, 1 epoch fine tuning GPT2-FT, and unsupervised GPT2-DA). In each case we benchmark against seed and learning rate sweeps of AdamW on the binary classification task. Please see Figures (c) - (f) in the rebuttal pdf for a sample of run time plots. CRONOS is able to outperform AdamW in both peak validation and in the time needed to reach it. In the GPT2-NFT (No Fine Tune) setting, the effects of CRONOS and AdamW are examined separately. Features extracted from a pre-trained GPT2 are given directly to each optimizer for individual classification. We experiment with 3 scenarios in the no fine-tune setting: immediate classification, mean-pool then classification, and attention-pooled then classification. In each case CRONOS achieves higher test accuracy than AdamW (with extensive grid search). \n \n\n3, 4.) Thank you for catching these! We will fix them in the revision.\n\n**References**\n\n[1] Kim and Pilanci, Convex Relaxations of ReLU Neural Networks Approximate Global Optima in Polynomial Time (2024)", "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": "Xfiuj1ATKj", "reviewer": "Reviewer_JCYR", "summary": "The authors introduce a reformulation of optimizing a two-layer neural network (one neuron in the second layer) to a convex optimization problem and suggest an algorithm to solve it. \nThis is done via reformulating a convex program based on sampled P activation patterns from the set of all possible ReLU activation patterns. The optimization is done from the space of P convex cones.  (I have some concerns/questions here).\n\nThen Proposition 3.1., shows how to take a step further to define it as a constrained generalized linear model that can be solved via ADMM (given by algorithm 1).\nThe first line of algorithm 1 requires solving a linear system (6) -- Solving (6) directly is too expensive due to large matrix sizes. \nThe Conjugate Gradient (CG) algorithm, needing only matrix-vector products, offers an alternative but converges slowly because of high condition numbers in typical machine learning data matrices, making CG impractically slow for repeated use.\n\nTo speed up CG convergence, they exploit the problem's approximate low-rank structure. The subproblem for v^{k+1} and s^{k+1} has a closed-form solution computable in O((n+d)P) time. The authors suggest applying the NysADMM algorithm, using Nyström Preconditioned Conjugate Gradient (NysPCG) to solve (6). NysPCG constructs a low-rank preconditioner for H+I, solving the linear system to \\delta-accuracy in O(log(1/\\delta)) iterations. \n\nThen CRONOS-AM is suggested for optimizing full DNN that integrates CRONOS with alternating minimization, allowing for the training of multi-layer networks with any architecture.\n\nThe theory seems to be robust and very interesting.", "questions": "* Consider changing the notation of relu it is so confusing specifically that the next character is multiplied and not added. \n\n* The authors define W^(1) and then use W_1 ?\n\n* I am missing the connection from (2) to (3):\n     A. Why sampling P activation patterns is enough? \n     B. Is it simple by sampling?\n\n* Define ||v||2,1 and provide an intuition on (4)\n\n* what are the reasonable conditions that (3) still has the same optimal solution as (2)?", "limitations": "See above", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper has many theoretical novelties and seems interesting from this perspective.\n\n2. The used tricks and formulation can open new directions for research. \n\n3. The paper aims to fill a big gap in deep learning --> theory to practice, which is a very important problem to tackle.", "weaknesses": "Small concerns:\n\n1. Writing: The authors assume the readers have robust knowledge in many fields and do not provide explanations, requiring so much time to understand and read the paper. \n\n2. The authors should be clear when defining their claims, e.g., Equation (1) does not define a classic 2-layer relu (with softmax). It is a one-hidden layer relu with a single neuron on the output layers. \n\nBig concerns:\n\n3. Experiments:  \na. What do the authors mean by binary classification on imagenet and cifar?!  \nb. Figure 2 seems to be wrong --  please add also the training accuracy to see if it is simply overfitting or what is happening.\n    * Adam is usually better than SGD, I encourage the authors to try the standard common parameters of ADAM.\n    * On what dataset is Figure 2 reporting the results? \n\nc. Figure 1 b is even worse, how come, known optimizers do not perform well on binary classification? and even reach 0 Accuracy\n\nd. NLP:\n   * What is the accuracy of GPT before running your algorithm on the same data, without any fine-tuning?\n   * I did not understand the baselines in the nlp task, do you use coronos in all of them,? if yes, why you did not compare your algorithm to other optimizers on the NLP task just simply to finetune or train one layer as a classifier?\n\nI don't feel that the experiments are convincing or even correct; I don't think that only cronos can optimize a binary classification problem in IMagenet and the other famous known algorithms fail.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weaknesses** \n\n1. Thank you for raising the point! We understand that the paper covers a wide range and utilizes ideas from diverse subjects such as deep learning, convex optimization, randomized numerical linear algebra, and high-performance computing. Unfortunately given the page constraints, we couldn't provide as much background information in the main paper as we would have liked. We are happy to include more background in the supplement of the revision with crossrefs in the main paper. If the reviewer has any thoughts on what would background info would be helpful, please let us know! \n\n2. Equation (1) does define a 2-layer ReLU networkin the regression setting, see [1, 2] .\n   The more familiar softmax case mentioned by the reviewer can also be covered by the convex reformulation, see [1, 3] for further details.  \n\n3a. By binary classification, we mean we look at two classes of CIFAR-10 and ImageNet and perform binary classification. \n        This is consistent with prior literature on convex neural networks such as [3, 4}. \n        \n3b. This occurred since the initial learning rate grid we used wasn't fine enough. \n      Our original grid included Adam's default learning rate of $10^{-3}$.\n       However this value doesn't work well in the architecture considered for these two classes of ImageNet.\n       The new plot is included in the rebuttal pdf Figures (a) and (b) with a finer grid.\n       Now, the competing methods don't all fail.\n\n3c.  Figure 2 shows the performance of Adam and SGD across different learning rates on CIFAR-10.\n       As per the previous response, the issue with figure 1b) has been revised with the finer grid in the rebuttal pdf. We have also revised the captions of our figures to indicate the dataset clearly.\n\n3d.  Thank you for raising this question! Since we utilize the features extracted from GPT2, the immediate output does not readily report a classification accuracy. We do perform experiments across 3 configurations of feature extraction with GPT2 (without fine-tuning) and benchmark against various trajectories of AdamW as a baseline. These three configurations are:\n\n*  Features extracted from pre-trained GPT2 followed by classification immediately with No Fine-Tuning (IMDb-NFT).\n*  Features extracted then followed by a mean-pool layer before classification.\n* Features extracted then followed by an attention-pooled layer followed by classification.\n \nIn all three settings above we benchmark against AdamW as a baseline, since this seems the most predominate in NLP applications. An example of these results have been included in the rebuttal pdf (Figures c -f), and we will definitely add more details to clarify experiment settings in the revision.\n\n**Questions**\n1.   Sure! We'll change this in the revision to avoid any confusion.\n2. Thanks for catching this!  We will fix it in the revision.\n3. Eq (2) is the standard non-convex training problem.\n   [2] proved that this admits an equivalent convex reformulation based on sampling patterns from the set $\\mathcal D_X$. Although this set is discrete, its cardinality grows as $\\mathcal O (r(n/r)^r)$, where $r$ is the rank of the data matrix. \nThis is extremely large unless the rank is small, so sampling all the patterns from this set is computationally infeasible.\n\n Eq (3) uses the same objective as the reformulation in [2], but only works with a subset of $P$ patterns from $\\mathcal D_X$. This result yields a convex optimization problem that is a computationally tractable alternative to Eq (2).\n    \nThis is a great question! Theorem 2.1 in [3] gives precise conditions under which the optimal values of Eq (1) and Eq (2) coincide. \nBut a simpler, more quantitative explanation is provided by Theorem 2.1 in the recent paper [5].\nThis theorem shows that as long as $P = \\mathcal O\\left(\\log(n)\\right)$, where $n$ is the number of data points, then\n\nOptimal value of Eq 2 $\\leq \\mathcal O(\\sqrt{\\log(n)}) \\times$ Optimal value of Eq 1.\n \nEvidently the optimal value of the subsampled convex reformulation Eq (2) is upper bounded by an extremely modest multiple of the optimal value of the non-convex problem Eq (1).\nIn this work, we only used $P = 10$ throughout all experiments, which was sufficient to obtain good results.\nIn the revision, we include a more careful discussion of these points. \n \nAs per your question on whether sampling is efficient, the answer is yes! \n We can obtain an element of $\\mathcal D_X$ by sampling a standard normal random vector $g \\in \\mathbb R^{d}$ and computing $\\textup{ReLU}(Xg)$, which costs $\\mathcal O(nd)$. Sampling $P$ patterns costs $\\mathcal O(ndP)$.\n\n4. Thank you for catching this! \n    We will make sure it is defined in the revision. \n \n    Sure, we can give intuition!\n    The $\\mathcal K_i$ in (3) are defined by linear inequality constraints. \n    By adding non-negative slack variables, these can be turned into linear equality constraints.\n    Then, organizing all the variables into the appropriate vectors $\\mathbf u, \\mathbf v,$ and $\\mathbf s$,\n    and compactly writing down the linear equality constraints, one arrives at (4).\n    As we note in the paper, this formulation was previously derived in [4], but we are happy to include a derivation in the supplement to help the paper be more self-contained.\n\n5.  Please see the response to question 3.\n\n**References**\n\n[1] Du et al., Gradient Descent Provably Optimizes Over-parameterized Neural Networks (2018)\n\n[2] Pilanci and Ergen, Neural Networks are Convex Regularizers: Exact Polynomial-time Convex Optimization Formulations for Two-layer Networks (2020)\n\n[3] Mishkin et al., Fast Convex Optimization for Two-Layer ReLU Networks: Equivalent Model Classes and Cone Decompositions, (2022)\n\n[4] Bai et al., Efficient global optimization of two-layer relu networks: Quadratic-time algorithms and adversarial training (2023)\n\n[5] Kim and Pilanci, Convex Relaxations of ReLU Neural Networks Approximate Global Optima in Polynomial Time (2024)", "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": "JTc94F45eM", "reviewer": "Reviewer_dmGV", "summary": "- This paper builds on work on reframing the optimization of 2-layer ReLU networks as convex optimization\n- The first contribution of this work is the use of the ADMM algorithm for solving the constrained optimization task from equation 4\n- Noting that need of repeatedly solve the inner problem in the ADMM method, they propose using a Nystrom-preconditioned solver to accelerate the inner problem, leading to Cronos\n- For deeper networks, the authors propose alternately minimizing the earlier model parameters (while freezing the final two layers) with D-adapted Adam, then optimizing the the final two layer parameters with with Cronos, leading to Cronos-AM\n- They show that this method can successfully optimize deep networks, leading to the first successful application of convex neural network application for deep networks", "questions": "- Would it be possible to have the results from Figure 1 and 2 in a table format for easier reading?\n- The appendix contains only the performance of Cronos on Fashion-MNIST on multiclass classification. It would be useful to include the performance of existing methods for a fair comparison\n- Are all other optimizers really achieving 0% accuracy on Binary ImageNet in Figure 1B? For binary classification achieving 0% would mean guessing the opposite class every since time, which seems like the classes were inverted. I am skeptical whether the plots in figure 1b are accurate.\n- Could we see the performance of D-adapted methods in figure 1 and 2?\n- What is the observed runtime of the algorithm?\n- The authors claim that the method is appealing because of the lack of hyperparameters, yet it seems that there are a number of hyperpameters associated with the Nystrom PCG algorithm (such as the rank and the tolerance). It would useful to include results showing the effect of these hyperparameters on the method's performance.", "limitations": "See the limitations section. Overall, I think this is an interesting line of work and successfully applying the convex formulation of two-layer networks for deep networks is a solid accomplishment, but the papers seems to be lacking proper comparisons to baselines and also does not thoroughly discuss aspects of the proposed methods (such as sensitivity of hyperparameters). I (the reviewer) am not particularly familiar with the work on convex-neural networks, however, so it is possible that I am undervaluing the theoretical contributions of this work.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- This work is the first work to successfully apply the convex reformulation for optimizing deep networks, albeit only the final two layers are optimized solely using this Cronos\n- The theory is robust and well supported\n- Compared to existing methods, Cronos does not require hyperameter tuning (such as learning rates), and appears to perform well compared to existing methods, albeit on binary classification tasks (for both vision and language)", "weaknesses": "- It seems that the use of D-adapted Adam for the majority of the network layers is necessary for the Cronos-AM to work. Therefore, it is unclear how important Cronos actually is\n- The authors claim that an advantage of Cronos is the lack of hyperparameters, but does not compare to D-adapted optimizers in figures 1 and 2, which similarly do not have hyperparameters\n- There is little to no discussion on the actual empirical runtime of the algorithm (for example, how many hours/minutes does the method take compared to existing methods?)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weaknesses**\n\n1. This is an excellent point that was unclear in the original submission. We have included figures for CIFAR-10 and ImageNet in the rebuttal document which include D-Adapted Adam (DAdam)  in the comparison. The difference in performance between CRONOS-AM and DAdam are visible on trajectory Figures (a) and (b). We have also summarized numerical comparisons in Tables 1 and 2 of the rebuttal document. It is evident that CRONOS-AM outperforms DAdam, which is shows the importance/value of CRONOS. \n\n We recommend DAdam for optimizing layers $1:L-2$ since it does not require tuning. This makes CRONOS-AM easier to use in practice. By applying the convex reformulation in the final two layers CRONOS-AM obtains better performance and faster convergence to solutions of better quality than DAdam. Alternatively, any other deep learning optimizer can be used for layers $1:L-2$.\n\nWe will include in the revision figures and plots similar to those in the rebuttal, to make the benefits of the convex reformulation in the last two layers clear. \n\n2.  Please see the preceding response.\n\n3.  Below is a table showing the runtime of CRONOS and CRONOS-AM to non-convex optimizers.\nNotably, the runtimes are comparable to standard methods, and we can conclude that the proposed methods have reasonable empirical complexity. In addition, we will include a more comprehensive table of runtimes in the revised document, summarizing NLP experiments with larger data sizes. \n## Optimizer Runtimes (s) on CIFAR-10 and ImageNet\n| Dataset | Cronos-AM | Adam | AdamW | D-Adapted Adam| SGD | Shampoo | Yogi |  \n| --------- | :---------:| :---------: | :---------: | :---------: | :---------: | :---------: | :---------: |\n| CIFAR-10 | 3.00 | 3.02 | 3.14 | 3.68 | 2.28 | 6.80 | 2.79|\n| ImageNet | 5.10 | 1.84| 3.14| 2.94|  2.48| 5.19| 1.98\n\n**Questions**\n1. Yes, this is a great idea! We agree this helps make things more clear, and shall include tables like those presented in the rebuttal pdf.\n\n2.  Yes, this is a good point! \n We plan to add these results to the revision. \n\n3. This occurred because our learning rate grid for tuning the competitors was not fine enough. We have re-run our experiments with a finer learning rate grid, and the revised plot replacing Figure 1.1b) is included in the rebuttal pdf Figures (a) and (b). The competing methods do not all fail. \n\n4. Yes! We have added this comparison. \n    Please see the rebuttal pdf.\n\n5.  Please see earlier response with table of runtimes, thank you. We will include a more comprehensive table of runtimes for all    experiments in the final document. \n\n6. The reviewer raises an excellent point! Nyström PCG does have two hyperparameters. \n    Fortunately, they are very easy to set.\n    For the tolerance, one can use any summable sequence to ensure convergence. \n    In particular, the popular ADMM solver SCS from [1] uses a sequence of $1/k^{1.2}$. \n    Another common heuristic is to use the geometric mean of the ADMM primal and dual residuals. This is what OSQP [2] uses, and also was used by [3].\n    We recommend using the summable tolerance as that is what the theory requires to obtain convergence. \n    Nevertheless, past work has found that both these strategies perform very similarly [3]. \n    We find this to be true in our experience as well.\n    We will include a discussion of how to set the tolerance in the revision.\n\n    The rank is also not difficult to set; in all our experiments, we found $r = 20$ works well in all our experiments.\n    This is consistent with works such as [3], which used a rank of $r = 50$ in all there experiments.\n    We also note that [4] which proposed the Nyström preconditioner also comes with an adaptive rank selection algorithm that provably selects the optimal rank.\n    Since we found it easy to select the rank, we did not use this algorithm.\n    In the revision, we will include a more thorough discussion about how to select $r$. \n    We will also mention the adaptive algorithm of [4] to select the rank.\n\n**References**\n\n[1] O’donoghue et al., Conic optimization via operator splitting and homogeneous self-dual embedding (2016)\n\n[2] Schubiger et al., GPU acceleration of ADMM for large-scale quadratic programming (2020)\n\n[3] Zhao et al., NysADMM: faster composite convex optimization via low-rank approximation (2022)\n\n[4] Frangella et al., Randomized nyström preconditioning (2023)", "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:
We would like to sincerely thank all the reviewers for their thoughtful feedback and suggestions, which will greatly improve the submission. 

In addition to our point-by-point response to each reviewer, we wish to underscore our contributions, describe the content of the rebuttal pdf, and address two key concerns shared by several reviewers.

**Contributions**

1. CRONOS is the first algorithm where convex neural networks have been successfully applied to large datasets with competitive performance.

2. CRONOS has strong theoretical convergence guarantees; we provide robust theoretical complexity analysis and convergence proof for the algorithm. This aims to bridge the gap between theory and practice in deep learning.

3. CRONOS-AM is the first algorithm to apply the convex reformulation to deeper networks and achieves excellent performance relative to standard optimizers. Our novel JAX implementation also provides the groundwork for efficient future experiments with fully leveraged GPU acceleration. 

Ultimately our goal is to take a meaningful step towards an innovative and alternative paradigm in deep learning through convex networks. Although much further work remains to be done, we believe this paper opens up exciting new avenues in research with potential for real-world applications. We are highly enthusiastic about exploring the capability of convex networks in enhancing the efficiency and interpretability of deep learning in future work.

**Contents of Rebuttal pdf** 
1. Figures (a) and (b) show the median performance of optimizers across a grid of learning rates with 95\% quantiles for CIFAR-10 and ImageNet. Please note that these plots now contain DAdapted-Adam (DAdam), and competitors no longer do poorly 
on ImageNet due to the finer learning rate grid search. Tables 1 and 2 present the range in peak validation across the grid (except for CRONOS-AM and DAdam, which do not require a learning rate parameter).
CRONOS-AM outperforms DAdam on both tasks and performs comparably to the best-tuned first-order optimizer. 
On ImageNet, Shampoo does best by a decent margin. We attribute this to Shampoo being an approximate second-order optimizer. Properly tuned, Shampoo may yield better performance than purely first-order optimizers like CRONOS-AM and Adam for certain tasks. The plots and the tables show competing methods exhibit a high degree of variance in the learning rate—poor selection can yield non-convergent behavior. 
In contrast, CRONOS-AM does not exhibit these weaknesses and performs comparably to the best-tuned competitor.

2. Figures (c)-(f) compare CRONOS to AdamW on two GPT2 configurations. The first is (GPT2-NFT) which extracts embeddings from pre-trained GPT2 with no fine-tuning. The second configuration GPT2-FT extracts the embeddings after 1 epoch of finetuning with AdamW. Figures (c) and (e) compare CRONOS to tuned AdamW on time. In both cases, CRONOS reaches the best validation faster than AdamW.
Tuned AdamW fails to catch up for both configurations, with a particularly large gap ($\approx 4\%$) on GPT2-NFT.
Figures (d) and (f) plot several AdamW trajectories along with CRONOS against a number of epochs for AdamW and ADMM iterations for CRONOS. The more translucent the curve for AdamW indicates larger deviation from median trajectory. Both plots show AdamW is extremely sensitive to the learning rate selection. Tables 3 and 4 give the same information as Tables 1 and 2 for the GPT2 models. 
Both support CRONOS's superior performance in terms of robustness to hyperparameter tuning, run time, and peak accuracy.   

**Response to General Concerns**

1.  Several reviewers were concerned that in the original submission, most of the competing baselines performed poorly in the ImageNet experiment. This was due to the grid we searched over for the learning rate being inadequate. The original grid searched over learning rates ranging from $10^{-3.5}$ to $10^{-1.5}$. Notably, this grid includes the widely used default of $10^{-3}$ for methods like Adam, AdamW, and SGD employed in frameworks like PyTorch. Unfortunately, learning rates near the default perform poorly in this example, so a deeper search was required. Ultimately, we feel this underscores how difficult it can be to get the right learning rate.
    Our finer search ultimately resolves this issue - the competing methods no longer do poorly. 
    CRONOS-AM still performs favorably relative to the competition. 
    Please see the rebuttal pdf for more details and the new resulting plots, thank you. 

2. There was a lack of a clear baseline for the original GPT2 experiments. We conducted 14 NLP experiments grouped into 3 categories (no train GPT2-NFT, 1 epoch fine tuning GPT2-FT, and unsupervised GPT2-DA). In each experiment, we benchmark against varying seed and learning rate sweeps of AdamW as a baseline. Please see the rebuttal pdf Figures (c) - (f) for a sample of these performance plots. 
We select AdamW as our baseline since it is unanimously utilized in most NLP tasks. 
Our results show that CRONOS outperforms tuned AdamW in both settings in terms of speed and accuracy. We will include all such results for the GPT2 experiments in the revision, along with a table of run-time comparisons. 

Thank you once again to the reviewers for your insightful questions and valuable feedback. We look forward to engaging with you during the discussion phase!

Meta-review: This paper introduces a novel way to train two-layer neural networks by a reformulation to a convex optimization problem. The CRONOS and CRONOS-AM algorithms are developped accordingly to solve the reformulated problem. The authors establish convergence guarantees and complexity analyses for the proposed algorithms. Experiment results on large datasets of image and language tasks to verify the effectiveness of the algorithm. Although the reviewers point out weaknesses of the paper such as insufficient explanations on the implementation in the experiments, all reviewers support accepting the paper after the rebuttal. The authors should add clarifications and discussions in the revised version of the paper according to the reviews and author-reviewer discussions.
Decision: Accept (poster)