Paper: FOSI: Hybrid First and Second Order Optimization (ICLR-2024)

Reviewer comment (Evaluation):
The method requires computing the extreme eigenvalues and vectors of the Hessian by the Lanczos algorithm, which raises much more computational cost per iteration than first-order methods.

Author response:
While the average FOSI iteration takes longer, FOSI requires much fewer iterations to reach the same loss compared to the base optimizer. This can be visualized in Fig 5 in the supplementary material. FOSI's ability to use fewer steps is due to its utilization of additional second-order information in each iteration, thereby facilitating more precise descent steps. We stress that the crucial factor is the time to convergence, not the number of internal steps (iterations).

Reviewer comment (Experiments):
It is unknown how to choose \( k \) and \( l \) in the proposed algorithm, and neither the theoretical analysis nor the empirical study includes evaluation on how \( k \) and \( l \) affect the behavior of the methods.

Author response:
Since our focus was for FOSI as a drop-in replacement without additional tuning, we did not tune $k$ or $\ell$. Instead, for optimizing non-convex losses like DNNs, we recommend setting $k$ to a small number (less than the number of classes), and  $\ell$ to 0. This is supported by the common observation that the number of extreme eigenvalues is generally small, on the order of the number of classes (see for example reference [R1] below), and small eigenvalues predominantly exhibit negative values in real DNN settings. It is worth noting that even if $k$ is lower than the actual number of outlier eigenvalues, the distance between these large eigenvalues can potentially be large, hence the effect on the condition number is still considerable.  
Similarly, $\alpha$ can simply be set to a small conservative constant such as 0.01, since the bottleneck for convergence is the base optimizer used for $f_2$, not the Newton method optimizer for $f_1$.  

In practice, as supported by our experiments, it is perfectly fine to simply set $k=10$, $\ell=0$ and $\alpha=0.01$ across a wide range of tasks and architectures.  

In our future research, we aim to develop innovative techniques that dynamically adjust $k$ and $\ell$ based on observed improvements in the effective condition number. This approach will enable us to strike a balance between computational cost and optimization performance. However, we acknowledge that this research extends beyond the scope of the current paper.

Reviewer comment (Experiments):
The parameters are too many; not only \( k \) and \( l \) need to be chosen, but Algorithm 2 also requires the parameter for learning rates \(\alpha\), the learning rate for the base optimizer.

Author response:
Since our focus was for FOSI as a drop-in replacement without additional tuning, we did not tune $k$ or $\ell$. Instead, for optimizing non-convex losses like DNNs, we recommend setting $k$ to a small number (less than the number of classes), and  $\ell$ to 0. This is supported by the common observation that the number of extreme eigenvalues is generally small, on the order of the number of classes (see for example reference [R1] below), and small eigenvalues predominantly exhibit negative values in real DNN settings. It is worth noting that even if $k$ is lower than the actual number of outlier eigenvalues, the distance between these large eigenvalues can potentially be large, hence the effect on the condition number is still considerable.  
Similarly, $\alpha$ can simply be set to a small conservative constant such as 0.01, since the bottleneck for convergence is the base optimizer used for $f_2$, not the Newton method optimizer for $f_1$.  

In practice, as supported by our experiments, it is perfectly fine to simply set $k=10$, $\ell=0$ and $\alpha=0.01$ across a wide range of tasks and architectures.  

In our future research, we aim to develop innovative techniques that dynamically adjust $k$ and $\ell$ based on observed improvements in the effective condition number. This approach will enable us to strike a balance between computational cost and optimization performance. However, we acknowledge that this research extends beyond the scope of the current paper.

Reviewer comment (Experiments):
There are some optimizers which also use the hybrid directions of first and second-order methods, or partial information of the Hessian; the author may need to compare and discuss them [1, 2].

Author response:
Concerning Sophia, we emphasize that, despite the paper's title containing "Second-order Optimizer," Sophia bears a close resemblance to Adam, as both approximate the Hessian diagonal. The distinction lies in Adam's utilization of the second moment of the gradient and sqrt for a conservative approximation, whereas Sophia employs the gradient of predictions and clipping. As such, we can expect Sophia to suffer from similar issues as Adam for non-diagonal Hessians (see Section 4.3 and Section B.2.1). We will add this discussion to the paper.  

DRSOM is a first-order method, specifically an approach for determining the most accurate step size for gradient methods incorporating Polyak's momentum. To our knowledge, DRSOM has not been applied to DNNs, but for problems with $n<1000$ dimensions and small datasets. We have been unable to find a Python implementation to support a direct comparison for (our much larger) tasks.

Reviewer comment (Experiments):
The authors should provide some discussion and numerical evaluation on how to choose \( k \) and \( l \).

Author response:
Since our focus was for FOSI as a drop-in replacement without additional tuning, we did not tune $k$ or $\ell$. Instead, for optimizing non-convex losses like DNNs, we recommend setting $k$ to a small number (less than the number of classes), and  $\ell$ to 0. This is supported by the common observation that the number of extreme eigenvalues is generally small, on the order of the number of classes (see for example reference [R1] below), and small eigenvalues predominantly exhibit negative values in real DNN settings. It is worth noting that even if $k$ is lower than the actual number of outlier eigenvalues, the distance between these large eigenvalues can potentially be large, hence the effect on the condition number is still considerable.  
Similarly, $\alpha$ can simply be set to a small conservative constant such as 0.01, since the bottleneck for convergence is the base optimizer used for $f_2$, not the Newton method optimizer for $f_1$.  

In practice, as supported by our experiments, it is perfectly fine to simply set $k=10$, $\ell=0$ and $\alpha=0.01$ across a wide range of tasks and architectures.  

In our future research, we aim to develop innovative techniques that dynamically adjust $k$ and $\ell$ based on observed improvements in the effective condition number. This approach will enable us to strike a balance between computational cost and optimization performance. However, we acknowledge that this research extends beyond the scope of the current paper.

Reviewer comment (Writing):
In the statement (step 2m Lemma 1): "The preconditioner P is symmetric and PD:.", is it possible to use a different font for the matrix "P", which has no conflict of P (positive) in "PD"? Or vice-versa.

Author response:
We are committed to clear notations and will differentiate more explicitly between the preconditioner’s P and PD for positive definiteness.

Reviewer comment (Reproducibility):
Hope the code of the paper, to implement FOSI for any application of Optimization, will be released soon by authors.

Author response:
Absolutely! The FOSI repository is ready to go, showcasing solid implementations in both Jax and Python. Its API conforms with commonly used packages in these frameworks (Optax and TorchOpt) to facilitate seamless drop-in replacement. Upon the lifting of anonymization requirements, FOSI will be promptly released. You can visit the anonymized repository at https://anonymous.4open.science/r/fosi-497F.

Reviewer comment (Reproducibility):
One of my major concern is that the memory consumption and computational complexity are very high especially for large-scale neural networks. This will limit the usage of the proposed method. Besides, it is not clear how to handle the communication cost and the computation of $V$ in the distributed setting.

Author response:
One of the key areas of focus in our future research will be addressing the memory overhead associated with FOSI, as well as parallelization of the algorithm. The first enhancement is to run the ESE procedure on the CPU in parallel to the training process occurring on the GPU. By carefully managing the synchronization between these processes, we aim to reduce the computation bottleneck of computing second-order information and overall training time. This will also reduce FOSI’s memory consumption on the GPU by moving the ESE scratch space memory to the more plentiful CPU memory.

Concerning the execution of the ESE in a distributed setting, our FOSI implementation already supports parallel computation of the dominant $hvp$ operation across multiple GPUs when available (our experiments were performed on a single GPU, however). $hvp$ computation involves gradient computations, allowing for efficient parallelization over the batch axis. We will include this discussion in the final version.

Reviewer comment (Experiments):
In comparison with second-order methods, it is better to compare the proposed with KLBFGS but not LBFGS because the former method is designed specifically for deep learning tasks. It is suggested to add comparison with Shampoo[1] or NGPlus[2], which has good effectiveness in practice.

Author response:
Block diagonal preconditioners such as KBFGS (Goldfarb et al., 2020), Shampoo and its variants [1], and NGPlus [2] have similar limitations as diagonal preconditioners, though to a lesser extent. They may not fully capture the problem's rotation. In contrast, by splitting the problem into two subspaces as done by FOSI, we obtain a full low-rank representation of the Hessian for the first subspace ($\hat{V}$), which captures both the rotation and curvature of the sub-problem. This contributes to accuracy and stability of the optimization, particularly as it is based on extreme eigenvalues and vectors that can be approximated more accurately. We will clarify the relevant discussion in the Related Work section.

In particular, KLBFGS is similar to K-FAC as both algorithms employ a Kronecker factored block-diagonal approximation of the Hessian, and both perform similarly (see Goldfarb et al., 2020). As our experiments show, while K-FAC is indeed better than classic LBFGS, its performance is still substantially worse than SGD, HB, Adam, and their FOSI variants.

Concerning the specific paper by Anil et al. [1], this paper proposes a distributed variant of the Shampoo optimizer (Gupta et al., 2018) that also uses parallel CPU computation. As noted above, while we intend to explore such optimizations for FOSI in our future research, we consider it beyond the scope of the current work.

Finally, we wish to stress that FOSI provides a new alternative formulation for hybrid 2nd-order optimizers: while prior work estimates the inverse preconditioner (the Hessian), FOSI estimates the preconditioner directly. Inverse preconditioner approaches such as K-FAC, K-BFGS and Shampoo have been developed and improved over years of research, and there are now many proposed second-order optimizers based on their principles with various targeted optimizations. We therefore opted to compare our approach with the two most representative ones so that we can focus on the difference between our direct preconditioner approach and inverse preconditioner approaches, rather than specific optimization.

Reviewer comment (Experiments):
It is better to tune the hyper-parameters of ADAM in numerical experiments to make the baseline strong enough.

Author response:
The primary focus of our paper is to demonstrate FOSI's improvement over the base optimizer, without additional tuning, and even when the base optimizer is not perfectly tuned. This aligns with our drop-in replacement concept for FOSI. Rather than tuning Adam, we show in Section 4.3 and Appendix B 2.2, figure 9, that FOSI improves performance across a wide range of learning rate and momentum. Also, given Adam's inherent adaptability, it tends to perform well across a relatively broad spectrum of hyperparameters.

Reviewer comment (Experiments):
Can randomized numerical algebra be combined in ESE procedure? This will help reduce the computational cost.

Author response:
Thank you for this suggestion. We are actively exploring multiple avenues to enhance the efficiency of ESE computation. It's noteworthy that, in line with the reviewer's suggestion of randomized numerical algebra, our current approach already incorporates elements of this concept. Specifically, we utilize a set of samples to approximate the Hessian in the $hvp$ computation. However, we acknowledge there is considerable potential for further optimization in this direction, including the exploration of techniques such as warm starting the Lanczos algorithm and other more intricate methods.