Title: Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner
Venue: NeurIPS 2025
Accepted: yes
Keywords: neural network training algorithms

optimization

Shampoo

Kronecker-factored

preconditioner

grafting

approximation

adaptive gradient methods

stochastic optimization
Abstract: The recent success of Shampoo in the AlgoPerf contest has sparked renewed interest in Kronecker-factorization-based optimization algorithms for training neural networks. Despite its success, Shampoo relies heavily on several heuristics such as learning rate grafting and stale preconditioning to achieve performance at-scale. These heuristics increase algorithmic complexity, necessitate further hyperparameter tuning, and lack theoretical justification. This paper investigates these heuristics from the angle of Frobenius norm approximation to full-matrix Adam and decouples the preconditioner's eigenvalues and eigenbasis updates. We show that grafting from Adam mitigates the staleness and mis-scaling of the preconditioner's *eigenvalues* and how correcting the eigenvalues directly eliminates the need for learning rate grafting. To manage the error induced by infrequent *eigenbasis* computations, we propose an adaptive criterion for determining the eigenbasis computation frequency motivated by terminating a warm-started QR algorithm. This criterion decouples the update frequency of different preconditioner matrices and enables us to investigate the impact of approximation error on convergence. These practical techniques offer a principled angle towards removing Shampoo's heuristics and developing improved Kronecker-factorization-based training algorithms.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the practical tricks of the optimizer Shampoo to address the critical problem of stale updates. The paper demonstrates theoretically that the stale update of Shampoo can lead to eigenevalue scale mismatch, resulting in significantly worse performance, which can be fixed by tricks like grafting and correcting eigenvalues (e.g., SOAP). Then the authors discuss the effect of stale eigenbasis and propose an adaptive way to update the eigenbasis to solve the problem. Empirical results are also provided to verify the claims.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. The paper is well-written and basically clear for understanding.\n2. The viewpoint of eigenvalue scale mismatch for explaining the effectiveness of grafting and SOAP is interesting, providing a novel and valid perspective for the motivation of these Shampoo variants. \n3. The study of the influence of the preconditioner staleness is detailed and solid, providing a comprehensive understanding of the problem.\n\nWeakness:\n1. The theoretical perspective in Section 3 on why grafting is useful may not be that rigorous and convincing. The bound in Lemma 1 and Proposition 1 is very loose, since the largest entry value/singular value can be much larger than the smallest ones. I don't think this bound is valid for justifying that the eigenvalue scale mismatch is the problem of Shampoo without grafting. \n2. In Section 4, the adaptive update rule of the eigenbasis is claimed to propose for reducing the effort of manually tuning the update frequency. However, it turns out that the adaptive threshold $\\tau$ also needs tuning and I don't see strong evidence in that the optimal choice can transfer over models and scales. This may make the idea not that powerful.\n\nQUESTIONS:\n1. Based on the explanation for why grafting is useful for Shampoo, I am curious why don't we just use $\\sqrt{mn}$ instead of $||U_t^{Adam}||_F$ in grafting for a matrix $W\\in\\mathbb{R}^{m\\times n}$, since they should have similar scales. Have you tried this kind of trivial scaling?\n2. I want to double check what the wall-clock time stands for in e.g., Figure 3. Is it the total training time including forward and backward or any other measure? Also, could you post the proportion of the optimizer step time in the whole training process? Based on the figure, the optimizer step computation takes the major time costs, which seems a bit strange, even though matrix operations are required.\n3. The observation that $L_t$ seems to update more than $R_t$ looks interesting. Do you have more evidence for this in other models or other scales? Also, any intuition for explaining this would be welcome.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper revisits practical tricks that have become standard when training with Shampoo: spacing out the eigendecomposition of its Kronecker factors and applying learning-rate grafting to match Adam’s step size. It shows that grafting compensates for eigenvalue issues: staleness between decompositions and a scale bias tied to layer shape, while the eigenvectors largely stay relevant. Replacing those stale eigenvalues each iteration with Adam-style second-moment estimates removes the need for grafting, and a simple Frobenius-norm test decides when the costlier eigenbasis must be refreshed. Experiments confirm that this “purified” variant matches or surpasses the grafted baseline with less compute.\n\nSTRENGTHS AND WEAKNESSES:\nThe paper addresses a highly motivated problem. The paper offers an interesting explanation of why grafting is needed in standard Shampoo. By separating the Frobenius-norm effect of eigenvalues from the rotational effect of eigenvectors, they show that grafting merely rescales updates that have become too small or too large because the eigenvalues are both out-of-date. This clarifies a long-standing practical mystery.\n\nWhile your analysis of grafting’s role in Shampoo and the supporting experiments are insightful, it’s still unclear what new algorithmic insight EShampoo contributes beyond SOAP, which already applies an Adam-style diagonal second moment in the fixed Shampoo eigenbasis.\n\nQUESTIONS:\nCould you please clarify how Eigenvalue-Corrected Shampoo meaningfully differs from SOAP? Appendix B of Anil et al. (2020) notes that the Frobenius-optimal diagonal Kronecker approximation in a fixed Shampoo eigenbasis is precisely the Adam second moment, and SOAP implements this by keeping the Shampoo basis (refreshed every F steps) while applying an Adam/Adafactor-style diagonal update at each iteration, so it seems EShampoo follows the exact same recipe.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work studies the Shampoo optimizer and two key heuristics (learning rate grafting and preconditioning frequency) important for its competitive performance. Based on empirical results and theoretical motivation, the authors claim that the learning rate grafting addresses stale and mis-scaled eigenvalues of the preconditioner, and this need can be removed by decoupling the eigenvalue updates from the eigenbasis. The authors also propose a method tracking approximation-error of the Kronecker factors to adaptively determine eigenbasis update frequency. Their method also allows one to track which components of the network require more frequent eigendecomposition computations, showing that biases and layer normalization parameters require more frequent updates.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n* The paper is clearly written and easy to follow. The relevant background and algorithms were presented particularly well.\n* Connecting learning rate grafting to previously known modifications of Shampoo (using Shampoo^2 with trace scaling, low PF, eigenvalue correction a la SOAP) is a nice result.\n* Investigations about determining eigenbasis computation frequency, its effect on different network parameters, and plotting tradeoffs between efficiency and performance is an important direction.\n\nWeaknesses:\n* The empirical setting is limited to two architectures on a single dataset and it would’ve been particularly interesting to see how the eigenbasis frequency results in eg. Figure 4 change for other domains, relating to previous work similarly identifying key network parameters requiring adaptivity due to issues like class imbalance for the last layer for instance [1,2].\n* The authors’ proposed adaptive QR algorithm provides efficiency gains but some other baselines may better inform whether this adaptive control based on thresholding the approximation error is necessary - please see Questions.\n\nReferences:\n\n[1] Kunstner, Frederik, et al. \"Heavy-tailed class imbalance and why adam outperforms gradient descent on language models.\" Advances in Neural Information Processing Systems 37 (2024): 30106-30148.\n\n[2] Zhao, Rosie, et al. \"Deconstructing what makes a good optimizer for language models.\" arXiv preprint arXiv:2407.07972 (2024).\n\nMinor: \n* Authors should add a reference to Appendix C after presenting Lemma 1 and Proposition 1 in the main paper.\n\nQUESTIONS:\n* How would the authors’ proposed adaptive criterion of the QR algorithm compare to having frequency warmup (starting with frequency = 1 and roughly gradual increase to F = 50) on all of the parameters, or a fraction of the parameters (in terms of loss vs wall-clock time)? \n* Do the authors have more intuition about the observed need for frequent updating for layer norm parameters? If the LayerNorms did not have learnable parameters, would the matrix layers then require more frequent computations (the LayerNorm parameters act as some output scale for the given layer)?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}