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

Weaknesses

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

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

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

2. Please see the preceding response.

3. Below is a table showing the runtime of CRONOS and CRONOS-AM to non-convex optimizers. Notably, 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.

Optimizer Runtimes (s) on CIFAR-10 and ImageNet

Questions

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

2. Yes, this is a good point! We plan to add these results to the revision.

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

4. Yes! We have added this comparison. Please see the rebuttal pdf.

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

6. The reviewer raises an excellent point! Nyström PCG does have two hyperparameters. Fortunately, they are very easy to set. For the tolerance, one can use any summable sequence to ensure convergence. In particular, the popular ADMM solver SCS from [1] uses a sequence of $1/k^{1.2}$. 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]. We recommend using the summable tolerance as that is what the theory requires to obtain convergence. Nevertheless, past work has found that both these strategies perform very similarly [3]. We find this to be true in our experience as well. We will include a discussion of how to set the tolerance in the revision.

   The rank is also not difficult to set; in all our experiments, we found $r = 20$ works well in all our experiments. This is consistent with works such as [3], which used a rank of $r = 50$ in all there experiments. 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. Since we found it easy to select the rank, we did not use this algorithm. In the revision, we will include a more thorough discussion about how to select $r$. We will also mention the adaptive algorithm of [4] to select the rank.

References

[1] O’donoghue et al., Conic optimization via operator splitting and homogeneous self-dual embedding (2016)

[2] Schubiger et al., GPU acceleration of ADMM for large-scale quadratic programming (2020)

[3] Zhao et al., NysADMM: faster composite convex optimization via low-rank approximation (2022)

[4] Frangella et al., Randomized nyström preconditioning (2023)

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.

Weaknesses

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

2. Equation (1) does define a 2-layer ReLU networkin the regression setting, see [1, 2] . The more familiar softmax case mentioned by the reviewer can also be covered by the convex reformulation, see [1, 3] for further details.

3a. By binary classification, we mean we look at two classes of CIFAR-10 and ImageNet and perform binary classification. This is consistent with prior literature on convex neural networks such as [3, 4}.

3b. This occurred since the initial learning rate grid we used wasn't fine enough. Our original grid included Adam's default learning rate of $10^{-3}$. However this value doesn't work well in the architecture considered for these two classes of ImageNet. The new plot is included in the rebuttal pdf Figures (a) and (b) with a finer grid. Now, the competing methods don't all fail.

3c. Figure 2 shows the performance of Adam and SGD across different learning rates on CIFAR-10. 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.

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

• Features extracted from pre-trained GPT2 followed by classification immediately with No Fine-Tuning (IMDb-NFT).
• Features extracted then followed by a mean-pool layer before classification.
• Features extracted then followed by an attention-pooled layer followed by classification.

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

Questions

1. Sure! We'll change this in the revision to avoid any confusion.
2. Thanks for catching this! We will fix it in the revision.
3. Eq (2) is the standard non-convex training problem. [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. This is extremely large unless the rank is small, so sampling all the patterns from this set is computationally infeasible.

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

This is a great question! Theorem 2.1 in [3] gives precise conditions under which the optimal values of Eq (1) and Eq (2) coincide. But a simpler, more quantitative explanation is provided by Theorem 2.1 in the recent paper [5]. This theorem shows that as long as $P = \mathcal O\left(\log(n)\right)$, where $n$ is the number of data points, then

Optimal value of Eq 2 $\leq \mathcal O(\sqrt{\log(n)}) \times$ Optimal value of Eq 1.

Evidently 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). In this work, we only used $P = 10$ throughout all experiments, which was sufficient to obtain good results. In the revision, we include a more careful discussion of these points.

As per your question on whether sampling is efficient, the answer is yes! 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)$.

4. Thank you for catching this! We will make sure it is defined in the revision.

   Sure, we can give intuition! The $\mathcal K_i$ in (3) are defined by linear inequality constraints. By adding non-negative slack variables, these can be turned into linear equality constraints. Then, organizing all the variables into the appropriate vectors $\mathbf u, \mathbf v,$ and $\mathbf s$, and compactly writing down the linear equality constraints, one arrives at (4). 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.

5. Please see the response to question 3.

References

[1] Du et al., Gradient Descent Provably Optimizes Over-parameterized Neural Networks (2018)

[2] Pilanci and Ergen, Neural Networks are Convex Regularizers: Exact Polynomial-time Convex Optimization Formulations for Two-layer Networks (2020)

[3] Mishkin et al., Fast Convex Optimization for Two-Layer ReLU Networks: Equivalent Model Classes and Cone Decompositions, (2022)

[4] Bai et al., Efficient global optimization of two-layer relu networks: Quadratic-time algorithms and adversarial training (2023)

[5] Kim and Pilanci, Convex Relaxations of ReLU Neural Networks Approximate Global Optima in Polynomial Time (2024)

Weaknesses

1.) This is an excellent point! It is important to have guidance on how to select the hyperparameter $P$ in practice. Fortunately, large $P$ is not necessary to ensure the optimal values are comparable. 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( \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)$. This is an extremely modest multiple. Therefore from a practical standpoint, only a modest number of neurons is needed to ensure good performance. 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.

In the revision we will be sure to mention this point. 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.

2.) Below is a table showing the runtime of CRONOS and CRONOS-AM to non-convex optimizers. Notably, 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.

Optimizer Runtimes (s) on CIFAR-10 and ImageNet

Questions

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

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

3, 4.) Thank you for catching these! We will fix them in the revision.

References

[1] Kim and Pilanci, Convex Relaxations of ReLU Neural Networks Approximate Global Optima in Polynomial Time (2024)

** Weaknesses **

1. Thank you for pointing this out! We have fixed this in the revision.

Questions

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

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