Paper: Jorge: Approximate Preconditioning for GPU-Efficient Second-Order Optimization (ICLR-2024)

Reviewer comment (Theory):
The equality on line -2 on page 4 is not valid: (A + B)^r =/= A^r (I + A^{-1} B)^r, unless A and B commute. Is there a way to fix the equation?

Author response:
Thank you for bringing this to our attention. You are correct, and we appreciate your observation. The equality (A + B)^r =/= A^r (I + A^{-1} B)^r is valid only when A and B commute. We acknowledge this limitation and will make the necessary clarification.

Reviewer comment (Experiments):
The authors' main claim is that the Taylor approximation is sufficient to compute the preconditioners — do you have some direct evidence of this (i.e. some illustration of the difference between the original preconditioners and the Jorge approximations)?

Author response:
While Jorge computes an approximation of Shampoo's preconditioner calculation, our experimental results, as illustrated in Figure 3 (left), demonstrate that it matches the sample efficiency of Shampoo. This provides empirical evidence supporting the sufficiency of our approximation in computing the preconditioners. We acknowledge the suggestion for a more fine-grained analysis, such as a direct comparison between the approximate and exact preconditioners, and we plan to incorporate this in our final manuscript to further strengthen this claim.

Reviewer comment (Experiments):
In most Shampoo implementations, the matrix roots are also computed via an iterative process that only involves matrix multiplications. How does your approximation compare with this process, for example by limiting the number of iterations to a fixed value?

Author response:
The implementation of Shampoo we compared with [1] includes support for an iterative scheme known as the Coupled Inverse Newton's iteration. We acknowledge the importance of comparing Jorge's iteration times with this scheme, and we plan to incorporate this comparison in our final manuscript for a more comprehensive evaluation.

Reviewer comment (Experiments):
You say that (1 - b)/b X_L is chosen adaptively to ensure convergence of the Taylor approximation — what values of b do you get as a result? Note that if b is small then effectively the previous gradients will be forgotten very quickly, which may lead to instability (so the b also serves a different purpose than the one you use it for).

Author response:
While a very low value of b can lead to gradients being forgotten quickly, we did not observe any convergence issues in any of our experiments (Figures 2 and 3). We agree that it might be useful to have a hyperparameter that serves as a lower bound on the value that our adaptive scheme can assign to b.

Reviewer comment (Theory):
The authors need to analyze computation and memory complexity for Jorge. While Jorge avoids computing matrix inversion, it also introduces additional operations (lines 5-8 in Algorithm 2). It is necessary to analyze the complexity of these additional matrix multiplications compared to the matrix inversion in line 11 in Algorithm 1. Although the authors provide a wall-clock comparison between Jorge and Shampoo, it would be more insightful to have a theoretical complexity analysis.

Author response:
We value the reviewer's suggestion and will address it by incorporating a theoretical complexity analysis for computation and memory in our final manuscript.

Reviewer comment (Theory):
There is no convergence analysis. Besides empirical analysis, it is also important to analyze convergence and show how the approximation and related parameters affect the convergence speed.

Author response:
To the best of our knowledge, convergence analysis for adaptive optimizers like Shampoo has primarily been limited to convex functions. While we can conduct convergence analysis for Jorge in the context of convex functions, its utility for neural network training—characterized by highly non-convex loss landscapes—remains uncertain.

Reviewer comment (Experiments):
In experiments, the authors use step learning rate decay for Jorge, SGD, and Shampoo, but use a cosine decay for Adam. This makes it difficult to compare convergence speed with different learning rate schedulers. While different optimizers might prefer different schedulers, it would be more straightforward to compare optimizers with the same scheduler.

Author response:
We used the cosine decay schedule with AdamW on the findings of its authors, who demonstrated its superior generalization compared to step decay and constant learning rate schedules [1].

Reviewer comment (Evaluation):
According to Table 2, AdamW achieves validation accuracy close to Jorge. However, there is no convergence plot like Figure 2 showing the convergence speed of Adam compared to Jorge. Providing such a plot in the paper would make the comparison more complete.

Author response:
In Figure 3, we present a comparison of the convergence speed between AdamW and Jorge in terms of the number of epochs. The corresponding wall clock times for convergence in these experiments are detailed in Table 3. Hence, we have assessed the convergence speeds of Jorge and AdamW from both the perspective of epochs and wall clock times.

Reviewer comment (Experiments):
There is a lack of necessary ablation studies. For instance, it would be useful to investigate how the preconditioner update frequency affects convergence and final model performance.

Author response:
We plan to augment our final manuscript by incorporating additional comparisons with more second-order methods and including ablation studies to provide a more comprehensive analysis.

Reviewer comment (Experiments):
There is a need for more baselines. For reducing the complexities of using second-order methods, optimizers such as K-FAC [1] and mL-BFGS [2] also target the same goal. It would be good to compare Jorge against these methods in terms of per-iteration and wall-clock convergence performance.

Author response:
We plan to augment our final manuscript by incorporating additional comparisons with more second-order methods and including ablation studies to provide a more comprehensive analysis.

Reviewer comment (Evaluation):
Evaluation to other related methods is limited, where only a single second-order optimizer is compared to. The authors reference the recent SGD L-BFGS papers, but it should also be explained why these methods aren't compared to (i.e., implementations of these SGD L-BFGS have not been released). However, the following second-order SGD optimizer should be compared to in the paper: Yao, Zhewei, et al. "Adahessian: An adaptive second order optimizer for machine learning." proceedings of the AAAI conference on artificial intelligence. Vol. 35. No. 12. 2021.

Author response:
We acknowledge the reviewer's observation regarding the limited comparison with other second-order optimizers in our current manuscript. We agree that a more comprehensive evaluation involving a broader set of second-order optimizers is necessary. We plan to address this gap in our final draft by including comparisons with additional second-order optimizers like Adahessian and KAISA.

Reviewer comment (Experiments):
For ResNet50+Imagenet in the Heo et al. paper, both AdamW and SGD achieve ~76% accuracy, whereas in the current paper, AdamW saturates at ~70% (in Figure 2). Furthermore, Figure 2 clearly shows SGD and Jorge benefiting from the authors’ decay schedule, which is not a fair comparison for AdamW (i.e., SGD is further well-tuned, while AdamW is not). A fairer comparison for all methods would be to use the same hyperparameters for both SGD and AdamW from Heo et al. (2021), while also determining the cause of performance loss for AdamW observed in this work.

Author response:
We will investigate why AdamW did not reach the target validation accuracy of 76% for ResNet50 with a batch size of 1024. Interestingly, we were able to reach 76% validation accuracy with AdamW with a batch size of 256 (Figure 3, left). We opted for the cosine learning rate schedule with AdamW based on the findings of its authors, who demonstrated its superior generalization compared to step decay and constant learning rate schedules [1].

Reviewer comment (Evaluation):
Jorge requires running SGD prior to each step (due to Jorge's use of grafting SGD). However, Jorge's reported runtime in Table 3 is lower than SGD itself. Thus, Jorge's reported timings seem inaccurate; because Jorge is dependent on grafting SGD, its runtime should reflect this.

Author response:
Indeed, Jorge's per-iteration times may be slower than SGD due to the grafting process. However, Table 3 does not display per-iteration times; instead, it illustrates the total wall clock time for convergence. Notably, Jorge necessitates 25-40% fewer iterations than SGD to converge to a target validation metric, as illustrated in Figure 3. Consequently, the overall wall clock times are substantially smaller for Jorge, despite being slower on a per-iteration basis

Reviewer comment (Experiments):
Given the ubiquity and importance of transformer architectures used across applications (not only in NLP, but vision as well), Jorge's impact on transformer training should also be demonstrated. A smaller issue is, similarly, Adam/AdamW is widely used, even for vision (e.g., vision transformers, ViT), and the inclusion of this optimizer for bootstrapping Jorge's hyperparameters is warranted for current deep learning.

Author response:
We recognize the significance of showcasing Jorge's performance on transformer architectures, given their widespread application. To enhance the experimental section of our paper, we plan to include results on transformer architectures in the final manuscript. Moreover, considering the prevalent use of AdamW, particularly in transformer training, we intend to implement a bootstrapping strategy to derive hyperparameters from AdamW in our experimental evaluations.

Reviewer comment (Evaluation):
For Table 3, please include the (average) number of epochs required to achieve target performance.

Author response:
We will include the epochs for convergence in Table 3, which can be inferred from Figure 3 as well.

Reviewer comment (Writing):
The statement “For example, SGD remains the optimizer of choice in computer vision due to its better final validation accuracy, even though Adam converges faster initially” is a very debatable claim and should be revised. In particular, Adam is widely used as the out-the-gate optimizer for vision transformers (e.g., ViT) as well as CLIP. Furthermore, even in extensive benchmark studies, mixed results refute this statement, e.g., see: Kumar, Ananya, et al. "How to fine-tune vision models with sgd." arXiv preprint arXiv:2211.09359 (2022).

Author response:
We will adjust the claim regarding SGD being better than AdamW to be more nuanced and considerate.

Reviewer comment (Experiments):
What is the update frequency used in Shampoo when comparing wall-clock time? Is possible to use a second-order update interval for Shampoo such that it runs at a similar speed to Jorge to achieve the target accuracy?

Author response:
In line with Shi et al. [1], we use a second-order update interval of 50 for Shampoo. 

 It is possible to adjust the second-order update interval for Shampoo to achieve a similar iteration time as Jorge. For example, for ResNet-50 on 16 GPUs (as shown in Table 4, row 1), setting Shampoo's update frequency to approximately 1500 would result in similar iteration times as Jorge with an update frequency of 50. However, it's crucial to note that such large update intervals for Shampoo may impact convergence significantly, as demonstrated by Anil et al. [2] (Figure 4c of their paper).

Reviewer comment (Evaluation):
Could you explain why Jorge can enjoy well-tuned hyper-parameters from SGD while Shampoo cannot?

Author response:
We would like to clarify that Shampoo, as demonstrated by Shi et al. [1] and Anil et al. [2], can indeed borrow certain hyperparameters from well-tuned first-order optimizers such as SGD and Adagrad. Our contribution lies in the specific insights we provide in the paper regarding the determination of optimal weight decay and learning schedules from SGD. As highlighted in our study, this process is nuanced and not straightforward, and our findings contribute to the understanding of how these hyperparameters can be effectively set for Jorge.

Reviewer comment (Experiments):
Is Jorge possible to apply to NLP tasks like LLM pertaining or fine-tuning?

Author response:
We are currently actively investigating the applicability of Jorge to NLP tasks, including language model pretraining and fine-tuning. As part of our ongoing research, we aim to provide a comprehensive evaluation of Jorge's performance in the context of NLP applications.