Rethinking Bregman Divergences in Kronecker-Factored Optimizers
Mirrored from arXiv — Machine Learning for archival readability. Support the source by reading on the original site.
arXiv:2606.00542v1 Announce Type: new
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.
More from arXiv — Machine Learning
-
BitsMoE: Efficient Spectral Energy-Guided Bit Allocation for MoE LLM Quantization
Jun 2
-
DAStatFormer: A Hybrid Multibranch Transformer with Statistical Feature Integration for DAS-Based Pattern Recognitions
Jun 2
-
Hoeffding Concept Bottleneck Models with Applications to Overhead Images
Jun 2
-
From Demonstrations to Rewards: Test-Time Prompt Optimization for VLM Reward Models
Jun 2
Discussion (0)
Sign in to join the discussion. Free account, 30 seconds — email code or GitHub.
Sign in →No comments yet. Sign in and be the first to say something.