pith. sign in

arxiv: 2606.00542 · v2 · pith:3UJDWOE6new · submitted 2026-05-30 · 💻 cs.LG

Rethinking Bregman Divergences in Kronecker-Factored Optimizers

Bing Liu , Wenjie Zhou , Chengcheng Zhao This is my paper

Pith reviewed 2026-06-28 18:48 UTC · model grok-4.3

classification 💻 cs.LG
keywords Bregman divergenceKronecker factorizationcovariance approximationsecond-order optimizationShampoo optimizerpreconditionerspectral analysissubspace optimization
0
0 comments X

The pith

Different Bregman divergences distribute Kronecker approximation errors differently across the covariance spectrum.

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

Shampoo-style optimizers approximate gradient covariances with Kronecker factors. When the covariance is not exactly Kronecker-structured, approximation error must go somewhere. The paper shows that the Frobenius, von Neumann, and LogDet divergences each spread this error across the eigenvalue spectrum in distinct patterns. The Kronecker factors themselves are produced by residuals that are weighted according to the chosen divergence. These spectral preferences, together with the observation that the leading eigenspace aligns better with the Hessian than the noisy tail, motivate a subspace-aware optimizer that applies eigenvalue preconditioning only on top and isotropic acceleration below.

Core claim

When a covariance matrix is projected onto Kronecker structure under a Bregman divergence, the unavoidable approximation error is allocated differently across the spectrum by the Frobenius, von Neumann, or LogDet divergence. The resulting Kronecker factors are governed by divergence-weighted residuals rather than the raw error. This produces preconditioners with different emphases on reliable versus unreliable parts of the spectrum. The top covariance eigenspace aligns substantially better with the Hessian while the tail is noisier, which directly motivates splitting the space into an eigenvalue-preconditioned top subspace and an isotropically accelerated bottom subspace.

What carries the argument

Kronecker projection under Bregman matrix divergences, with factors realized through divergence-weighted residuals.

If this is right

  • The approximation error is allocated according to divergence-specific weighting rather than uniformly.
  • Each divergence produces Kronecker factors whose spectral behavior follows from its own weighted residuals.
  • Eigenvalue-based preconditioning is reliable only in the top subspace where alignment with the Hessian holds.
  • An adaptive isotropic acceleration term is appropriate for the tail where eigenvalue estimates are unreliable.

Where Pith is reading between the lines

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

  • The weighted-residual perspective may extend to other structured low-rank or factored approximations beyond Kronecker products.
  • A per-subspace divergence choice could be tested to match the reliability profile of each block.
  • The alignment observation could be checked on additional architectures to see whether the subspace split generalizes.

Load-bearing premise

The top eigenvectors of the covariance matrix align substantially better with the Hessian than those in the tail of the spectrum.

What would settle it

Compute the principal angles or overlap between the top covariance eigenvectors and the corresponding Hessian eigenvectors on a trained model; if alignment is no stronger in the top block than in the tail, the justification for the proposed subspace split disappears.

read the original abstract

Shampoo-style optimizers approximate gradient covariance matrices using Kronecker-factored structures. Recent work~\cite{lin2026understanding} showed that such approximations can be viewed as projections under Bregman matrix divergences, leading to different Kronecker-factored preconditioners. However, it remains unclear what role the choice of divergence plays when the covariance is not exactly Kronecker-factored. We study this question through the spectrum of the covariance matrix. We show that Frobenius, von Neumann, and LogDet divergences distribute the unavoidable Kronecker approximation error differently across the covariance spectrum. We further show that their Kronecker factors are governed by divergence-weighted residuals rather than the raw approximation error, explaining how these spectral preferences are realized in the resulting preconditioners. Empirically, we observe that the top covariance eigenspace is substantially better aligned with the Hessian matrix, while the tail spectrum is much noisier and unreliable. Motivated by these findings, we propose a subspace-aware Kronecker optimizer that applies eigenvalue-based preconditioning in the top subspace and uses an adaptive isotropic acceleration constant in the bottom subspace.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

2 major / 2 minor

Summary. The paper analyzes the effect of Bregman matrix divergences (Frobenius, von Neumann, LogDet) on how Kronecker approximation error is distributed across the spectrum of a covariance matrix that is not exactly Kronecker-factored. It shows that the resulting Kronecker factors are governed by divergence-weighted residuals rather than raw error. Empirically, the top covariance eigenspace is observed to align better with the Hessian than the noisy tail; this motivates a subspace-aware optimizer that applies eigenvalue-based preconditioning to the top block and adaptive isotropic acceleration to the bottom block.

Significance. If the spectral weighting analysis is correct, the work supplies a principled explanation for why different divergences produce distinct preconditioners and could inform divergence selection in practice. The proposed subspace split is a direct consequence of the claimed alignment difference; its practical value hinges on whether that difference is robustly quantified and whether the resulting optimizer yields measurable gains over standard Shampoo variants.

major comments (2)
  1. [Abstract (final paragraph)] Abstract (final paragraph) and the empirical motivation section: the claim that 'the top covariance eigenspace is substantially better aligned with the Hessian matrix, while the tail spectrum is much noisier and unreliable' is load-bearing for the subspace split in the proposed optimizer, yet the manuscript provides no quantitative metrics (principal angles, trace overlap, cosine similarity), no specific models or datasets, and no controls for finite-sample bias or architecture dependence. Without these, the justification for eigenvalue preconditioning in the top block versus isotropic acceleration in the tail does not follow from the spectral analysis.
  2. [§3] §3 (theoretical development of weighted residuals): the statement that Kronecker factors are 'governed by divergence-weighted residuals rather than the raw approximation error' is central, but the derivation must explicitly show how the weighting arises from the Bregman projection for each divergence and confirm that the weighting is not an artifact of the particular matrix factorization chosen.
minor comments (2)
  1. [§2] Notation for the three divergences and their associated residuals should be introduced once with a single table or equation block rather than redefined inline in multiple sections.
  2. [§5] The experimental section should report the precise hyper-parameter settings used for the baseline Shampoo variants and the new subspace-aware method so that the claimed improvements can be reproduced.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for the constructive comments. We address each major point below and indicate where revisions will be made to strengthen the manuscript.

read point-by-point responses

  1. Referee: [Abstract (final paragraph)] Abstract (final paragraph) and the empirical motivation section: the claim that 'the top covariance eigenspace is substantially better aligned with the Hessian matrix, while the tail spectrum is much noisier and unreliable' is load-bearing for the subspace split in the proposed optimizer, yet the manuscript provides no quantitative metrics (principal angles, trace overlap, cosine similarity), no specific models or datasets, and no controls for finite-sample bias or architecture dependence. Without these, the justification for eigenvalue preconditioning in the top block versus isotropic acceleration in the tail does not follow from the spectral analysis.

    Authors: We agree that the alignment claim requires quantitative support to justify the subspace split. In the revised manuscript we add a dedicated empirical subsection reporting principal angles, cosine similarities, and trace overlaps between the top covariance eigenvectors and the Hessian. Results are shown for ResNet-18 trained on CIFAR-10 and a 6-layer transformer on WikiText-2, with controls varying batch size to assess finite-sample effects and architecture dependence. These additions directly substantiate the motivation for eigenvalue preconditioning on the top block and isotropic acceleration on the tail. revision: yes

  2. Referee: [§3] §3 (theoretical development of weighted residuals): the statement that Kronecker factors are 'governed by divergence-weighted residuals rather than the raw approximation error' is central, but the derivation must explicitly show how the weighting arises from the Bregman projection for each divergence and confirm that the weighting is not an artifact of the particular matrix factorization chosen.

    Authors: We accept that the weighting derivation needs to be shown explicitly. The revised Section 3 now contains the full stationarity conditions for the Bregman projections under Frobenius, von Neumann, and LogDet divergences, deriving the residual weighting factors directly from each optimality equation. We further add a short argument establishing that the weighting depends only on the chosen divergence and the Kronecker constraint, independent of any particular factorization algorithm, by repeating the projection under an alternative parameterization of the factors. revision: yes

Circularity Check

0 steps flagged

No significant circularity; theoretical claims derive from divergence definitions and matrix spectrum properties

full rationale

The paper's core claims—that Frobenius, von Neumann, and LogDet divergences distribute Kronecker approximation error differently across the covariance spectrum and that factors are governed by divergence-weighted residuals—are presented as following from the definitions of the divergences and the spectral decomposition of the covariance matrix. No equations or steps are shown that reduce these results to fitted parameters, self-citations, or inputs by construction. The empirical alignment observation between top eigenspace and Hessian is stated as an independent motivation for the subspace optimizer rather than a load-bearing derivation step. The cited prior work on Bregman projections is external and does not create a self-referential chain. The derivation chain remains self-contained against external matrix properties.

Axiom & Free-Parameter Ledger

0 free parameters · 1 axioms · 1 invented entities

The central claims rest on standard properties of Bregman divergences for matrix projections (from the cited work) and on an empirical observation about Hessian alignment that is not independently derived in the abstract.

axioms (1)
  • domain assumption Bregman matrix divergences define valid projections onto the set of Kronecker-factored matrices
    Invoked via the recent work cited in the abstract to frame the approximations.
invented entities (1)
  • subspace-aware Kronecker optimizer no independent evidence
    purpose: Applies eigenvalue-based preconditioning only in top subspace and adaptive isotropic acceleration in bottom subspace
    Proposed based on the reported empirical spectral alignment findings.

pith-pipeline@v0.9.1-grok · 5720 in / 1352 out tokens · 25733 ms · 2026-06-28T18:48:53.252555+00:00 · methodology

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 · 5 canonical work pages · 1 internal anchor

  1. [1]

    ASGO : Adaptive structured gradient optimization

    Kang An, Yuxing Liu, Rui Pan, Yi Ren, Shiqian Ma, Donald Goldfarb, and Tong Zhang. ASGO : Adaptive structured gradient optimization. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025

    2025

  2. [2]

    Alex Damian, Eshaan Nichani, and Jason D. Lee. Self-stabilization: The implicit bias of gradient descent at the edge of stability. In The Eleventh International Conference on Learning Representations, 2023

    2023

  3. [3]

    Turner, and Hao-Jun Michael Shi

    Runa Eschenhagen, Aaron Defazio, Tsung-Hsien Lee, Richard E. Turner, and Hao-Jun Michael Shi. Purifying shampoo: Investigating shampoo's heuristics by decomposing its preconditioner. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025

    2025

  4. [4]

    Shampoo: Preconditioned stochastic tensor optimization

    Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, pages 1842--1850. PMLR, 2018

    2018

  5. [5]

    Gradient Descent Happens in a Tiny Subspace

    Guy Gur-Ari, Daniel A Roberts, and Ethan Dyer. Gradient descent happens in a tiny subspace. arXiv preprint arXiv:1812.04754, 2018

    work page internal anchor Pith review Pith/arXiv arXiv 2018

  6. [6]

    Muon: An optimizer for hidden layers in neural networks, 2024

    Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https://kellerjordan.github.io/posts/muon/

    2024

  7. [7]

    Yang, Zachary Nado, Sourabh Medapati, Philipp Hennig, Michael Rabbat, and George E

    Priya Kasimbeg, Frank Schneider, Runa Eschenhagen, Juhan Bae, Chandramouli Shama Sastry, Mark Saroufim, BOYUAN FENG, Less Wright, Edward Z. Yang, Zachary Nado, Sourabh Medapati, Philipp Hennig, Michael Rabbat, and George E. Dahl. Accelerating neural network training: An analysis of the algoperf competition. In The Thirteenth International Conference on Le...

    2025

  8. [8]

    arXiv preprint arXiv:2510.16981 , year=

    Ahmed Khaled, Kaan Ozkara, Tao Yu, Mingyi Hong, and Youngsuk Park. Muonbp: Faster muon via block-periodic orthogonalization. arXiv preprint arXiv:2510.16981, 2025

    work page arXiv 2025

  9. [9]

    Adam: A method for stochastic optimization

    Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015

    2015

  10. [10]

    Limitations of the empirical fisher approximation for natural gradient descent

    Frederik Kunstner, Philipp Hennig, and Lukas Balles. Limitations of the empirical fisher approximation for natural gradient descent. Advances in Neural Information Processing Systems, 32, 2019

    2019

  11. [11]

    Lowe, Felix Dangel, Runa Eschenhagen, Zikun Xu, and Roger Baker Grosse

    Wu Lin, Scott C. Lowe, Felix Dangel, Runa Eschenhagen, Zikun Xu, and Roger Baker Grosse. Understanding and improving shampoo and SOAP via kullback-leibler minimization. In The Fourteenth International Conference on Learning Representations, 2026

    2026

  12. [12]

    Decoupled weight decay regularization

    Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019

    2019

  13. [13]

    Optimizing neural networks with kronecker-factored approximate curvature

    James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In International Conference on Machine Learning, pages 2408--2417. PMLR, 2015

    2015

  14. [14]

    A new perspective on shampoo's preconditioner

    Depen Morwani, Itai Shapira, Nikhil Vyas, Sham M Kakade, Lucas Janson, et al. A new perspective on shampoo's preconditioner. In The Thirteenth International Conference on Learning Representations, 2025

    2025

  15. [15]

    Does sgd really happen in tiny subspaces? In 13th International Conference on Learning Representations

    Minhak Song, Kwangjun Ahn, and Chulhee Yun. Does sgd really happen in tiny subspaces? In 13th International Conference on Learning Representations. International Conference on Learning Representations (ICLR), 2025

    2025

  16. [16]

    On the interplay between noise and curvature and its effect on optimization and generalization

    Valentin Thomas, Fabian Pedregosa, Bart Merri \"e nboer, Pierre-Antoine Manzagol, Yoshua Bengio, and Nicolas Le Roux. On the interplay between noise and curvature and its effect on optimization and generalization. In International Conference on Artificial Intelligence and Statistics, pages 3503--3513. PMLR, 2020

    2020

  17. [17]

    Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham M. Kakade. SOAP : Improving and stabilizing shampoo using adam for language modeling. In The Thirteenth International Conference on Learning Representations, 2025

    2025

  18. [18]

    The sharpness disparity principle in transformers for accelerating language model pre-training

    Jinbo Wang, Mingze Wang, Zhanpeng Zhou, Junchi Yan, Lei Wu, et al. The sharpness disparity principle in transformers for accelerating language model pre-training. In Forty-second International Conference on Machine Learning, 2026

    2026

  19. [19]

    Understanding warmup-stable-decay learning rates: A river valley loss landscape view

    Kaiyue Wen, Zhiyuan Li, Jason S Wang, David Leo Wright Hall, Percy Liang, and Tengyu Ma. Understanding warmup-stable-decay learning rates: A river valley loss landscape view. In The Thirteenth International Conference on Learning Representations, 2026

    2026

  20. [20]

    The alignment property of sgd noise and how it helps select flat minima: A stability analysis

    Lei Wu, Mingze Wang, and Weijie Su. The alignment property of sgd noise and how it helps select flat minima: A stability analysis. Advances in Neural Information Processing Systems, 35: 0 4680--4693, 2022

    2022

  21. [21]

    Reddi, Sanjiv Kumar, and Zhiyuan Li

    Shuo Xie, Tianhao Wang, Sashank J. Reddi, Sanjiv Kumar, and Zhiyuan Li. Structured preconditioners in adaptive optimization: A unified analysis. In Forty-second International Conference on Machine Learning, 2025

    2025

  22. [22]

    arXiv preprint arXiv:2601.08393 , year=

    Tian Xie, Haoming Luo, Haoyu Tang, Yiwen Hu, Jason Klein Liu, Qingnan Ren, Yang Wang, Wayne Xin Zhao, Rui Yan, Bing Su, et al. Controlled llm training on spectral sphere. arXiv preprint arXiv:2601.08393, 2026

    work page arXiv 2026

  23. [23]

    Fismo: Fisher-structured momentum-orthogonalized optimizer

    Chenrui Xu, Wenjing Yan, and Ying-Jun Angela Zhang. Fismo: Fisher-structured momentum-orthogonalized optimizer. arXiv preprint arXiv:2601.21750, 2026

    work page arXiv 2026

  24. [24]

    Bsfa: Leveraging the subspace dichotomy to accelerate neural network training

    Wenjie Zhou, Bohan Wang, Wei Chen, and Xueqi Cheng. Bsfa: Leveraging the subspace dichotomy to accelerate neural network training. In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, pages 18845--18860, 2025

    2025

  25. [25]

    Accelerating llm pre-training through flat-direction dynamics enhancement.arXiv preprint arXiv:2602.22681, 2026

    Shuchen Zhu, Rizhen Hu, Mingze Wang, Mou Sun, Xue Wang, Kun Yuan, and Zaiwen Wen. Accelerating llm pre-training through flat-direction dynamics enhancement. arXiv preprint arXiv:2602.22681, 2026

    work page arXiv 2026

  26. [26]

    The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects

    Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. In International Conference on Machine Learning, pages 7654--7663. PMLR, 2019

    2019