PaperlayerView exact revision on arXiv

Source-linked AI summary

How Bregman Divergences Shape Shampoo

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

arXiv:2610.08534v1cs.LG

TL;DR

The relationship between Shampoo’s divergence objective and its differing training behavior is not yet understood. The paper unifies divergence-based variants with a trace-power Bregman framework and analyzes finite-sample effects. It finds that SQ and KL better compensate for empirical second-moment underestimation and achieve lower GPT-2 validation losses than F and VN.

  • Problem

    How divergence objectives produce different Shampoo training performance remains unclear, despite the importance of understanding Shampoo’s principles for optimizer development.

    Source: §1 · p. 1

  • Method

    A trace-power Bregman family unifies divergence-based Shampoo variants and is used to study Kronecker approximation and finite-sample effects in preconditioning.

    Source: §1 · pp. 1–2

  • Results

    SQ and KL better compensate for empirical second-moment underestimation and achieve lower GPT-2 validation losses than F and VN, with KL’s advantage growing at larger widths.

    Source: §1 · p. 2, §7 · p. 10

  • Takeaways & Limitations

    Held-out gain and recovery errors provide practical metrics for evaluating new preconditioners.

    Source: §7 · p. 10

    The factor-stability analysis compares divergences under an exact Kronecker Gaussian model.

    Source: §4.2 · p. 6

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper asks why divergence choices yield different Shampoo training behavior, separating Kronecker approximation error from finite-sample errors in the empirical moment and factor estimates. It develops a unified framework and studies how these effects shape preconditioning and GPT-2 training.

1 INTRODUCTION

  • The analysis separates Kronecker approximation mismatch, finite-sample errors in moments and factor estimates, and their interaction during preconditioning.The introduction frames these as effects often conflated in practice.Source: §1 · pp. 1–2
  • The trace-power Bregman family unifies divergence-based Shampoo variants, and its square-root member recovers PSGD’s gradient-whitening objective as SQ-Shampoo.The family exponent determines which spectral regions are prioritized and whether over- or underestimation is penalized more.Source: §1 · p. 2
  • On GPT-2, empirical gradient second moments severely underestimate tail eigenvalues even when gradients outnumber dimensions.The same exponent qD controls finite-sample Kronecker-factor error, with KL achieving the lowest error.Source: §1 · p. 2
  • SQ and KL achieve lower GPT-2 validation losses than F and VN, with KL’s advantage growing at larger widths.Figure 1’s right panel plots GPT-2 validation loss relative to KL; lower is better.Source: §1 · p. 2, Figure 1 · p. 2
  • The framework connects divergence-dependent approximation and finite-sample behavior to practical criteria for designing Kronecker-factored optimizers.It aims to explain why divergence variants differ in training performance.Source: §1 · pp. 1–2

2 BACKGROUND

The paper presents Shampoo and PSGD as approaches to preconditioning gradient second moments, then describes Shampoo variants through Bregman divergence minimization over Kronecker factors. Their divergence choices have different training performance, whose relationship to the objective remains unexplained.

2 BACKGROUND

  • Shampoo variants approximate the gradient second moment with a Kronecker product and apply inverse-root preconditioning.Modern variants use p = 1/2, while original Shampoo uses p = 1/4.Source: §2 · p. 3
  • Shampoo maintains averages of matrix-gradient products, while PSGD learns a preconditioner that acts on the stochastic gradient.PSGD-Kron restricts its preconditioner to a Kronecker product and learns its factors through relative-gradient updates.Source: §2 · p. 3
  • A Bregman matrix divergence is formed from a strictly convex generator, and choices of generator yield Frobenius, von Neumann, or KL divergence.The KL divergence is also known as LogDet divergence.Source: §2 · p. 3
  • Bregman divergence minimization yields F-, VN-, and KL-Shampoo, but the relationship between divergence choice and empirical training performance remains unexplained.The paper notes that the variants perform differently when training neural networks.Source: §2 · p. 3

3 KRONECKER APPROXIMATION ERROR ALLOCATION

A trace-power Bregman family unifies Shampoo variants and PSGD-Kron, while its exponent determines how spectral approximation errors are weighted. Smaller exponents emphasize the spectrum’s tail and penalize underestimation more strongly.

3.1 A TRACE-POWER FAMILY OF KRONECKER PRECONDITIONERS

  • F-, VN-, and KL-Shampoo, plus PSGD-Kron’s gradient-whitening objective, belong to a trace-power Bregman family; q=1/2 yields SQ-Shampoo.Figure 2 shows the scalar generators in this family.Source: §3.1 · p. 4, Figure 2 · p. 4

3.2 SPECTRAL ERROR ALLOCATION UNDER LIMITED KRONECKER CAPACITY

  • A single Kronecker product generally cannot match the target exactly, so divergence choice determines how the remaining approximation errors are penalized.The analysis separates eigenvalue and eigenspace mismatch costs.Source: §3.2 · p. 4
  • Scaling both matrices by r multiplies eigenspace-mismatch cost by r^q, so larger q weights misalignment more at larger eigenvalue scales.The exponent also controls how strongly eigenvalue-mismatch costs depend on scale.Source: §3.2 · pp. 4–5
  • F and VN favor underestimation over reciprocal overestimation, KL shows the opposite preference, and SQ is symmetric.Smaller q places less emphasis on large eigenvalues.Source: §3.2 · p. 5
  • Smaller q_D gives more weight to the tail spectrum and penalizes underestimation more strongly, encouraging larger fitted eigenvalues.Those larger eigenvalues can help avoid excessive inverse-root amplification.Source: §3.2 · p. 5

4 FINITE-SAMPLE EFFECTS ON KRONECKER APPROXIMATION

The section examines two sources of finite-sample error: an unreliable empirical second-moment target and noisy estimates of Kronecker factors. In the studied settings, tail eigenvalues are underestimated, while smaller divergence exponents yield more stable factor estimates.

4.1 RELIABILITY OF THE EMPIRICAL SECOND-MOMENT TARGET

  • The reliability check estimates eigenvectors on one data split and compares the scales along those directions across independent splits.It assesses scale reliability, not whether those directions remain eigenvectors on held-out splits.Source: §4.1 · pp. 5–6
  • Held-out scale errors rise toward smaller eigenvalues, indicating increasingly severe tail underestimation across checkpoints and projections.Cross-fit errors stay near zero, supporting sampling noise as the source of the scale differences.Source: §4.1 · p. 6

4 FINITE-SAMPLE EFFECTS ON KRONECKER APPROXIMATION

  • Even with many gradients, empirical second moments can underestimate the spectral tail; smaller q_D improves factor estimation, with KL achieving the lowest expected relative error.The factor-estimation trend also holds empirically beyond the assumptions of the theoretical model.Source: §4.1 · p. 6, §4.2 · p. 7

4.2 FINITE-SAMPLE STABILITY OF KRONECKER FACTORS

  • The theoretical comparison fixes one factor under an exact Kronecker Gaussian model, whereas the joint-fit test uses synthetic non-Kronecker gradients.The real-gradient analysis examines recovery of large-sample inverse-root gains without Gaussian or Kronecker assumptions.Source: §4.2 · p. 7
  • With the opposite factor fixed under an exact Kronecker Gaussian model, smaller q_D yields more accurate factor estimates; KL has the largest effective rank and smallest error.On synthetic long-tailed non-Kronecker gradients, joint-fit mean errors follow the predicted ordering at every tested sample size.Source: §4.2 · p. 7, Figure 4 · p. 7

5 FROM KRONECKER APPROXIMATIONS TO PRECONDITIONING

Divergence choice shapes Kronecker fits and their preconditioning gains, while finite-sample effects determine how reliably those gains are recovered. KL, closely followed by SQ, best compensates for tail underestimation and provides more reliable gains from short gradient histories.

5.1 LARGE-SAMPLE KRONECKER FITS

  • SQ and KL produce larger fitted tail eigenvalues than F and VN, while F best matches the leading eigenvector and KL has higher overlap for most larger subspaces.The fitted spectra and eigenspaces follow the divergence-dependent preferences predicted in Section 3.2.Source: §5.1 · pp. 7–8, Figure 5 · p. 8

5.2 PRECONDITIONING GAINS OF LARGE-SAMPLE FITS

  • KL achieves the lowest held-out gain error, with SQ next; both stay closer to the held-out target across the spectrum than F and VN.The same pattern appears across the remaining checkpoints and projections.Source: §5.2 · pp. 8–9, Table 2 · p. 9

5.3 FEW-SAMPLE FITS AND ONLINE EMA

  • KL has the lowest recovery error, followed by SQ, VN, and F, across 40-sample fits and EMA; the same ordering holds for aggregate held-out gain error.The pattern holds across both checkpoint sources.Source: §5.3 · p. 9
  • Short-history reliability reflects two effects: empirical second moments underestimate tail scales, while lower-qD factor estimates are more stable with limited gradients.Together, these effects help explain why SQ and KL better compensate for empirical bias and recover their large-sample gains more reliably.Source: §5.3 · p. 10

6 GPT-2 TRAINING

The paper tests whether differences in preconditioning gains correspond to GPT-2 training outcomes. SQ and KL outperform F and VN at every tested width, with KL lower at the two largest widths.

6 GPT-2 TRAINING

  • SQ and KL achieve lower GPT-2 validation losses than F and VN at every width; their losses are similar at smaller widths, while KL is lower at the two largest widths.The comparison uses 12-layer GPT-2 models trained on FineWeb at widths 256–1,536, with 5,100 updates.Source: §6 · p. 10, Table 4

7 CONCLUSION

The conclusion links divergence-dependent approximation and finite-sample effects to preconditioning and GPT-2 outcomes. The appendices establish the framework’s relations to prior methods and clarify its analytical scope.

7 CONCLUSION

  • SQ and KL better compensate for empirical second-moment underestimation and recover large-sample gains more reliably from short histories; both also achieve lower GPT-2 validation losses than F and VN.The conclusion proposes held-out gain and recovery errors as practical metrics for evaluating new preconditioners.Source: §7 · p. 10

A RELATED WORK

  • The paper studies spectral error allocation and finite-sample stability across F, VN, SQ, and KL, complementing fixed-factor analysis with joint fits on synthetic non-Kronecker and real gradients.The related-work discussion situates this analysis alongside matrix-divergence Kronecker estimation and finite-sample estimation studies.Source: A RELATED WORK · p. 15

B.1 PROOF OF PROPOSITION 1

  • The square-root divergence and PSGD have the same objective up to positive scale and an additive constant, including under the Kronecker constraint.The q→1 and q→0 limits recover the VN and KL divergences, respectively.Source: §B.1 · p. 16

B.2 LIMITED KRONECKER CAPACITY

  • A single Kronecker product generally cannot exactly represent a gradient second moment unless that moment itself has Kronecker structure.The factor parameter count is smaller than that of a general symmetric matrix when m,n ≥ 2, so divergence choice allocates the remaining approximation error.Source: §B.2 · p. 18

B.3 SPECTRAL ERROR DECOMPOSITION AND KRONECKER APPROXIMATIONS

The spectral analysis separates eigenvalue mismatch from eigenspace misalignment and shows that divergence choice changes how these errors depend on scale and direction of mismatch.

B.3 SPECTRAL ERROR DECOMPOSITION AND KRONECKER APPROXIMATIONS

  • The divergence cost decomposes into eigenvalue mismatch for aligned eigenspaces and a separate cost for eigenspace misalignment.For Kronecker approximations, both mismatch terms retain the stated scaling relations and reciprocal-error comparisons.Source: §B.3 · p. 18, §B.3.4 · p. 20
  • F’s eigenvalue penalty scales quadratically under joint rescaling, whereas KL’s penalty is scale-invariant.The respective scaling relations are r^2 and unchanged.Source: §B.3.2 · p. 19
  • KL penalizes underestimation more than reciprocal overestimation, while VN penalizes overestimation more; SQ treats the reciprocal penalties equally.For F, the penalty is also greater for overestimation when r > 1.Source: §B.3.3 · p. 19

B.4 PROOF OF PROPOSITION 4

The appendix relates divergence-dependent factor-estimation error and local geometry to Shampoo update schemes and validation. It gives proof details, compares SQ and KL geometry, and describes offline and online fitting procedures.

B.4 PROOF OF PROPOSITION 4

  • The proof orders factor-estimation error through effective rank: effective rank decreases with qD, so larger exponents imply larger expected error, with equality for isotropic opposite-factor spectra.The argument estimates one factor with the other fixed, and the right-factor identity follows symmetrically.Source: §B.4 · p. 21

B.5 LOCAL GEOMETRY OF SQ AND KL

  • Only KL has a Hessian proportional to the Gaussian Fisher operator for every positive-definite S; SQ’s local geometry does not match it.At S = sI, the SQ-to-Fisher ratio varies as 2√s.Source: §B.5 · p. 21

C.1 ORIGINAL SHAMPOO

  • Original Shampoo uses marginal statistics to form inverse-root preconditioners, with p = 1/4 for the original update and p = 1/2 for the Kronecker inverse square root.Source: §C.1 · p. 22

C.2 PSGD-KRON, SQ-SHAMPOO, AND KL-SHAMPOO

  • The methods specify distinct factor updates: PSGD-Kron uses relative-gradient updates, while SQ- and KL-Shampoo use spectral EMA with divergence-specific statistics.The spectral EMA updates factor statistics and eigenvalue estimates over time and periodically refreshes eigenbases.Source: §C.2 · p. 22

C.3 DIVERGENCE-BASED FACTOR STATISTICS

C.4 OFFLINE FIXED-POINT FITS

  • Offline fits use a bilateral fixed-point rule with under-relaxation, initialized at identity factors and iterated to numerical convergence.For few-sample KL fitting, the criterion remains finite even when the empirical second moment is singular.Source: §C.4 · p. 23

C.5 ONLINE EMA UPDATES

  • Online EMA updates average factor accumulators and eigenvalue estimates, refresh eigenbases periodically, and precondition the current momentum before adding the gradient to the statistics.The EMA diagnostics use a basis-refresh period of five steps.Source: §C.5 · p. 23

D SYNTHETIC VALIDATION

  • Synthetic validation examines spectral error allocation and finite-sample stability using a smooth, long-tailed second moment that is not exactly Kronecker.Source: D SYNTHETIC VALIDATION · p. 23

D.1 EXPERIMENTAL SETUP

The experiments compare divergence-based fits on synthetic gradients and test directional gains in a stochastic quadratic. They also specify the real-gradient diagnostics used to assess spectra and gain behavior.

D.1 EXPERIMENTAL SETUP

  • The synthetic second moment combines three matrix-normal components with differing decay rates and orientations, plus isotropic noise, so it is generally not a single Kronecker product.Fits use a fixed pool of 4096 gradients and stop at a relative-factor-change threshold or 10,000 iterations.Source: §D.1 · p. 24

D.2 METRICS AND RESULTS

  • Synthetic fits showed mean stability errors ordered KL < SQ < VN < F at every tested sample size, while F most severely underestimated small eigenvalues.SQ and KL had larger fitted tail eigenvalues.Source: §D.2 · p. 24, Figure 7 · p. 25

D.3 DIRECTIONAL GAIN INTERVENTIONS ON A QUADRATIC

  • Replacing KL’s leading gains with F’s raised quadratic loss to 1.441, while replacing its tail gains raised loss to 0.03333; corresponding SQ replacements stayed close to KL.KL’s loss was 0.00844, and the SQ replacements yielded 0.00857 and 0.00855 for leading and tail directions, respectively.Source: §D.3 · p. 25, Figure 8 · p. 26

E REAL-GRADIENT EXPERIMENTAL DETAILS

  • Real-gradient diagnostics analyze query, value, and attention-output gradients from block 5 of a 12-layer GPT-2 at steps 100, 2500, and 4500.Second moments and eigendecompositions use FP64, with retained eigenvalues above the stated machine-precision threshold.Source: E REAL-GRADIENT EXPERIMENTAL DETAILS · p. 26, Table 8 · p. 26

E.1.1 FINITE-SAMPLE TAIL SELECTION

Independent-split comparisons show that empirical second moments tend to underestimate tail scales, with small source eigenvalues amplifying held-out ratios. The section then examines reliability, divergence-specific spectra and gain recovery across checkpoints, including Distributed Shampoo.

E.1.1 FINITE-SAMPLE TAIL SELECTION

  • Downward sampling error in a selected source eigenvalue raises its conditional expected held-out ratio, and the effect can be large in the tail because the error is divided by a small source eigenvalue.Cross-fit errors have medians and IQRs near zero, supporting sampling-induced selection bias rather than distribution shift.Source: §E.1.1 · pp. 26–27

E.1.2 HIGH-DIMENSIONAL SAMPLING AND TAIL SIGNALS

  • Tail directions can carry useful signals, such as directions associated with infrequent classes, but finite-sample noise can still severely underestimate their scales.Source: §E.1.2 · p. 27

E.2 RELIABILITY ACROSS CHECKPOINTS

  • Across checkpoints and projections from AdamW and Distributed Shampoo, held-out scale error is near zero in leading directions and strongly positive in the tail.Table 9 reports regional medians and defines positive error as underestimation.Source: Table 9 · p. 27

E.3 INVERSE-ROOT GAIN ACROSS CHECKPOINTS

  • Held-out gain errors are evaluated for large-sample fits, few-sample fits, and EMA states at steps 100, 2500, and 4500.Few-sample and EMA curves average over four and 16 repeats, respectively.Source: §E.3 · p. 28, Figure 10 · p. 28, Figure 11 · p. 29

E.4 KRONECKER FITS OF DIFFERENT DIVERGENCES

  • F and VN favor large eigenvalues, whereas SQ and KL yield larger fitted tail eigenvalues; KL penalizes relative errors equally across scales and underestimation more strongly.Small fitted eigenvalues can also increase gains when fitted eigenvectors differ from source eigenvectors.Source: §E.4 · pp. 30–31

E.5 GAIN RECOVERY WITH SHORT GRADIENT HISTORIES

  • Figure 14 compares large-sample fit recovery with 40-sample fits and EMA estimates for F, VN, SQ, and KL at step 2500.The caption identifies solid curves as 40-sample fits and dashed curves as EMA estimates.Source: Figure 14 · p. 31, §E.5 · p. 31

E.6 DIAGNOSTICS ON DISTRIBUTED SHAMPOO CHECKPOINTS

  • Distributed Shampoo diagnostics at step 2500 report held-out scale error, cross-fit controls, and held-out gain error for large-sample, 40-sample, and EMA fits.The 40-sample fits use four repeats; EMA uses 512 gradients and 16 repeats.Source: Figure 15 · p. 33, §E.6 · p. 33

F GPT-2 TRAINING DETAILS AND FULL RESULTS

The GPT-2 training details specify model, data, and optimizer settings, while the reported figure presents validation-loss trajectories from step 2,000 to the end of training. The supplied passages do not state the comparative outcome of those trajectories.

F GPT-2 TRAINING DETAILS AND FULL RESULTS

Related papers

Loading 2610.08534v1…