Combining Axes Preconditioners through Kronecker Approximation for Deep Learning

It would be great if you could include some large scale Computer Vision experiments, I believe it would add a lot of value to the paper.

Regarding your observation about momentum, I guess $\beta_2 \rightarrow 1$ would result in not incorporating any second order information into the preconditioners $L$ and $R$. Can you please clarify that, I guess you refer to preconditioner updates as they appear in algorithm 1, line 4 (without any $\beta_2$ parameter), right?

Thank you for your valuable review.

In the evaluation you introduce a momentum for L and R matrices for a fair comparison against the other optimizers, can you please say what is the impact of momentum terms in this case? Have you experimented without momentum?

We conducted an additional experiment without momentum. We used the optimal hyperparameter for CASPR and removed momentum (i.e, equivalent to beta2 tends to 1) for L and R for CASPR, and noticed that the performance is poor. This could be because in practice momentum is helpful as the curvature could change rapidly. We will conduct a full experiment with several hyperparameters in the next revision.

at page 4, Lemma 3.1, the Frobenius norm should be…

Thank you for finding this error, we corrected it.

at page 5, when you rewrite the preconditioner…

Yes, this $\frac{(L^{-1/2}\otimes I_n + I_m \otimes R^{-1/2})/2 + L_t^{-1/4} \otimes R_t^{-1/4}}{2}$ should be $\frac{(\tilde{L}^{-1/2}\otimes I_n + I_m \otimes \tilde{R}^{-1/2})/2 + \tilde{L}_t^{-1/4} \otimes \tilde{R}_t^{-1/4}}{2}$. We edited this.

the paper does not have experiments on computer vision tasks, such as ResNets, ViTs or other important architectures used for benchmarking in the literature

We will try to conduct more experiments on vision tasks in the next revision.

By $\beta_2 \rightarrow 1$, we mean to update $L_t$ and $R_t$ as $L_t = L_{t-1} + G_tG_t^{\top}$ and $R_t = R_{t-1}+ G_t^{\top}G_t$, so the entire history of gradient second order information is accumulated in $L_t$ and $R_t$ without exponential moving average. Please let us know if we understood your suggestion correctly.

Thank you for your valuable review.

abstract and introduction could be written more clearly, it is only mentioned what is achieved in the abstract and introduction but not how this will be done

Thank you for pointing this out, we will add more details regarding our method in the introduction of the next revision.

On page 3 it says the second-moments are computed per row of the gradient. Is it correct to say that this means it is a neuron-level conditioner? This is something that seems to exist in the context. See for example, Fig 1 "unit-wise" preconditioning in https://arxiv.org/pdf/2305.04684.pdf.

Yes, preconditioning each row individually can be imagined as a neuron-level preconditioner. We will cite the mentioned paper in the final draft. Specifically, this originates from a block-diagonal approximation $S_{R} = \{ blockdiag(R_1,R_2\ldots,R_m): R_i\succeq 0\in R^{n\times n},\forall i\in[m]\}$ of full-matrix statistic. Since maintaining this approximation and computing the corresponding preconditioner is infeasible in terms of memory and compute, we additionally develop a statistic/preconditioner following the constraint $S_{I\otimes R} = \{ blockdiag(R,R\ldots,R): R \succeq 0\in R^{n\times n}\}$ in Lemma 3.1, where the additional constraint is that all the blocks $R_1=R_2=\cdots=R_m=R$ are the same. Our work can be seen as extending such neuron level preconditioners by combining them through kronecker approximations. In practice this does give us a big gain as seen in our experimental results.

Are the 300 hyperparameters optimized for each method or only for caspr? What is the target for this hyperparameter tuning? Validation accuracy is the performance reported so is the grid search optimizing that value?

For ogbg-molpcba, we use 300 hyperparameters for Shampoo and AdamW and 200 for CASPR randomly sampled from the search space in Table 1 (Appendix A.4). The hyperparameter configuration which gives the best generalization performance is picked and reported individually for each method.

Why are there no error bars included? Could the experiments be run on at least 3 seeds? Without a single rerun, the results could be random.

We mention here the standard deviations averaged across 3 seeds for the Transformer on universal dependencies and OGBG-molpcba benchmarks:

We will try to add similar numbers for our language modeling benchmarks in the next revision, as they are more time consuming.

seems unclear how to extend this method to layers that are not fully-connected

Our method (as well as Shampoo) can be applied to any layer. For example, in our experiments with transformers we handle 3D parameters, such as attention layers by flattening the parameter (in the form of higher order tensors) into a 2D tensor by merging the dimensions. The attention layers have 3D parameter (m,n,k) which we precondition by merging the second and third dimensions to a 2D parameter (m, n*k) and then applying CASPR. For higher dimensional parameters, one can apply a similar transformation by merging more dimensions.

first page sentence ... "..to be at one end of this sequence" confusing

Thank you for pointing this out. Here, we added a forward reference to Lemma 3.2 to address the confusion.

Thank you for your valuable review. We will add the additional details about coupled Newton algorithm for matrix inverse pth root (Appendix D in [1]) in our final draft. Here we clarify a few questions:

how to choose the damping term ?

For our experiments, the inverse pth root operations were performed using single-precision floating-point format (float32). We determined the damping term, $\epsilon$, by scaling the largest eigenvalue of $\lambda_{max}(L)$ with $\epsilon'$ as outlined for the Shampoo algorithm in line 16 of Algorithm 1 on page 20 in [1]. This scaling ensures that the modified matrix $\tilde{L} = L + \epsilon' \lambda_{max}(L) I_m$, maintains an $\ell_2$-condition number not exceeding $1/\epsilon’$. We set $\epsilon'$ to $10^{-6}$ across all our experiments, which caps the condition number of $\tilde{L}$ at $10^{6}$, except in the case of the 14M parameter language model showcased in Figure 4a. In that instance, we experimented with $\epsilon'$ values in the set {$10^{-6}, 10^{-10}$}, as detailed in Table 3.

The proposed method may not work well in a low numerical precision setting such as float-32 or bfloat-16. The authors should explicitly discuss this point.

Thank you for catching this limitation. The inverse pth root operations of our experiments were conducted in float32. Nevertheless, conducting inverse pth root operations can be limiting in float32 as accurate inverses can only be obtained for matrices with condition number less than $\sim 10^6$. We leave conducting accurate inverse pth root operations tolerant to low precision as a future work.

[1] Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09018, 2020.