Paper: Combining Axes Preconditioners through Kronecker Approximation for Deep Learning (ICLR-2024)

Reviewer comment (Theory):
at page 4, Lemma 3.1, the Frobenius norm should be squared (the 2nd power is missing from the norm)

Author response:
Thank you for finding this error, we corrected it.

Reviewer comment (Writing):
at page 5, when you rewrite the preconditioner $X_t^{caspr}(1)$ in terms of $X_t^L$ and $X_t^R$, I believe the matrices $L$ and $R$ are missing a tilde, based on the definitions of the $X_t^{L/R}$ terms 4 lines above

Author response:
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.

Reviewer comment (Writing):
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

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

Reviewer comment (Reproducibility):
relatively limited benchmarking to other methods and missing error bars on the experiments, it seems like a single training run is reported in the figures. This is the major potential weakness I see in the paper and I am worried that the results are not reproducible. See my questions below (Are the figures created from multiple runs averaged? What is the target for tuning hyperparameters if the plots show validation performance?).

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


|Optimizer|Transformer on Universal dependencies - Accuracy|OGBG-molpcba- Test mAP|
|---------|------------------------------------------------|----------------------|
|Shampoo  |69.45 +/- 0.26%                                 |0.2843 +/- 0.0028     |
|CASPR    |69.76 +/- 0.14%                                 |0.2873 +/- 0.0017     |
|AdamW    |68.28 +/- 0.39%                                 |0.2701 +/- 0.0011     |

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

Reviewer comment (Writing):
first page sentence ... "..to be at one end of this sequence" confusing

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

Reviewer comment (Writing):
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.

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Evaluation):
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.

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


|Optimizer|Transformer on Universal dependencies - Accuracy|OGBG-molpcba- Test mAP|
|---------|------------------------------------------------|----------------------|
|Shampoo  |69.45 +/- 0.26%                                 |0.2843 +/- 0.0028     |
|CASPR    |69.76 +/- 0.14%                                 |0.2873 +/- 0.0017     |
|AdamW    |68.28 +/- 0.39%                                 |0.2701 +/- 0.0011     |

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