Pith. sign in

REVIEW 2 major objections 1 minor 26 references

Reviewed by Pith at T0; open to challenge.

T0 means a machine referee read the full paper against a public rubric. The mark states how deep the mechanical check went, never who wrote it. the ladder, T0–T4 →

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

T0 review · grok-4.3

A reparametrization of Shampoo preconditioners allows BFloat16 storage and partial subspace QR updates while forming a complete basis.

2026-06-29 22:14 UTC pith:XFNC4AWY

load-bearing objection The reparametrization for partial subspace QR updates lets Shampoo variants run under BFloat16 with lower cost, but the paper leaves the key claim about the hybrid basis unproven. the 2 major comments →

arxiv 2605.26327 v1 pith:XFNC4AWY submitted 2026-05-25 cs.LG

Reparametrizing Shampoo and SOAP for Subspace Basis Updates and BFloat16 Storage

Alan Milligan , Zikun Xu , Simon Lacoste-Julien , Felix Dangel , Wu Lin This is my paper
classification cs.LG
keywords Shampoo optimizerSOAPBFloat16QR decompositionpreconditionersubspace basis updatelow-precision trainingneural network optimization
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper introduces a reparametrization for Shampoo-based optimizers that store preconditioners in BFloat16 and update only a subspace of the basis via QR decomposition. The updated vectors are combined with unchanged ones to maintain a full basis. This cuts the cost of full-precision QR while avoiding the accuracy loss that BFloat16 normally causes in these methods. The change applies to KL-Shampoo, SOAP, and KL-SOAP and lets KL-SOAP reach or surpass KL-Shampoo performance under the lower-precision regime.

Core claim

The central claim is that reparametrizing the preconditioner to combine a subspace-updated basis (obtained via QR) with the remaining unchanged vectors produces a complete basis that supports BFloat16 storage, lowers computational cost, and removes the performance penalty BFloat16 otherwise imposes on Shampoo-style methods.

What carries the argument

Reparametrization that forms a complete basis by merging subspace-QR-updated vectors with unchanged basis vectors.

Load-bearing premise

Combining the partially updated basis vectors with the unchanged ones still produces a complete basis that keeps the original preconditioning effect intact.

What would settle it

Measure whether training loss or final accuracy on a fixed model and dataset drops when the reparametrized BFloat16 version replaces full FP32 QR on the entire basis.

Watch this falsifier — get emailed when new claim-graph text bears on it.

If this is right

  • The method works for any Shampoo variant that uses QR, including KL-Shampoo, SOAP, and KL-SOAP.
  • SOAP and KL-SOAP regain performance under BFloat16 storage.
  • KL-SOAP can match or exceed KL-Shampoo while using less memory and time.
  • Shampoo-based optimizers become both more memory-efficient and faster overall.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • The same subspace idea could be tested on other matrix-factorization preconditioners that rely on QR.
  • Further shrinking the updated subspace size might yield additional speed gains if the basis-stability property continues to hold.
  • Adoption would let practitioners train larger models with these second-order methods without doubling memory for FP32 storage.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Referee Report

2 major / 1 minor

Summary. The paper claims that reparametrizing Shampoo-based preconditioners (KL-Shampoo, SOAP, KL-SOAP) to update only a subspace of the QR basis vectors while combining them with unchanged vectors produces a complete orthonormal basis. This allows BFloat16 storage, reduces QR computational cost for large matrices, mitigates precision-induced degradation, and enables KL-SOAP to match or exceed KL-Shampoo performance.

Significance. If the hybrid-basis construction preserves orthonormality and the original column space (so that the inverse-square-root preconditioner remains valid), the method would make second-order optimizers more memory- and time-efficient for large-scale training; the explicit credit for a practical reparametrization that targets both storage and compute is a strength.

major comments (2)
  1. [reparametrization description] The reparametrization section: the assertion that 'combining updated basis vectors with unchanged ones forms a complete basis' lacks any algebraic derivation or verification that the resulting matrix Q satisfies Q^T Q = I, spans the same column space as the original QR factorization, or induces only a bounded perturbation on the preconditioner (typically involving the inverse square root of the factored matrix). This step is load-bearing for all claims about preserved optimizer dynamics and BFP16 stability.
  2. [method] No perturbation bound, stability argument, or explicit reparametrization formula (e.g., how the hybrid Q is assembled from the subspace QR output and the retained vectors) is supplied; without it the central performance claims cannot be confirmed to follow from the construction rather than from unanalyzed side effects.
minor comments (1)
  1. [Abstract] Abstract could briefly note the key algebraic property (orthonormality and span) that the reparametrization is intended to preserve.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for the careful review and constructive criticism. The points raised correctly identify that the reparametrization section would benefit from additional algebraic detail and analysis. We address each major comment below and will incorporate the requested material in the revision.

read point-by-point responses

  1. Referee: [reparametrization description] The reparametrization section: the assertion that 'combining updated basis vectors with unchanged ones forms a complete basis' lacks any algebraic derivation or verification that the resulting matrix Q satisfies Q^T Q = I, spans the same column space as the original QR factorization, or induces only a bounded perturbation on the preconditioner (typically involving the inverse square root of the factored matrix). This step is load-bearing for all claims about preserved optimizer dynamics and BFP16 stability.

    Authors: We agree that the manuscript currently asserts the completeness of the hybrid basis without supplying the algebraic derivation or the required verifications. In the revised manuscript we will add a dedicated subsection that (i) states the explicit construction of the hybrid Q, (ii) proves Q^T Q = I by direct computation on the partitioned blocks, (iii) shows that the column space is unchanged because the retained vectors already lie in the original space and the updated block is obtained from a QR on a subspace of the same space, and (iv) derives a first-order perturbation bound on the inverse-square-root preconditioner that is controlled by the subspace dimension and the BFloat16 rounding error. revision: yes

  2. Referee: [method] No perturbation bound, stability argument, or explicit reparametrization formula (e.g., how the hybrid Q is assembled from the subspace QR output and the retained vectors) is supplied; without it the central performance claims cannot be confirmed to follow from the construction rather than from unanalyzed side effects.

    Authors: We acknowledge that the current text does not provide the explicit assembly formula or any perturbation/stability argument. The revision will include (a) the precise formula that concatenates the QR output of the chosen subspace with the unchanged orthogonal vectors from the previous basis, and (b) a short stability argument showing that the deviation from the original preconditioner is bounded by a term proportional to the subspace fraction times machine epsilon in BFloat16. These additions will make the performance claims traceable to the construction. revision: yes

Circularity Check

0 steps flagged

New algorithmic reparametrization; no self-referential derivation or fitted predictions

full rationale

The paper proposes a subspace-based QR update that combines partial new basis vectors with unchanged ones to enable BFP16 storage and lower cost for Shampoo/SOAP preconditioners. This is presented as an explicit construction in the method section rather than a quantity derived from its own outputs, fitted parameters, or a self-citation chain. No equation reduces to a prior result by definition, no uniqueness theorem is imported from the authors' own prior work to force the choice, and the central claim (hybrid basis remains complete) is an assertion of the algorithm itself, not a prediction obtained by fitting. The derivation chain is therefore self-contained as an engineering modification; absence of a full algebraic proof for orthonormality is a correctness gap, not circularity.

Axiom & Free-Parameter Ledger

0 free parameters · 1 axioms · 0 invented entities

Abstract-only; no explicit free parameters, axioms, or invented entities are stated. The method implicitly assumes standard properties of QR decomposition and basis completeness after partial updates.

axioms (1)
  • domain assumption QR decomposition on a subspace produces vectors that can be combined with unchanged basis vectors to maintain a valid complete basis for the preconditioner.
    Stated in abstract as the mechanism that forms a complete basis; no proof or reference supplied.

pith-pipeline@v0.9.1-grok · 5744 in / 1199 out tokens · 21387 ms · 2026-06-29T22:14:41.216525+00:00 · methodology

0 comments
read the original abstract

Shampoo-based methods, such as KL-Shampoo and SOAP, have demonstrated strong performance in training neural networks and rely on QR decomposition. Because existing QR implementations require single-precision (FP32) arithmetic and remain computationally expensive, these methods become time- and memory-intensive when their preconditioning matrices are large. Moreover, using BFloat16 (BFP16) storage to reduce memory usage can degrade the performance of Shampoo-based methods. We propose a reparametrization of the preconditioner that supports BFP16 storage and forms a complete basis by combining updated basis vectors with unchanged ones. By updating only part of the basis through QR decomposition in a subspace, our approach reduces computational overhead while mitigating the performance degradation caused by BFP16 storage. Our approach applies broadly to Shampoo-based methods that employ QR decomposition, including KL-Shampoo, SOAP, and KL-SOAP. In particular, it improves the performance of SOAP and KL-SOAP under BFP16 storage, enabling KL-SOAP to match or exceed KL-Shampoo. Overall, our approach makes Shampoo-based methods more memory- and time-efficient.

Figures

Figures reproduced from arXiv: 2605.26327 by Alan Milligan, Felix Dangel, Simon Lacoste-Julien, Wu Lin, Zikun Xu.

Figure 1
Figure 1. Figure 1: Top left: Performance of full-basis and subspace KL-Shampoo (with BFP16 storage) on Llama3 (598M) trained on FineWeb 4.1B, eval￾uated using the hyperparameters tuned for full￾basis KL-Shampoo over 120 runs. Muon is also tuned independently over 120 runs. Top right: Computational cost for a weight matrix in Rd1×d2 with d1 = d2 = 6000, and in a subspace with dsub = Bd1, on an H100. Bottom: Test loss and per-… view at source ↗
Figure 2
Figure 2. Figure 2: Side-by-side comparison of KL-Shampoo (left) and its reparameterized variant (right). The reparameterized variant stores Pi := Q⊤ i SiQi instead of Si ; given equivalent initial state, the two variants produce equivalent iterates (see section 3.1), justifying the term reparameterization. Color code. Symbols in red mark where the two variants differ. Symbols in blue mark the rotated factors G˜′ i = Q⊤ i G′ … view at source ↗

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

26 extracted references · 3 canonical work pages

  1. [1]

    The polar express: Optimal matrix sign methods and their application to the M uon algorithm

    Noah Amsel, David Persson, Christopher Musco, and Robert M Gower. The polar express: Optimal matrix sign methods and their application to the M uon algorithm. In International Conference on Learning Representations (ICLR), 2026

    2026

  2. [2]

    Anderson, Z

    E. Anderson, Z. Bai, C. Bischof, S. Blackford, J. Demmel, J. Dongarra, J. Du Croz, A. Greenbaum, S. Hammarling, A. McKenney, and D. Sorensen. LAPACK Users' Guide . 1999

    1999

  3. [3]

    Scalable Second Order Optimization for Deep Learning

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

    work page Pith review arXiv 2002

  4. [4]

    Dynamic ordering for a parallel block- J acobi SVD algorithm

    Martin Be c ka, Gabriel Ok s a, and Marian Vajter s ic. Dynamic ordering for a parallel block- J acobi SVD algorithm. Parallel Computing, 2002

    2002

  5. [5]

    JAX : composable transformations of P ython+ N um P y programs, 2018

    James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, and Skye Wanderman-Milne. JAX : composable transformations of P ython+ N um P y programs, 2018

    2018

  6. [6]

    Benchmarking neural network training algorithms, 2023

    George E Dahl, Frank Schneider, Zachary Nado, Naman Agarwal, Chandramouli Shama Sastry, Philipp Hennig, Sourabh Medapati, Runa Eschenhagen, Priya Kasimbeg, Daniel Suo, et al. Benchmarking neural network training algorithms, 2023

    2023

  7. [7]

    Purifying S hampoo: Investigating S hampoo's heuristics by decomposing its preconditioner

    Runa Eschenhagen, Aaron Defazio, Tsung-Hsien Lee, Richard E Turner, and Hao-Jun Michael Shi. Purifying S hampoo: Investigating S hampoo's heuristics by decomposing its preconditioner. In Advances in Neural Information Processing Systems (NeurIPS), 2025

    2025

  8. [8]

    Eschenhagen, A

    Runa Eschenhagen, Anna Cai, Tsung-Hsien Lee, and Hao-Jun Michael Shi. Clarifying S hampoo: Adapting spectral descent to stochasticity and the parameter trajectory. arXiv preprint arXiv:2602.09314, 2026

    work page arXiv 2026

  9. [9]

    The cyclic J acobi method for computing the principal values of a complex matrix

    George Elmer Forsythe and Peter Henrici. The cyclic J acobi method for computing the principal values of a complex matrix. Transactions of the American Mathematical Society, 1960

    1960

  10. [10]

    Matrix computations

    Gene H Golub and Charles F Van Loan. Matrix computations. JHU press, 2013

    2013

  11. [11]

    S hampoo: Preconditioned stochastic tensor optimization

    Vineet Gupta, Tomer Koren, and Yoram Singer. S hampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning (ICML), 2018

    2018

  12. [12]

    Mixed precision algorithms in numerical linear algebra

    Nicholas J Higham and Theo Mary. Mixed precision algorithms in numerical linear algebra. Acta Numerica, 2022

    2022

  13. [13]

    MAGMA : Matrix algebra on GPU and multicore architectures, 2026

    Innovative Computing Laboratory . MAGMA : Matrix algebra on GPU and multicore architectures, 2026. URL https://developer.nvidia.com/magma. Accessed: 2026-05-06

    2026

  14. [14]

    NanoGPT (124M) in 3 minutes

    Keller Jordan. NanoGPT (124M) in 3 minutes . https://github.com/KellerJordan/modded-nanogpt, 2024. Accessed: 2025-06

    2024

  15. [15]

    Accelerating neural network training: An analysis of the AlgoPerf competition

    Priya Kasimbeg, Frank Schneider, Runa Eschenhagen, Juhan Bae, Chandramouli Shama Sastry, Mark Saroufim, Boyuan Fend, Less Wright, Edward Z Yang, Zachary Nado, et al. Accelerating neural network training: An analysis of the AlgoPerf competition. In The Thirteenth International Conference on Learning Representations, 2025

    2025

  16. [16]

    Lowe, Felix Dangel, Runa Eschenhagen, Zikun Xu, and Roger B

    Wu Lin, Scott C. Lowe, Felix Dangel, Runa Eschenhagen, Zikun Xu, and Roger B. Grosse. Understanding and improving S hampoo and SOAP via K ullback- L eibler minimization. In International Conference on Learning Representations (ICLR), 2026

    2026

  17. [17]

    Muon is scalable for llm training

    Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, Yanru Chen, Huabin Zheng, Yibo Liu, Shaowei Liu, Bohong Yin, Weiran He, Han Zhu, Yuzhi Wang, Jianzhou Wang, Mengnan Dong, Zheng Zhang, Yongsheng Kang, Hao Zhang, Xinran Xu, Yutao Zhang, Yuxin Wu, Xinyu Zhou, and Zhilin Yang. Muon is sca...

    2025

  18. [18]

    A new perspective on S hampoo's preconditioner

    Depen Morwani, Itai Shapira, Nikhil Vyas, Eran Malach, Sham Kakade, and Lucas Janson. A new perspective on S hampoo's preconditioner. In International Conference on Learning Representations (ICLR), 2024

    2024

  19. [19]

    cuSOLVER Library , 2026

    NVIDIA Corporation . cuSOLVER Library , 2026. URL https://docs.nvidia.com/cuda/cusolver/index.html. Accessed: 2026-05-06

    2026

  20. [20]

    PyTorch : An imperative style, high-performance deep learning library

    Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch : An imperative style, high-perf...

    2019

  21. [21]

    Build a large language model (from scratch)

    Sebastian Raschka. Build a large language model (from scratch). https://github.com/rasbt/LLMs-from-scratch, 2024. Accessed: 2025-10

    2024

  22. [22]

    Benchmarking optimizers for large lan- guage model pretraining.arXiv preprint arXiv:2509.01440,

    Andrei Semenov, Matteo Pagliardini, and Martin Jaggi. Benchmarking optimizers for large language model pretraining. arXiv preprint arXiv:2509.01440, 2025

    work page arXiv 2025

  23. [23]

    A distributed data-parallel PyTorch implementation of the distributed S hampoo optimizer for training neural networks at-scale, 2023

    Hao-Jun Michael Shi, Tsung-Hsien Lee, Shintaro Iwasaki, Jose Gallego-Posada, Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, and Michael Rabbat. A distributed data-parallel PyTorch implementation of the distributed S hampoo optimizer for training neural networks at-scale, 2023

    2023

  24. [24]

    The block J acobi method for computing the singular value decomposition

    Charles Van Loan. The block J acobi method for computing the singular value decomposition. Technical report, Cornell University, 1985

    1985

  25. [25]

    SOAP : Improving and stabilizing S hampoo using A dam for language modeling

    Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham M Kakade. SOAP : Improving and stabilizing S hampoo using A dam for language modeling. In International Conference on Learning Representations (ICLR), 2025

    2025

  26. [26]

    Convergence analysis of the parallel classical block J acobi method for the symmetric eigenvalue problem

    Yusaku Yamamoto, Zhang Lan, and Shuhei Kudo. Convergence analysis of the parallel classical block J acobi method for the symmetric eigenvalue problem. JSIAM Letters, 2014

    2014