Paper: Adam through a Second-Order Lens (ICLR-2024)

Reviewer comment (Novelty):
The paper lacks novelty and originality because it only combines Adam and K-FAC in a facile way and thus doesn't provide better results for most of the tasks mentioned, such as UCI Energi/Protein, Fashion-MNIST and CIFAR-10 (in these cases, K-FAC and/or original Adam are better than the proposed method because the generalization performance VS runtime is not competitive as stated in the abstract).

Author response:
Since AdamQLR is closely related to its Adam base, we would not expect AdamQLR to display dramatic performance improvements over Adam. The major contribution of our work is that an _Untuned_ version of AdamQLR performs comparably to a _Tuned_ version of Adam. This indicates a significant cost saving which does not appear on our figures: that of performing hyperparameter optimisation. We will update our paper to make this narrative clearer.

While we appreciate our combination of Adam and K-FAC may seem “facile” at first glance, we emphasise our motivating observation that the heuristics used in K-FAC are _essential_ to its ability to perform stable training, rather than providing slight quality improvements. This is a very different relationship from the use of e.g. weight decay in SGD, despite both being ‘heuristics’. From this perspective, we argue an evaluation of these techniques applied to Adam is of interest and value to the ML community, and extends beyond a trivial augmentation of Adam.

Reviewer comment (Experiments):
The evaluation was performed on small tasks and the usage of Rosenbrock function does not add any value to the paper since the tasks that involve Neural Networks are much more complicated.

Author response:
Even though our experiments in Figure 2 aren’t as large-scale as some other applications, we think they strike a balance between being large enough to display features of the neural network training problem, but small enough to easily perform multiple repetitions of the experiments to reinforce our results. The Rosenbrock function is included largely as a proof-of-concept, but we think it’s valuable to visualise the trajectory followed by AdamQLR as compared to our benchmarks.

Reviewer comment (Presentation):
The paper does not have any tables that contain accuracies for classification tasks, making it extremely difficult to figure out what the final accuracies are only by looking at the plots.

Author response:
We will add tables of final performance figures to the paper, though we note Figure 4 (Appendix B.1.1) plots accuracy evolutions for our experiments on classification tasks.

Reviewer comment (Novelty):
Overall, the paper does not meet the novelty and originality requirements for ICLR.

Author response:
To our knowledge, our submission is the first work to apply heuristics from second-order optimisation to a gradient-based method such as Adam. It addresses the relevant question of how these heuristics might aid first-order algorithms, especially considering their fundamental importance to second-order techniques. Further, it describes how the tuned performance of a widely-used optimiser (Adam) can be replicated with AdamQLR without the cost of hyperparameter optimisation. Recalling the ICLR guidance that “a lack of state-of-the-art results does not by itself constitute grounds for rejection”, we feel our work is sufficiently novel and original to be of interest to the ICLR community. We would welcome further discussion to help us understand the root of your disagreement.

Reviewer comment (Writing):
I wonder how to get the matrix $C$ in Eq. 1.

Author response:
The curvature matrix $\mathbf{C}$ denotes any estimate of objective curvature we may wish to use. AdamQLR only uses this matrix in the context of a product with an arbitrary vector. Using the Jacobian-vector product trick (Pearlmutter, 1994; reference in paper), we can compute such products with the Fisher information matrix (which we use) or the Hessian matrix (a common alternative) without having to compute or store the whole matrix. This allows us to compute the exact matrix-vector products according to the definition of our chosen curvature matrix. The functions `learning_rate_from_hessian()` and `learning_rate_from_fisher()` in `optimisers.py` from our Supplementary Material contain our implementations in JAX. We briefly mention this approach in the final paragraph of Section 3.3, and will expand our explanation to make this clearer.

Reviewer comment (Theory):
What is the principle of setting $\omega_{dec}$ and $\omega_{inc}$, and why $\lambda$ is adjusted when $\rho$ larger than 3/4 or smaller than 1/4?

Author response:
In principle, we seek $\omega_\textrm{dec}$ and $\omega_\textrm{inc}$  to be sufficiently close to 1 to avoid destabilising the optimiser with dramatically-varying damping, but sufficiently far from 1 to allow prompt rectification of any undesired behaviour from the optimiser. We note that a major thrust of our contribution is that an _Untuned_ version of AdamQLR (using the default hyperparameters we suggest) performs comparably to a _Tuned_ version of Adam, which avoids placing any burden of selecting these factors on the end-user – we will update our paper to make this clearer. In any case, our sensitivity study in Figure 13 (Appendix B.2.5) suggests AdamQLR is robust to the setting of these factors.

Our thresholds for $\rho$ in Equation 2 were taken directly from the original K-FAC settings (Martens and Grosse, 2015; reference in paper), since we considered them to be components of the damping strategy we were porting to Adam. Their purpose is to detect when the loss change predicted by the model is much larger than the actual change, and intervene quickly to apply damping to mitigate the model error. The intention is to bias the decision towards decreasing damping if the optimisation is proceeding stably, since in the vicinity of a local minimum excessive damping would slow down convergence. We will add a more verbose discussion of this phenomenon to the paper.

Reviewer comment (Evaluation):
I think the main results are from the figure 2. But the figure is not very clear for me, maybe you can list the training loss, test loss, convergence steps, and generalization gap ( |training_loss - test_loss| ) in a table. From this figure. I'm not very clear whether the proposed method can solve the overfitting issue and improve the generalization. So I think you can analyze the generalization gap.

Author response:
Figure 2 is indeed the crux of our results section, and our claims about generalisability are based on comparisons from these graphs. We will add a table of final results as you suggest.

Reviewer comment (Experiments):
The training task is too simple and the results on complex tasks (such as ImageNet) is not very strong.

Author response:
Even though the tasks in Figure 2 aren’t as large-scale as some other applications, we think they strike a balance between being large enough to display features of the neural network training problem, but small enough to easily perform multiple repetitions of the experiments to reinforce our results. The reviewer may also be interested in the Penn Treebank/GPT-2 results we will be adding in response to other feedback.

Reviewer comment (Evaluation):
Loss value in figure 2 is not great enough to compare the performance of different methods and maybe you should provide the accuracy value.

Author response:
We show classification accuracies in Figure 4 (Appendix B.1.1), since we agree there is reason to be interested in both loss and accuracy performance. We will update the main body to point out this figure more clearly. Regarding the performance similarities: since we apply second-order heuristics to the core Adam algorithm, we don’t expect to massively beat it – our main contribution is that the _Untuned_ version of AdamQLR performs comparably to the _Tuned_ version of Adam, meaning AdamQLR (Untuned) can be applied with dramatically reduced computational effort.