How Bregman Divergences Shape Shampoo

Bing Liu; Wenjie Zhou; Chengcheng Zhao; Hongtao Zhang; Boao Kong; Felix Dangel; Wu Lin

cs.LGOct 6, 2026

How Bregman Divergences Shape Shampoo

Authors: Bing Liu, Wenjie Zhou, Chengcheng Zhao, Hongtao Zhang, Boao Kong, Felix Dangel, Wu Lin

Organizations: Zhejiang University · University of the Chinese Academy of Sciences · Peking University · Concordia University & Mila · University of Central Florida

Abstract

Understanding the principles behind Shampoo has recently guided the development of more effective neural network optimizers. These methods learn a preconditioner by optimizing the Frobenius or Kullback-Leibler (KL) divergence against the gradient second moment. In this work, we investigate how the choice of divergence shapes preconditioning, which remains unclear and blocks further improvements. To do so, we develop a unified Bregman divergence framework that connects all popular divergences, allowing us to study them jointly. Through empirical spectral analysis of gradient second moments, we examine how divergence choice shapes Kronecker approximation and interacts with finite-sample error in preconditioning. We find that some divergences can better compensate for finite-sample underestimation of the empirical second moment, helping explain the differing behavior of their corresponding Shampoo variants. We further validate this explanation through GPT-2 pretraining experiments. By connecting divergence choice to practical training behavior, we believe our framework provides principled guidance for understanding the foundations of, and further improving, Shampoo.

Figures & tables

Panel of Table 1 — Table 1: Overview of the divergence properties studied in this work. The exponent $q_{\mathrm{D}}$ indexes the trace-power family; spectral weighting and error asymmetry determine how approximation error is allocated, while few-sample error measures the stability of Kronecker factor estimation.

Figure 1: Divergences shape preconditioning and training. Left and middle: A GPT-2 checkpoint ( output_proj ) at step

2{,}500

, ordered by decreasing eigenvalue. Left: Empirical second-moment (

C_{0}

) eigenvalues become less reliable toward the tail, increasingly underestimating held-out scales ( Section 4.1 ). Middle (smoothed curves, asinh y-axis for better resolution near zero): KL stays closest to the held-out target (dashed zero line), followed by SQ, compensating

C_{0}

’s scale errors ( Section 5 ). Right: GPT-2 validation loss relative to KL with asinh y-axis; lower is better ( Section 6 ).

Figure 2: The family of scalar generators

f_{q}(x)

inducing trace-power Bregman divergences we study through one unifying lens. Shown with an asinh vertical scale.

Figure 3: Left: Normalized penalties

\smash{\frac{d_{\mathrm{D}}(\lambda,r\lambda)}{d_{\mathrm{D}}(1,r)}}=\lambda^{q_{\mathrm{D}}}

for

r\neq 1

. Smaller

q

places less emphasis on large eigenvalues; KL is scale-independent. Right:

d_{\mathrm{D}}(1,r)

versus

r

, with a log horizontal scale and an asinh vertical scale. F and VN favor underestimation over reciprocal overestimation, KL shows the opposite preference, and SQ is symmetric.

Figure 4: Left: Effective rank

r_{\mathrm{eff}}(R_{\star}^{q})

for eigenvalues

\rho_{j}=j^{-1.8}

. Right: Mean squared relative Frobenius error of the jointly fitted

S=L\otimes R

( Equation 100 ) on non-Kronecker gradients over

20

repeats, relative to each method’s

4096

-sample fit. Both results favor smaller

q_{\mathrm{D}}

, consistent with the theoretical prediction in Proposition 4 .

Figure 5: Large-sample Kronecker fits. Left: empirical and fitted eigenvalues. Right: overlap of the leading subspaces. The fitted spectra and eigenspaces follow the divergence-dependent preferences predicted in Section 3.2 .

Figure 6: Smoothed log held-out gain error for the output projection at AdamW step

2{,}500

. Closer to zero is better; with asinh vertical scale. Left: large-sample fit; middle:

40

-sample fits (

4

repeats); right: EMA (

16

repeats). Black solid: large-sample

\smash{C_{0}^{-\nicefrac{{1}}{{2}}}}

in all panels; dashed: the held-out target.

Optimizer	Method	Large-sample	40-sample	EMA
AdamW	F	$0.9759\,\pm\,0.4623$	$2.4945\,\pm\,0.8265$	$0.7645\,\pm\,0.3453$
	VN	$0.3328\,\pm\,0.1769$	$0.5549\,\pm\,0.2986$	$0.2897\,\pm\,0.1473$
	SQ	$0.1467\,\pm\,0.0573$	$0.2175\,\pm\,0.0893$	$0.1294\,\pm\,0.0412$
	KL	$\mathbf{0.1024}\,\pm\,0.0182$	$\mathbf{0.1561}\,\pm\,0.0297$	$\mathbf{0.1071}\,\pm\,0.0203$
Distributed Shampoo	F	$0.8801\,\pm\,0.3234$	$2.6922\,\pm\,0.5674$	$0.7008\,\pm\,0.2267$
	VN	$0.3353\,\pm\,0.1109$	$0.6207\,\pm\,0.1966$	$0.2806\,\pm\,0.0935$

Table 2: Absolute held-out gain error ( Equation 20 ,

\downarrow

). KL has the lowest mean error, followed by SQ. Each experiment contributes the median

|e_{D,i}|

over all directions and repeats; entries show mean

\pm

sample SD across nine checkpoint-projection pairs per training optimizer.

	AdamW		Distributed Shampoo
Divergence	Few-sample	EMA	Few-sample	EMA
F	$1.668\,\pm\,0.601$	$0.173\,\pm\,0.070$	$1.967\,\pm\,0.430$	$0.154\,\pm\,0.017$
VN	$0.259\,\pm\,0.141$	$0.082\,\pm\,0.032$	$0.339\,\pm\,0.121$	$0.087\,\pm\,0.017$
SQ	$0.089\,\pm\,0.042$	$0.063\,\pm\,0.020$	$0.077\,\pm\,0.021$	$0.048\,\pm\,0.008$
KL	$\mathbf{0.066}\,\pm\,0.020$	$\mathbf{0.044}\,\pm\,0.011$	$\mathbf{0.060}\,\pm\,0.015$	$\mathbf{0.040}\,\pm\,0.009$

Table 3: Recovery error ( Equation 21 ,

\downarrow

). KL has the lowest mean error, followed by SQ. Each experiment contributes the median

|\gamma_{D}^{\mathrm{rec}}|

over all reference directions and repeats; entries show mean

\pm

sample SD across nine checkpoint-projection pairs per training optimizer.

Method / width	256	512	768	1024	1536
F	3.8281	3.5082	3.3993	3.3528	3.3111
VN	3.7573	3.4377	3.2876	3.1979	3.0972
SQ	3.7450	3.4173	3.2661	3.1824	3.0817
KL	3.7475	3.4180	3.2660	3.1769	3.0758

Table 4: Final GPT-2 validation loss on FineWeb after

5,100

updates (

\downarrow

); one run per configuration. SQ and KL outperform F and VN at every width. Bold marks the lowest loss in each column.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Population objective	Factor updates discussed here
PSGD-Kron	SQ divergence	Original relative-gradient steps; other geometries are available
SQ-Shampoo	SQ divergence	Spectral EMA of SQ statistics
KL-Shampoo	KL divergence	Spectral EMA of KL statistics

Appendix

Table 5: Objectives and update schemes discussed here. Objectives are expressed in

S

, up to positive scale and additive constants; the PSGD row refers to gradient whitening.

Method	$\Psi_{L}(G;L,R)$	$\Psi_{R}(G;L,R)$
F-Shampoo	$GRG^{\top}/\Tr(R^{2})$	$G^{\top}LG/\Tr(L^{2})$
VN-Shampoo	$GG^{\top}/\Tr(R)$	$G^{\top}G/\Tr(L)$
SQ-Shampoo	$GR^{-\nicefrac{{1}}{{2}}}G^{\top}/\Tr(R^{\nicefrac{{1}}{{2}}})$	$G^{\top}L^{-\nicefrac{{1}}{{2}}}G/\Tr(L^{\nicefrac{{1}}{{2}}})$
KL-Shampoo	$GR^{-1}G^{\top}/n$	$G^{\top}L^{-1}G/m$

Appendix

Table 6: Instantaneous factor statistics for the Shampoo variants.

Setting	$m\times n$	$N$	Additional parameters
Smooth long-tailed	$16\times 24$	$4096$	scale $=10^{-1}$ , $\sigma=0.02$

Appendix

Table 7: Settings for the synthetic experiments.

Figure 7: Smooth synthetic setting. From left to right: empirical and fitted spectra, cumulative spectral error, and leading-subspace overlap. Finite-sample stability is shown in Figure 4 (right).

Figure 8: Quadratic training loss averaged over

3

seeds. Left: four methods. Middle and right: replacing KL gains from step

900

. F gains raise loss in either region; SQ gains leave it close to KL.

Setting	Specification
Gradient pools	Four splits per experiment; $20{,}480$ FP32 mini-batch gradients per split, assigned by round-robin interleaving
Gradient collection	Batch size $64$ ; sequence length $512$
Large-sample fit	All $20{,}480$ gradients in split $0$
Few-sample fit	$40$ gradients from split $0$ ; four repeats
Fixed-point solver	Relaxation $\alpha=0.3$ ; tolerance $10^{-8}$ ; at most $1{,}000$ steps
Online EMA	$16$ disjoint chronological repeats from split $0$ , each with $512$ gradients; $\beta_{2}=0.95$ ; eigendecomposition every five steps

Appendix

Table 8: Shared settings for real-gradient diagnostics.

AdamW checkpoints
Metric	Top (0–10%)	Middle (10–40%)	Lower (40–75%)	Tail (75–100%)
$\|\xi\|$ ( $\downarrow$ )	$0.027\,\pm\,0.016$	$0.124\,\pm\,0.037$	$0.431\,\pm\,0.030$	$1.009\,\pm\,0.120$
Signed $\xi$	$0.010\,\pm\,0.030$	$0.124\,\pm\,0.037$	$0.431\,\pm\,0.030$	$1.009\,\pm\,0.120$
Cross-fit $\xi^{\mathrm{cf}}$	$-0.008\,\pm\,0.004$	$-0.006\,\pm\,0.004$	$-0.006\,\pm\,0.005$	$-0.007\,\pm\,0.006$
Distributed Shampoo checkpoints
Metric	Top (0–10%)	Middle (10–40%)	Lower (40–75%)	Tail (75–100%)

Appendix

Table 9: Held-out scale error across all four eigenvalue regions. Zero indicates agreement; positive

\xi

indicates underestimation, and larger

|\xi|

indicates greater mismatch. Entries show mean

\pm

sample SD across nine checkpoint-projection pairs per training optimizer, using each’s regional median.

Panel of Table 10 — Table 10: Held-out scale error. Rows: steps 100, 2500, 4500. Bands show split-comparison IQRs.

Panel of Table 11 — Table 11: Held-out gain error on AdamW checkpoints at step 100. Rows: large-sample fit, 40-sample fit, EMA with eigendecomposition every five steps.

Panel of Table 12 — Table 12: Held-out gain error on AdamW checkpoints at step $2500$ . Rows: large-sample fit, 40-sample fit, EMA with eigendecomposition every five steps.

Panel of Table 13 — Table 13: Held-out gain error on AdamW checkpoints at step 4500. Rows: large-sample fit, 40-sample fit, EMA with eigendecomposition every five steps.

Panel of Table 14 — Table 14: Large-sample spectra. Rows: steps 100, 2500, 4500.

Panel of Table 15 — Table 15: Large-sample fit recovery on AdamW checkpoints at step $2500$ . Rows: F, VN, SQ, KL. Solid: 40-sample fits; dashed: EMA with eigendecomposition every five steps.

Panel of Table 16 — Table 16: Distributed Shampoo diagnostics at step $2500$ . Top: held-out scale error and cross-fit controls. Remaining rows: held-out gain error for large-sample fits, $40$ -sample fits (four repeats), and EMA ( $512$ gradients, $16$ repeats, eigendecomposition every five steps).

Width	Params (M)	LR	F	VN	SQ	KL
256	22.3	0.0024	3.8281	3.7573	3.7450	3.7475
512	63.5	0.0024	3.5082	3.4377	3.4173	3.4180
768	123.6	0.0018	3.3993	3.2876	3.2661	3.2660
1024	202.5	0.0012	3.3528	3.1979	3.1824	3.1769
1536	417.0	0.0008	3.3111	3.0972	3.0817	3.0758

Appendix

Table 17: Final validation loss after

5,100

steps; one run per configuration.

Figure 9: Validation loss from step

2,000

at all five widths.

Explore similar work

Understanding and Improving Shampoo and SOAP via Kullback-Leibler Minimization
Sep 3, 2025Wu Lin, Scott C. Lowe, Felix Dangel +3Deep Learning Optimization Neural Network Optimization
1Vector Institute, Toronto, Canada · University of Cambridge, Cambridge, United Kingdom · 3Microsoft, United States +1
Rethinking Bregman Divergences in Kronecker-Factored Optimizers
May 30, 2026Bing Liu, Wenjie Zhou, Chengcheng ZhaoSecond-Order Optimization Adaptive Gradient Methods
College of Control Science and Engineering, Zhejiang University · State Key Laboratory of AI Safety, Institute of Computing Technology, Chinese Academy of Sciences
DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers
Feb 2, 2026Ionut-Vlad Modoranu, Philip Zmushko, Erik Schultheis +2Deep Learning Optimization GPU Acceleration
Institute of Science and Technology Austria · Lancaster University · Red Hat AI

Sep 3, 2025, stat.MLJ/K move · Enter open · S save

Understanding and Improving Shampoo and SOAP via Kullback-Leibler Minimization

Wu Lin, Scott C. Lowe, Felix Dangel, Runa Eschenhagen+2

1Vector Institute, Toronto, Canada · University of Cambridge, Cambridge, United Kingdom · 3Microsoft, United States+1

Deep Learning Optimization Neural Network Optimization Second-Order Optimization Adaptive Gradient Methods

Open paper AskPDF arXiv

Shampoo and its efficient variant, SOAP, employ structured second-moment estimations and have shown strong performance for training neural networks (NNs). In practice, however, Shampoo typically requires step-size grafting with Adam to be competitive, and SOAP mitigates this by applying Adam in Shampoo's eigenbasis -- at the cost of additional memory overhead from Adam in both methods. Prior analyses have largely relied on the Frobenius norm to motivate these estimation schemes. We instead recast their estimation procedures as covariance estimation under Kullback-Leibler (KL) divergence minimization, revealing a previously overlooked theoretical limitation and motivating principled redesigns. Building on this perspective, we develop

\textbf{KL-Shampoo}

and

\textbf{KL-SOAP}

, practical schemes that match or exceed the performance of Shampoo and SOAP in NN pre-training while achieving SOAP-level per-iteration runtime. Notably, KL-Shampoo does not rely on Adam to attain competitive performance, eliminating the memory overhead introduced by Adam. Across our experiments, KL-Shampoo consistently outperforms SOAP, Shampoo, and even KL-SOAP, establishing the KL-based approach as a promising foundation for designing structured methods in NN optimization. An implementation of KL-Shampoo/KL-SOAP is available at https://github.com/yorkerlin/KL-Methods

Abstract

Explore similar work

Understanding and Improving Shampoo and SOAP via Kullback-Leibler Minimization

Rethinking Bregman Divergences in Kronecker-Factored Optimizers

DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers