CRONOS: Enhancing Deep Learning with Scalable GPU Accelerated Convex Neural Networks

Authors: miria  feng
zachary  frangella
mert  pilanci
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 08:11:22

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.

Main paper: 
 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 as stochastic gradient descent (SGD), offer no guarantees of producing more than an approximate stationary point Arjevani et al. [2023] , which may be considerably suboptimal Ge et al. [2015] . The effectiveness of methods such as SGD and Adam are therefore reliant on heavily conditioned training environments. As a result, most models require extensive hyperparameter tuning of the optimizer to train successfully. This leads to expensive iterations in high compute settings with variable performance depending on optimizer selection and problem domain Yao et al. [2021] . Scaling laws Rosenfeld [2021] 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 of 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 to develop algorithms which converge to the global minimum in polynomial time. However, the resulting convex program is a constrained high-dimensional linear model that is intensely difficult to solve at scale. Other more recent works of Mishkin et al. [2022] and Bai et al. [2023] have made progress in solving this problem with small downsampled versions of MNIST Noever and Noever [2021] and CIFAR-10 Krizhevsky and Hinton [2010] , yet were unable to scale to 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 a fast, efficient and nearly hyperparameter-free method for training twolayer 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 to 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 Boyd 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 a 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 as ImageNet Recht et al. [2019] and large language modeling tasks with GPT-2 Budzianowski 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 et al. [2011] to train two-layer ReLU neural networks with global convergence guarantees.• Using this as a primitive, we extend our algorithm to effectively train multi-layer networks of arbitrary architecture by combining CRONOS with alternating minimization.• We demonstrate the efficient practical applications of our method on real-world large scale image and language datasets, including ImageNet and IMDb.• Our experiments are implemented in JAX -a functional programming paradigm for significant GPU acceleration. Results demonstrate performance speedups and successfully overcome the problem of memory bottlenecks to tackle large data.• Our theoretical analysis proves the convergence guarantees to global minimum under mild assumptions of our proposed algorithms. 
related work: 
 CRONOS builds upon key ideas from previous literature on convex neural networks, ADMM, randomized numerical linear algebra and synergy with JAX implementation. Please see Appendix A for detailed discussions of related work according to each of these three areas of interest.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: Yes Justification: [Yes] , we have a paragraph discussing limitations at the end.Guidelines:• The answer NA means that the paper has no 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 (e.g., independence assumptions, noiseless settings, model well-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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairness. • 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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
background: 
 This section introduces the convex reformulation of two-layer neural networks. We formalize this in two steps: 3.1) foundational definition of convex ReLU neural networks and 3.2) its equivalence as a linearly constrained Generalized Linear Model (GLM) with group lasso regularization.Given a dataset X ∈ R n×d , a two-layer ReLU multilayer perceptron (ReLU-MLP) with weights W (1) ∈ R m×d , w (2) ∈ R m outputs the prediction:EQUATIONHere (x) + = max{x, 0} denotes the ReLU activation function.Given targets y ∈ R n , the network in (1) is typically trained by minimizing the following non-convex loss function:EQUATIONwhere ℓ : R n → R is the loss function, and β ≥ 0 is the regularization strength. It is typically challenging to solve (2), 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 to obtain good performance. This dramatically contrasts with the convex optimization framework, where algorithms come with strong convergence guarantees and involve minimal hyperparameters. Fortunately, it is 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 ≥ m * , for some m ≤ n + 1. Therefore, reformulating (2) as a convex program does not result in loss of information or generality.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 as separating hyperplanes, which essentially multiply the rows of 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 asD X = D = diag (1(Xv ≥ 0)) : v ∈ R d .The cardinality of D X grows as |D X | = O (r(n/r) r ), where r := rank(X) Pilanci and Ergen [2020] . Given D i ∈ D X , the set of vectors v for which (Xv) + = D i Xv, is given by the following convex cone:K i = {v ∈ R d : (2D i -I)Xv ≥ 0}.Learning directly based on the enumeration of D X is impractical due to the exponential size of D X Mishkin et al. [2022] . Instead, we are motivated to work with the following convex program, based on sampling P activation patterns from D X :EQUATIONAlthough (3) only works with a subsampled version of the convex reformulation of Pilanci and Ergen [2020] , it can be shown under reasonable conditions that (3) still has the same optimal solution as (2) Mishkin et al. [2022] . Therefore we can confidently work with the tractable convex program in (3).Prior derivations Bai et al. [2023] have shown that (3) may be reformulated as a linearly constrained composite convex program: Proposition 3.1. Define the matricesF i = D i X and G i = (2D i -I)X, where i ∈ [P ].Then by introducing the constraints u i = v i , z i = w i , where i ∈ [P ], and appropriate slack variables s 1 , . . . , s P , t 1 , . . . , t P , (3) can be reformulated as:min (u,v,s) ℓ(F u, y) + β||v|| 2,1 + 1(s ≥ 0) s.t. I 2dP G u - v s = 0. (4) where u, v, s ∈ R 2dP , F ∈ R n×2dP , G ∈ R 2nP ×2dPThe reformulation in ( 4) is essentially a very large constrained generalized linear model (GLM) with a group lasso penalty Yuan and Lin [2006] . The data matrix F associated with (4) has dimensions n × 2dP , and the constraint matrix has dimensions 2(n + d)P × 2dP . Although the sizes of F and the constraint matrix seem intractable, they are highly structured. Each F i and G i consists of the data matrix X multiplied by a diagonal matrix of 0's and 1'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 i and G i then also inherit this approximate low-rank structure in (4), which CRONOS exploits and solves via fast matrix-vector products on GPU acceleration.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: We have provided the code for our experiments in a zip file.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
cronos: convex neural networks via operator splitting: 
 This section introduces the CRONOS algorithm for solving (4). CRONOS is a scalable, convex optimization-based learning method capable of handling large data and utilizes GPU acceleration for enhanced performance.Algorithm 1 ADMM for Convex ReLU NetworksRequire: penalty parameter ρ repeat 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 al. [2011] . 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 (4) requires solving the u-subproblem at each iteration:u k+1 = argmin u 1 2 ∥F u -y∥ 2 + ρ 2 ∥u -v k + λ k ∥ 2 2 + ρ 2 ∥Gu -s k + ν k ∥ 2 v k+1 s k+1 = argmin v,s β∥v∥ 2,1 + 1(s ≥ 0) + ρ 2 ∥u k+1 -v + λ k ∥ 2 ▷ Primal update λ k+1 ← λ k + γα ρ (u k+1 -v k+1 ) ▷ Dual λ update ν k+1 ← ν k + γα ρ (Gu k+1 -s k+1 ) ▷ Dual ν update until convergence Equation (EQUATIONThis solution then requires solving the following linear system:EQUATIONb k = 1 ρ F T y + v k -λ k + G T s k -ν k .Since F is a n × 2dP matrix and G is a 2nP × 2dP matrix, the cost of solving (6) via a direct method is O(nd 2 P 4 + d 3 P 3 ). This is prohibitively expensive since n and dP are typically large. A natural method is to solve (6) inexactly via the Conjugate Gradient (CG) algorithm, which only requires matrix-vector products (matvecs) with F, F T , G and G T . Under these conditions ADMM will still converge, since the sequences of subproblem errors are summable Eckstein and Bertsekas [1992] .However the number of iterations required by CG to achieve δ-accuracy is O(κ(H + I) log (1/δ)), where κ(H + I) is the condition number of H + I. This results in slow convergence for machine learning problems, since F and G are closely related to the data matrix X, which is approximately low-rank Udell and Townsend [2019] . Given that the linear system in (6) must be solved at each iteration, CG's slow convergence renders it infeasible for real-world application.Algorithm 2 CRONOS Require: penalty parameter ρ, forcing sequence{δ k } ∞ k=1 , rank parameter r [U, Λ] = RandNyströmApprox(F T F + G T G, r)▷ Compute using Algorithm 3 repeat Use Nyström PCG (Algorithm 4) to find u k+1 that solves (5) within toleranceδ k v k+1 ← prox β ρ ∥•∥2 (u k+1 + λ k ) ▷ Primal v update s k+1 ← (Gu k+1 + ν) + ▷ Slack s update λ k+1 ← λ k + γα ρ (u k+1 -v k+1 ) ▷ Dual λ update ν k+1 ← ν k + γα ρ (Gu k+1 -s k+1 ) ▷ Dual ν update until convergenceWe exploit the approximate low-rank structure in our problem matrices to 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 v k+1 and s k+1 has a closed-form solution that may be computed in O ((n + d)P ) time. NysADMM employs the Nyström Preconditioned Conjugate Gradient (NysPCG) from Frangella et al. [2023a] to 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 is constructed with the appropriate rank, it can be shown that Nyström PCG solves the linear system in (6) to δ-accuracy within O(log(1/δ)) iterations (Proposition 6.3). The dependence on the condition number is therefore eliminated, and (5) can be solved quickly at each iteration. Details of the NysPCG algorithm and construction of the preconditioner are presented in Appendix B.1.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.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. [2018] , a high-performance numerical computing library designed to accelerate machine learning research. This framework provides an efficient way to perform array operations, automatic differentiation, and optimization of numerical processes. Leveraging Just-In-Time (JIT) compilation capabilities through XLA (Accelerated Linear Algebra) allow us 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 to fully utilize parallel processing capabilities. The combination JAX and JIT compilation effectively enables us to scale to high-dimensional datasets such as Food (267 × 267 × 3), ImageNet (512 × 512 × 3), and the IMDb language dataset. Our experiments are summarized in Section 7. 
cronos-am: deep networks via alternating minimization: 
 In this section, we introduce the CRONO-AM algorithm for training a 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:minimize θ ∥F(θ; X) -y∥ 2 + α 2 L-2 i=1 ∥W i ∥ 2 F + β 2 ∥W L-1 ∥ 2 F + ∥w L ∥ 2 ,where θ = (vec(W 1 ), vec(W 2 ), . . . , w L ) ∈ R p is the vector of the network weights, with W i denoting the weights for the ith layer, while X is the data matrix, and y are the labels. We can rewrite this objective as: minimizeθ 1:L-2 ,W L-1 ,w L F 1:L-2 (θ 1:L-2 , X)W T L-1 + w L -y 2 + α 2 L-2 i=1 ∥W i ∥ 2 F + β 2 ∥W L-1 ∥ 2 F + ∥w L ∥ 2 ,where θ 1:L-2 = (vec(W 1 ), vec(W 2 ), . . . vec(W L-2 )), F 1:L-1 consists of the first L -2 layers. Let us write X(θ 1:L-2 ) = F 1 (θ 1:L-2 , X), which may be viewed as a transformed data matrix. We shall sometimes write X for brevity. Hence, we can write the output of the network as:F(w; X) = X(θ 1:L-2 )W T L-1 + w L .Thus, the training problem may be rewritten as:minimizeθ 1:L-2 ,w L-1 ,w L X(w 1:L-2 )W T L-1 + w L -y 2 + α 2 L-2 i=1 ∥W i ∥ 2 F + β 2 ∥w L-1 ∥ 2 + ∥w L ∥ 2 .When θ 1:L-2 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: minimizeEQUATIONEquation ( 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 w 1:L-2 and (u, v, s). To handle the convex minimization, we can apply CRONOS. For the nonconvex portion, we propose utilizing DAdapted-Adam [Defazio and Mishchenko, 2023] , as it does not require setting the learning rate. we call this algorithm CRONOS-AM (CRONOS-Alternating Minimization). Pseudocode for CRONOS-AM is given in Algorithm 5. 
theoretical analysis of cronos: 
 In this section, we establish the rate of convergence for CRONOS in solving (4), and provide an overall computational complexity estimate. Additionally, we show that the linear system for the u k -update may be solved by PCG at a rate independent of the condition number.6.1 F ′ i s and G ′ i s inherit approximate low-rank structure of XWe begin by showing the spectrum of X T X upper bounds the spectrum ofF T i F i and G T i G i Proposition 6.1 (F i and G i are approximately low-rank if X is). Let i ∈ [P ], and j ∈ [d]. Then the eigenvalues of F T i F i and G T i G i satisfy: max{λ j (F T i F i ), λ j (G T i G i )} ≤ λ j (X T X).Proposition 6.1 shows that if X is an approximately low-rank matrix, so are all the F i 's and G i 's. Most data matrices in machine learning exhibit fast polynomial or exponential spectral decay and so are well-approximated 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 ′ i s and the G ′ i s, it will also be approximately low-rank, which motivates the following assumption. Assumption 6.2 (Approximate low-rank structure). Let λ j (H) denote the jth eigenvalue of H. Then the eigenvalues values of H satisfy:λ j (H) = O j -2β , β > 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.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 = O log(1 ζ ) . Then after t = O log( 1 δ ) iterations, Nyström PCG outputs a point u k+1 satisfying∥(H + I)u k+1 -b k ∥ ≤ δ,with probability at least 1 -ζ.For fixed ζ > 0, Proposition 6.3 shows that with a rank of r = O(1), NysPCG solves the usubproblem within O log( 1 δ k ) iterations. Thus, PCG solves the linear system quickly. Proposition 6.3 is consistent with practice, NysPCG with rank of r = 20 and 10-40 PCG iterations work well across all problem instances. Therefore solving (5) is not a bottleneck for CRONOS.The following theorem shows CRONOS converges to the global minimum of (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 r = O log( 1ζ ) , and at the kth CRONOS iteration run Nyström PCG with tolerance δ k to convergence. Then, with probability at least 1 -ζ the following statements hold:1. After K iterations, CRONOS' output satisfies ℓ(F ūK , y) + β∥v K ∥ 2,1 + 1(s K ≥ 0) -p ⋆ = O(1/K), I 2dP G ūK - vK sK = O(1/K).2. The total complexity of CRONOS to produce an ϵ-suboptimal point of (4) isC CRONOS = Õ ndP 2 ϵTheorem 6.4 shows CRONOS converges ergodically at an O(1/K)-rate. When the minimum of ( 4) and (2) coincide 1 , Theorem 6.4 guarantees CRONOS is within ϵ of the global minimum of (2) after O(1/ϵ) iterations. By comparison, in the worst-case SGD on (2) can only be guaranteed to output an ϵ-approximate stationary point after O(1/ϵ 4 ) iterations Arjevani et al. [2023] . SGD's output may also fail to be close to a local minimum, and can be highly suboptimal Ge et al. [2015] . Therefore, CRONOS offers much stronger convergence guarantees than SGD.Theorem 6.4 is 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 ϵ > 0, then the cost of the state-of-the-art ADMM approach from Bai et al. [2023] , is at least O(nd 2 P 4 + d 3 P 3 ) since it solves the linear system derived from Eq. ( 5) exactly. In contrast, the dependence of CRONOS upon n, d, and P offers a significant improvement, since the cost grows linearly with n and d. This enables our scalability to very high-dimensional datasets in both vision and language.Theorem 6.4 as presented is pessimistic about the overall convergence speed of CRONOS. For sparsity-promoting convex regularizers such as the group lasso penalty, ADMM is known to reach the manifold containing the support in finite time, after which it converges linearly Liang et al.[2017], Yuan et al. [2020] . Therefore, CRONOS convergence is practically much faster than what Theorem 6.4 suggests. Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: We specify that the hyperparameters for the competing methods are selected via random search, and what grid they are selected from. We provide the values of the hyperparameters used for our method in the appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiments: 
 In this section, we empirically evaluate the efficacy of CRONOS and CRONOS-AM on classification tasks. We first experiment with vision data in multi-layer perceptron (MLP) and convolutional neural network (CNN) architectures 2 . We find that without the necessity of tuning hyperparameters, 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 F presents detailed discussion of the experimental setup and additional numerical results. Details on each dataset used can be found in Appendix F.3.We evaluate the performance of CRONOS-AM for binary classification with a 4-layer multilayer perception whose architecture is provided in the Appendix F.5. We train the model on three different datasets CIFAR-10, Food, and ImageNet 3 . CRONOS-AM is benchmarked against several of the most popular optimizers in deep learning: SGD with Polyak momentum (SGD) [Sebbouh et al., 2021] , Adam [Kingma and Ba, 2014] , AdamW [Loshchilov and Hutter, 2017] , Shampoo [Gupta et al., 2018] , Yogi [Zaheer et al., 2018] and D-adapted Adam (DAdam) [Defazio and Mishchenko, 2023] .For each competing method we consider 5 learning rates, selected randomly from a logarithmic grid with range [10 -5.5 , 10 -1.5 ].Fig. 1 plots the median trajectory across different learning rates of each competing method along with the 5 th and 95 th quantiles. Fig. 1 shows CRONOS-AM either achieves the best or 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 presents 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 fair 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. Our experiments explore sentiment classification with the IMDb dataset. For all experiments, we use the pretrained GPT2 architecture with 12 transformer blocks and an 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 9 batches of size (1500 x 60 x 768) for each of the 2 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. 2 compares CRONOS to tuned AdamW 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 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. 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 3 shows CRONOS significantly outperforms tuned AdamW. It reaches higher validation accuracy faster 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. 2c 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 AdamW, 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 Appendix F, and show promising directions for future work.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: We report quantiles for methods across different learning rates. For the deepnetwork experiments, we ran each method for the 3 different random seeds and found the conclusions to be the same.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• 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. 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 a 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
conclusion: 
 We introduce CRONOS, the first algorithm to successfully apply convex neural networks to highdimensional 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: a novel algorithm which utilizes 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.19% 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.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 NLP 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. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: [NA] Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: Yes we have used JAX and FLAX and cited them in our bibliography, as well as during citations the main 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 or dataset.• The authors should state which version of the asset is used and, if possible, include a 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. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA]Justification: [NA] 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 is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
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: [NA] 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
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: [NA] Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on 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. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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. 
 
 Appendix: 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 nonconvex optimization. Bengio et al. [2005] is an early work in this direction; they treat 2-layer neural networks as convex models but require the first layer to be fixed. Random features regression is a similar approach, where the first layer is randomly initialized, and only the last layer is trained Rahimi and Recht [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 to standard neural networks Chizat et al. [2019] , Ghorbani et al. [2021] . This makes them less attractive relative to the convex reformulation of (3). Frangella et 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 in 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.In this section, we give the algorithms from Frangella et al. [2023a] for the randomized Nyström approximation and Nyström PCG. Algorithm 3 constructs a low-rank approximation Ĥ = U ΛU T to the input matrix H. Ĥ is then used to form the Nyström preconditioner P , which (along with its inverse) is given by:P = 1 λr + µ U ( Λ + µI)U T + (I -U U T ), P -1 = ( λr + µ)U ( Λ + µI) -1 U T + (I -U U T ).Require:psd matrix H ∈ S + d (R), sketch size s Ω = randn(d, s) ▷ Gaussian test matrix Ω = qr(Ω, 0) ▷ thin QR decomposition Y = HΩ ▷ s matvecs with H ν = eps(norm(Y, 2)) ▷ compute shift Y ν = Y + νΩ ▷ add shift for stability C = chol(Ω T Y ν ) ▷ Cholesky decomposition B = Y ν /C ▷ triangular solve [U, Σ, ∼] = svd(B, 0) ▷ thin SVD Λ = max{0, Σ 2 -νI} ▷ remove shift, compute eigs return Nyström approximation Ĥ = U ΛU TIn CRONOS, we set µ = 1, 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 3, is that if the time to compute a matrix-vector product with H is T mv , the asympotic cost is O(T mv r) Zhao et al. [2022] . The preconditioner is then deployed with Nyström PCG (Algorithm 4) at each iteration to efficiently solve (5).Require: psd matrix H, righthand side r, initial guess x 0 , regularization parameter ρ, sketch size s, tolerance ε[U, Λ] = RandomizedNyströmApproximation(H, s) w 0 = r -(H + ρI)x 0 y 0 = P -1 w 0 p 0 = y 0 while ∥w∥ 2 > ε do v = (H + ρI)p 0 α = (w T 0 y 0 )/(p T 0 v) x = x 0 + αp 0 w = w 0 -αv y = P -1 w β = (w T y)/(w T 0 y 0 ) x 0 ← x, w 0 ← w, p 0 ← y + βp 0 , y 0 ← y end while return approximate solution x B.1.1 Setting hyperparameters in Nyström PCGNyström PCG has two parameters: the rank r and the tolerance ε. The tolerance is easy to set -any summable sequence suffices to ensure CRONOS will converge. Thus, we recommend setting the PCG tolerance to ε k = k -1.2 at each iteration. This is consistent with popular ADMM solvers such as SCS, which uses the same sequence in its CG solver [O'donoghue et al., 2016] .The rank is also not difficult to set. In our experiments, a rank of r = 20 worked well uniformly. So, we recommend this as a default value. However, if 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 r is on the order of the effective dimension with high probability. This guarantees that Nyström PCG will converge to tolerance ε in O (log(1/ε)) iterations (see Appendix C).We present pseudo-code for CRONOS-AM in Algorithm 5. For DAdapted-Adam we run the algorithm for 1-epoch to perform the approximate minimization.Require: penalty parameter ρ, forcing sequence {δ k } ∞ k=1 , rank parameter r repeatApproximately minimize (7) with respect to (u, v, s) using CRONOS to obtain (u k+1 , v k+1 , s k+1 ) ▷ Run CRONOS for 5 iterations Approximately minimize (7) with respect to w 1:L-2 using DAdapted-Adam to obtain w k+1 1:L-2 . ▷ Run DAdapted-Adam for 1 epoch until convergenceC.1 Proof of Proposition 6.1Proof. Recall by definition, that F i = D i X i and G i = (2D i -I)X, where D i ∈ R n×n is a diagonal matrix where D ii ∈ {0, 1}. Consequently,D 2i ⪯ I and (2D i -I) 2 ⪯ I. Thus, by the conjugation rule, we conclude:X T D 2 i X ⪯ X T X and X T (2D i -I) 2 X ⪯ X T X. Recalling A ⪯ B implies λ j (A) ≤ λ j (B), and that F T i F i = X T D 2 i X and G T i G i = X T (2D i - I) 2 X, the desired claim follows from the preceding display.The key quantity in establishing fast convergence of Nyström PCG is to select the rank to be on the level of the effective dimension of the matrix defining the linear system Frangella et al. [2023a] . Given µ > 0 and a symmetric positive definite matrix H, the effective dimension is given by:d µ eff (H) = trace H(H + µI) -1 . Roughly speaking, d µ eff (H)gives a smoothed count of the eigenvalues of H larger than the parameter µ > 0. We now recall the following result which bounds d µ eff (H) when the eigenvalues decay polynomially. Lemma C.1 (Effective dimension under polynomial decay.). Let µ > 0, β > 1/2 and let H ∈ R N ×N be a symmetric positive semidefinite matrix. Suppose the eigenvalues of H satisfy:λ j (H) = O(j -2β ), ∀j ∈ [P ]. Then, d µ eff (H) ≤ Cµ -1 2β, where C > 0 is some constant.For a proof, see Section C.1 of Bach [2013] . With these preliminaries out of the way, we now commence with the proof of Proposition 6.3.Proof. Let Ĥ denote the randomized Nyström approximation of H output by Algorithm 3 and define E = H -Ĥ. By Lemma C.1 with µ = 1 and Assumption 6.2, we have thatd 1 eff (H) = O(1). As r = O log 1 ζ = O d 1 eff (H) + log 1 ζ, it follows from Lemma A.7 of Zhao et al. [2022] , thatP (∥E∥ ≤ 2) ≥ 1 -ζ.We also obtain from Lemma 5.4 of Frangella et al. [2023a] that λ r ( Ĥ) ≤ 1. Thus, applying Proposition 4.5 of Frangella et al. [2023a] gives:κ(P -1/2 (H + I)P -1/2 ) ≤ 1 + λ r ( Ĥ) + ∥E∥ ≤ 4.Moreover, from the convergence analysis for CG, we know that Golub and Van Loan [2013] ∥u k+1 -ũk+1 ∥ H+I ≤ 2 κ(P -1/2 (H + I)P -1/2 ) -1EQUATIONNow,∥(H + I)u k+1 -b∥ = ∥(H + I)(u k+1 -ũk+1 )∥ ≤ λ 1 (H + I)∥u k+1 -ũk+1 ∥ H+I ,Hence, by the last display with (8), we reach∥(H + I)u k+1 -b∥ ≤ λ 1 (H + I)∥ũ k+1 ∥ H+I (1/3) t ,The desired claim immediately follows from the last display.C.3 Proof of Theorem 6.4CRONOS convergence is a consequence of the following theorem, which is a simplified version of Theorem 1 of Frangella et al. [2023b] specialized to the setting of Algorithm 2. Theorem C.2 (Simplified Theorem 1, Frangella et al. [2023b] ). Let p ⋆ denote the optimum of (4).For each K ≥ 1, denote ūK+1 = 1 K t+1 k=2 u k , vK+1 = 1 K K+1 k=2 v k , and sK+1 = 1 K K+1k=2 s k ,where {x k } k≥1 and {z k } k≥1 are the iterates produced by Algorithm 2 with forcing sequence {δ k } k≥1 with u 1 = 0, v 1 = 0, s 1 = 0, λ 1 = 0, and ν 1 = 0. Let ũk+1 denote the exact solution of (5) at iteration k. Suppose that ∥u k+1 -ũk+1 ∥ ≤ δ k for all k ∈ [K]. Then, the suboptimality gap satisfiesℓ(F ūK+1 , y) + β∥v K+1 ∥2,1 + 1(s K+1 ≥ 0) -p ⋆ = O(1/K).Furthermore, the feasibility gap satisfiesI 2dP G ūK+1 - vK+1 sK+1 = O 1 K .Consequently, after O(1/ϵ) iterations, ℓ(F ūK+1 , y) + β∥v K+1 ∥2,1 + 1(s K+1 ≥ 0) -p ⋆ ≤ ϵ, andI 2dP G ūK+1 - vK+1 sK+1 ≤ ϵ.Proof. From Theorem C.2, it suffices to show that for all k ∈ [K] that∥u k+1 -ũk+1 ∥ ≤ δ k .To this end, observe as r = O log 1 ζ and Nyström PCG at the kth iteration runs for O log 1 δ k , it follows from Proposition 6.3, that with probability 1-ζ for all k ∈ [K] the output u k+1 of Nyström PCG satisfies:∥(H + I)u k+1 -b k ∥ ≤ δ k , ∀k ∈ [K]. Consequently, with probability at least 1 -ζ: ∥u k+1 -ũk+1 ∥ ≤ ∥u k+1 -ũk+1 ∥ H+I = ∥(H + I)(u k+1 -ũk+1 )∥ (H+I) -1 = ∥(H + I)u k+1 -b k ∥ (H+I) -1 ≤ ∥(H + I)u k+1 -b k ∥ ≤ δ k .Thus, we can invoke Theorem C.2 to conclude the proof of the first statement.We now prove the second statement. Recall that the cost of multiplying a vector by H costs O ndP 2 . 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 r > 0, the cost of forming the Nyström approximation is O(ndP 2 r) (see Appendix B.1). In Theorem 6.4, r = O log 1 ζ , so the total cost of constructing the preconditioner is Õ(ndP 2 ). By Proposition 6.3, with probability at least 1 -ζ, the total cost of solving (5) for K iterations is given by:K j=1 O ndP 2 log 1 δ k ≤ KO ndP 2 log 1 δ K = Õ ndP 2 ϵ ,where the last equality follows as K = O (1/ϵ) . Thus, with probability at least 1 -ζ the total complexity of Algorithm 2 isC CRONOS = Õ ndP 2 ϵ .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. ( 3) coincide with the optimal value of the non-convex problem Eq. ( 2)? Pilanci and Ergen [2020] established that the optimal value of the full convex reformulation agrees with that of Eq. ( 2), but this is computationally intractable. It is not apparent apriori that the optimal value of the subsampled convex program (3) will agree with the optimal value of Eq. ( 2). Mishkin et 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 of Kim and Pilanci [2024] provides a more quantitative bound on the gap between the optimal values of ( 3) and (2). Let p ⋆ MLP-NCVX denote the optimal value of (2), and p ⋆ MLP-CVX denote the optimal value of (3). Theorem 2.1 in Kim and Pilanci [2024] establishes that if m = Ω (log(n)) and P = m/2, then with high probability:p ⋆ MLP-NCVX ≤ p ⋆ MLP-CVX ≤ C log(2n)p ⋆ MLP-NCVX ,where C > 0 is some absolute constant. This result shows that for P = Ω(log(n)), p ⋆ MLP-CVX deviates from p ⋆ MLP-NCVX by a modest logarithmic factor. Moreover, as two-layer ReLU networks are universal approximators, it is often the case that the network is capable of perfectly fitting the training data, in which case p ⋆ MLP-CVX = p ⋆ MLP-NCVX = 0. Combining the preceding result with Theorem 6.4, we see that CRONOS will deliver a solution with comparable objective value to p ⋆ MLP-NCVX in polynomial time. When p ⋆ MLP-NCVX = 0 or the conditions of Mishkin et al. [2022] hold, CRONOS can produce a solution arbitrarily close to p ⋆ MLP-NCVX in polynomial time.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. [2022] applies to losses beyond the least-squares loss. Instead of solving (5), NysADMM replaces ℓ(F u) with a (regularized) second-order Taylor expansion about the current iterate u k , to obtain the subproblem.argmin u {ℓ(F u k ) + (u -u k )F T ∇ℓ(F u k , y) + 1 2 (u -u k ) T (F T ∇ 2 ℓ(F u k )F + σI)(u -u k ) + ρ 2 ∥u -v k + λ k ∥ 2 2 + ρ 2 ∥Gu -s k + ν k ∥ 2 } (9)The solution of (9) may be found by solving the following linear system:(H k + I)u k+1 = b k ,whereH k = 1 ρ F T ∇ 2 ℓ(F u k )F + G T 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 to Algorithm 2, the convergence of Algorithm 6 can be established using results from Frangella et al. [2023b] .Require: penalty parameter ρ, forcing sequence {δ k } ∞ k=1 , rank parameter r, preconditioner update frequencyf repeat if k is a multiple of f then [U, Λ] = RandNysAppx( 1 ρ F T ∇ 2 ℓ(F u k )F + G T G, r) ▷ Update Nyström preconditioner every f iterations end ifUse Nyström PCG to find u k+1 that solves (9) within toleranceδ k v k+1 ← prox β ρ ∥•∥2 (u k+1 + λ k ) ▷ Primal v update s k+1 ← (Gu k+1 + ν) + ▷ Slack s update λ k+1 ← λ k + γα ρ (u k+1 -v k+1 ) ▷ Dual λ update ν k+1 ← ν k + γα ρ (Gu k+1 -s k+1 ) ▷ Dual ν update until convergenceIn this section, we include additional details relevant to our experimental design. We also provide additional numerical results not present within the main paper.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 OS.In all our experiments CRONOS (including when used as a subproblem solver in CRONOS-AM is 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 ρ is varied from 0.001 to 1 depending upon the task. For CRONOS-AM DAdapted-AdamW is always run for 1 epoch to get the non-convex weights.In practice, ρ 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 a common strategy employed in practice by ADMM solvers to set ρ [Stellato et 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 ρ > 0 [Frangella et al., 2023b] . So while a judicious choice of ρ may lead to 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 if the learning rate is not within some small band, non-convergent behavior can arise.Fashion MNIST. The Fashion-MNIST dataset is a collection of 70,000 grayscale images, each of size 28x28 pixels, categorized into 10 classes. It 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 of 2 classes on Cifar10, usually in a downsampled setting. Food-5k Dataset. The Food-5K dataset is a collection of 5,000 images with binary classification label of '1 = food' and '-1 = 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 ImageNet512. 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. It consists of only negative and positive reviews, without neutral. We perform all experiments with a batch size of 32, and pad all input sequences to 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% validation accuracy.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. We select AdamW as our baseline since it 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 2 batches only, in order to evaluate the efficacy of our method on low data regime settings.IMDB-FT. 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.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 in real world scenarios. We 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. This examines 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. We recognize that domain-specific tasks generally perform well when the feature distributions of the domains are similar. Therefore we examine the effectiveness of initializing from the broadly pretrained GPT2 model on the unlabeled IMDb dataset. Experimental results show promising directions for future work.Below in Fig. 3 , we report results for various optimizers on Food-5k.Here we test the performance of CRONOS-AM on a 4-layer CNN consisting of 2 convolutional layers with average pooling, followed by two densely connected layers. The goal of this experiment is to demonstrate the feasibility of applying CRONOS-AM to architectures aside from MLPs. Fig. 4 gives results for CIFAR-10 and ImageNet-171. On CIFAR-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-AM beyond MLPs. For future work, it would be interesting to explore the effectiveness of CRONOS-AM on other popular architectures.To see how CRONOS performs in the multi-class setting, we experiment on the Fashion MNIST dataset [Xiao et al., 2017] . This yields a classification problem with 10 unique classes. We consider a two-layer ReLU MLP with 64 neurons. For comparison we also consider Adam, AdamW, SGD+Momentum, Shampoo, and Yogi. To fully exploit GPU acceleration, the batchsize for these methods is set to 10, 000. As in the main paper, we consider a random grid of five learning rates.The results for the experiment are provided in Table 5 . The table shows CRONOS delivers the best performance by significant margin relative to its non-convex competitors. Moreover, we see the Justification: We provide the theory outlined in the abstract. Our experiments compare our method to popular competing optimizers across a grid of learning rates and show comparable or better performance.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 in 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 as long as it is clear that these goals are not attained by the paper.For a detailed discussion when this occurs, please see Appendix D. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 1.The symbols W1j, w2 are not consistent with their definition in line 77.
2. Figure 2 is not referenced in the paper.
questions: questions: How is unsupervised training implemented on IMDC-DA?
limitations: limitations: The authors should conduct experiments on more datasets to test the effectiveness of the proposed algorithms.
paper_summary: paper_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.

The 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.

The 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.
strength_weakness: strength_weakness: 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.
    * 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$.
    * 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.
    * 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.
    * The selection of $P$ directly influences the computation cost as the authors analyzed in Section 6.4.

However, 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. 

2. Following the first point, the authors did not mention the actual training time comparison of the proposed method and normal deep learning optimization methods.
questions: questions: 1. How were the experiments on Imagenet actually conducted?
    * Why were they binary classification tasks as described in Line 599?
    * 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?

2. 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?

3. 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.

4. In line 91, it should be "provided $m \geq m^*$, for some $m^* \leq n + 1$.
limitations: limitations: N/A
paper_summary: paper_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. 
This 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).

Then 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).
The first line of algorithm 1 requires solving a linear system (6) -- Solving (6) directly is too expensive due to large matrix sizes. 
The 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.

To 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. 

Then 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.

The theory seems to be robust and very interesting.
strength_weakness: strength_weakness: Small concerns:

1. 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. 

2. 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. 

Big concerns:

3. Experiments:  
a. What do the authors mean by binary classification on imagenet and cifar?!  
b. Figure 2 seems to be wrong --  please add also the training accuracy to see if it is simply overfitting or what is happening.
    * Adam is usually better than SGD, I encourage the authors to try the standard common parameters of ADAM.
    * On what dataset is Figure 2 reporting the results? 

c. Figure 1 b is even worse, how come, known optimizers do not perform well on binary classification? and even reach 0 Accuracy

d. NLP:
   * What is the accuracy of GPT before running your algorithm on the same data, without any fine-tuning?
   * 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?

I 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.
questions: questions: * Consider changing the notation of relu it is so confusing specifically that the next character is multiplied and not added. 

* The authors define W^(1) and then use W_1 ?

* I am missing the connection from (2) to (3):
     A. Why sampling P activation patterns is enough? 
     B. Is it simple by sampling?

* Define ||v||2,1 and provide an intuition on (4)

* what are the reasonable conditions that (3) still has the same optimal solution as (2)?
limitations: limitations: See above
paper_summary: paper_summary: - This paper builds on work on reframing the optimization of 2-layer ReLU networks as convex optimization
- The first contribution of this work is the use of the ADMM algorithm for solving the constrained optimization task from equation 4
- 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
- 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
- They show that this method can successfully optimize deep networks, leading to the first successful application of convex neural network application for deep networks
strength_weakness: strength_weakness: - 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
- 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
- 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?)
questions: questions: - Would it be possible to have the results from Figure 1 and 2 in a table format for easier reading?
- 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
- 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.
- Could we see the performance of D-adapted methods in figure 1 and 2?
- What is the observed runtime of the algorithm?
- 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: 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.

Comments: title: Thank you for your reply
Thank you for your reply
Post rebuttal
comment: We thank the reviewer for their reply and for raising their score. We are happy to hear we addressed your concerns.

The table shows the total time each optimizer takes to train the model. The reviewer is correct in noting that the training time is quite fast! The main reason for this is that all optimizers use JAX's just-in-time compilation (JIT) along with GPU acceleration, making all the methods extremely fast. Additionally since these experiments consider binary classification with two classes, the number of training examples isn't massive. Therefore training for 50 epochs happens very fast when using JIT on a 4090 GPU with about 83 TFLOPS. Without Jax's JIT compilation the runtimes would not be this fast.
I appreciate the additional experiments and they have addressed my concerns. I will raise my score to a 5.

Additionally, what is the runtime table measuring? Could the authors just give a table showing the total runtime of the entire training procedure? Unless I am misunderstanding, 3.00s is too fast to train a model, so the numbers in that table do not seem correct.
Thank you for your response and for raising your score. We are happy to hear that our reply addressed your main concerns.
Thank you for the detailed response.
It addressed most of my comments.

I will increase the score as I see the main contributions of this paper are the theoretical ones. However, I still think that the experiments are limited (even if prior work does that).