Improving Adaptive Moment Optimization via Preconditioner Diagonalization

I think there are two major issues with this work:

1. On the theoretical side, the line “Since $\Sigma_{\tau}$ is a diagonal matrix, we have $\text{vec}(\Sigma_{\tau}) \, \text{vec}(\Sigma_{\tau})^{\top}$ is almost diagonal (off-diagonal elements are mostly zero).” is the crux of their proof. But this is clearly not fully correct since $\text{vec}(\Sigma_{\tau}) \, \text{vec}(\Sigma_{\tau})^{\top}$ is a rank-1 matrix. The only matrix which is both rank 1 and diagonal is a matrix with a single non-zeros entry on its diagonal. I agree with the authors that this matrix is more diagonal than the original matrix due to $(\Sigma_{\tau}$ having a strong decay along the diagonal.

   There is in fact a bigger issue with this approach: the matrix the authors begin with $C(G_{\tau}) = \text{vec}(G_{\tau}) \, \text{vec}(G_{\tau})^{\top}$ is itself is a rank 1 matrix. There is a trivial way to rotate this matrix to make it exactly diagonal which is to use a basis with one of the elements being $\text{vec}(G_{\tau})$. If the authors believe diagonalization is the main motivation they should use any such basis. But under this basis applyling Adam on the gradient (assuming it remains nearly constant which while not true is anyways assumed by the authors since frequency > 1) will just lead to gradient normalization per layer since all of the projected gradient will lie in the coordinate corresponding to the $\text{vec}(G_{\tau})$ basis element.

2. One the empirical side:

   A. The authors do not give details of their hyperaprameters ablations for their results.

   B. Since the optimizer has connections to Shampoo the authors should compare it to Shampoo (see https://github.com/facebookresearch/optimizers/tree/main/distributed_shampoo for an implementation).

There are some other issues with the discussion of prior and related works which I can expand on once the authors address these two points.

1. Although the method proposed by the paper clearly lies within the class of second order methods, no second order baselines have been provided in the work. The paper needs to compare to other second order methods such as Shampoo, KFAC, SOAP as baselines.

2. Also, in section 3.3, the paper makes the claim that the preconditioner basis provided by SOAP does not have any theoretical basis. However, SOAP uses the same basis as Shampoo, for which the theoretical basis in terms of optimal Kronecker approximation has already been established in previous works [1,2].

   [1] - Gupta et al. 2018, Shampoo: Preconditioned Stochastic Tensor Optimization 2018

   [2] - Morwani et al. 2024, A New Perspective on Shampoo's Preconditioner,

Overall, without proper comparison to other existing second order methods, it is unclear if there is any significant contribution of this work.

I will describe the weaknesses of each of the three main contributions listed in the introduction (lines 068-078):

1. Regarding the proposed method itself, the paper ignores highly relevant previous work: Liu et al. (2018) [1] and George et al. (2018) [2] also propose to rotate the parameter space in a way that the preconditioner is approximately diagonal. Liu et al. propose this to improve the continual learning method EWC and George et al. aim to improve the optimization performance of K-FAC; they call the resulting method EK-FAC. One difference to this paper is how the transformation matrix is chosen since they use the eigenbasis of the K-FAC preconditioner, which is very closely related to the here proposed transformation matrix (the left- and right-singular vector of the gradient matrix, or equivalently, the eigenvectors of Shampoo’s preconditioner without any accumulation), but not exactly the same. The choice of the transformation matrix proposed in this paper has also been mentioned in Appendix B of Anil et al. (2020) [3] (the only difference being whether we use accumulation to compute Shampoo’s preconditioner), who describe how to adopt the EK-FAC method to use the eigenvectors of Shampoo’s instead of K-FAC’s preconditioner as the transformation matrix. Finally, the authors mention SOAP by Vyas et al. (2024), which is concurrent work. However, I think the statement that “their approach seems to lack an interpretable framework to explain its effectiveness” (c.f. lines 371-372) is incorrect: since Vyas et al. discuss the related work mentioned here, which already provided a framework to explain the effectiveness of this class of methods, they inherit all the explanations that are valid for any of these works. I will increase my score if an accurate discussion of all of this previous work is added.

2. Regarding the convergence analysis in section 4, it seems like the stated result is already contained in the cited work by Liang et al. (2024): the case of AdaDiag corresponds to the results in section 4.1+4.2 in Liang et al. (here the transformed gradient is $P_t^T \nabla \mathcal{L}(W_t)$) and the case of AdaDiag++ is described in section 4.3 where the results are generalized to general linear operators (including the explicitly mentioned case where the transformed gradient is $P_t^T \nabla \mathcal{L}(W_t) Q_t$). The choice of the projection matrices differs between the two works since Liang et al. are motivated by memory-efficient optimization, but unless I’m missing something this doesn’t seem to affect the convergence analysis significantly.

3. While the experimental results are promising, they are limited in multiple ways: a) All results seem to focus on either final validation performance or convergence in iterations. However, in the chosen optimization settings it is crucial to evaluate the optimization performance in wall-clock time as well. You mention that the computational overhead is less than 10% in the text, but don't show any experiments for this. b) To substantiate claims that AdaDiag variations “can substantially enhance the convergence speed of modern adaptive optimizers” (abstract), I would expect to see a wider range of benchmarks, beyond the two datasets and three architecture considered here. c) Alternatively, if the claims were restricted to the image classification and language model settings considered here, I would like to see a slightly more thorough evaluation, e.g. runs with multiple random seeds, plots that compare the runs in wall-clock time as well, either more tuning of all methods or a more elaborate justification why you would expect the Adam baseline to be strongly tuned already, and additional ablations, e.g. how the performance depends on the hyperparameter $T$ which is introduced on top of Adam's hyperparameters.

[1] Liu et al. (2018). Rotate your Networks: Better Weight Consolidation and Less Catastrophic Forgetting

[2] George et al. (2018). Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis

[3] Anil et al. (2020). Scalable Second Order Optimization for Deep Learning

The proposed method requires extra memory consumption and extra computational overhead brought by SVD. More rigorous discussion on these aspects is needed. See below.