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

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

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 - The theory is robust and well supported - 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 - 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

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

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

2

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.

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.

Strengths

1. The paper has many theoretical novelties and seems interesting from this perspective. 2. The used tricks and formulation can open new directions for research. 3. 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: 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

* 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)?

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

See above

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.

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

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

Rating

5

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

N/A

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.

Strengths

1. Two neural network optimization methods, CRONOS and CRONOS-AM, are proposed. 2. The authors provide theoretical complexity analysis and convergence proof of the algorithms. 3. 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. 4. 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. 2. Figure 2 is not referenced in the paper.

Questions

How is unsupervised training implemented on IMDC-DA?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors should conduct experiments on more datasets to test the effectiveness of the proposed algorithms.

Post rebuttal

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