🤖 machine learning

Rethinking Bregman Divergences in Kronecker-Factored Optimizers

This paper analyzes how different Bregman divergences distribute Kronecker approximation errors across the covariance spectrum, revealing that top eigenspaces are reliable while the tail is noisy, which motivates a new subspace-aware optimizer that combines eigenvalue-based preconditioning with adaptive isotropic acceleration.

Original authors: Bing Liu, Wenjie Zhou, Chengcheng Zhao

Published 2026-06-02

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Bing Liu, Wenjie Zhou, Chengcheng Zhao

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Navigating a Foggy Mountain

Imagine you are trying to guide a massive, heavy robot (an AI model) down a mountain to find the lowest valley (the best possible performance). The robot is too big to turn on a dime, so it needs a map to tell it which way is "down."

In the world of AI, this map is called a preconditioner. It tells the robot how steep the ground is in every direction so it can take the perfect step size.

The problem? The mountain is huge, and drawing a perfect, detailed map of every single rock and pebble is impossible and too slow. So, engineers use a shortcut: they approximate the map using a "Kronecker-factored" structure. Think of this as trying to describe a complex 3D landscape by only looking at two separate 2D slices (like looking at a shadow from the side and a shadow from the front) and combining them.

The Problem: The "Shadow" Isn't Perfect

The authors point out a fundamental flaw in this shortcut: The shadow never perfectly matches the real object.

Because the real mountain (the data) is complex, you can never perfectly reconstruct it from just two 2D slices. There will always be "approximation error"—parts of the map that are slightly wrong.

For a long time, researchers thought all these shortcut methods were basically the same, just using slightly different math formulas (called Bregman Divergences) to measure how "wrong" the shadow was. The paper asks: If we can't get a perfect map, does it matter which formula we use to measure the error?

The Discovery: Different Formulas, Different Blind Spots

The authors discovered that yes, it matters a lot. Different formulas treat the "mistakes" in the map differently:

Frobenius (The "Big Picture" Fan): This formula cares mostly about the big, obvious features of the mountain (the large peaks and deep valleys). It ignores the tiny pebbles. It's like a cartographer who only draws the major highways and ignores the side streets.
LogDet (The "Microscope"): This formula is obsessed with the tiny details. It gets very sensitive to the small, noisy parts of the map. It tries to fix every tiny pebble, even if that pebble is just a random rock that doesn't represent the true shape of the mountain.
von Neumann (The "Middle Ground"): This sits somewhere in between, caring about the big features but not as aggressively as the first one.

The Key Insight: The authors found that the "big features" (the top spectrum of the data) usually align well with the true shape of the mountain (the Hessian). However, the "tiny details" (the bottom spectrum) are often just noise—random static that doesn't actually help the robot navigate.

If you use the "Microscope" (LogDet) approach, you end up trying to steer the robot based on random noise, which makes it wobble and move slower. If you use the "Big Picture" approach, you focus on the reliable parts.

The Solution: The "Subspace-Aware" Optimizer

Instead of trying to fix the whole map with one formula, the authors propose a new strategy called BregTop. They split the map into two zones:

The "High-Confidence" Zone (Top Spectrum):
- What it is: The big, clear features of the mountain.
- Strategy: Use a precise, eigenvalue-based map here. This tells the robot exactly how to turn and how fast to go based on the real terrain.
- Analogy: This is like using a GPS with high-definition satellite imagery for the main highways.
The "Noisy" Zone (Bottom Spectrum):
- What it is: The tiny, unreliable details that are mostly just static.
- Strategy: Stop trying to map these details! Instead, just give the robot a steady, uniform push. Don't try to steer around every tiny pebble; just keep moving forward at a consistent speed.
- Analogy: This is like driving through a foggy field where you can't see the ground. Instead of trying to dodge every invisible rock, you just keep driving straight at a safe, constant speed.

The Results

The authors tested this new "split strategy" (BregTop) against the old methods (like standard Shampoo and its variants) on a language model training task.

The Outcome: The new method reached the target performance faster (in fewer steps) than the others.
Why: By trusting the reliable "big picture" data and ignoring the noisy "tiny details," the robot didn't waste energy wobbling around. It moved more efficiently down the mountain.

Summary

The paper argues that when we can't perfectly map the complex world of AI data, we shouldn't try to fix every tiny error. Instead, we should:

Trust the big, clear patterns.
Ignore the noisy, unreliable details by treating them simply and uniformly.

This "know what to trust and what to ignore" approach leads to faster and more efficient AI training.

Technical Summary: Rethinking Bregman Divergences in Kronecker-Factored Optimizers

Problem Statement
Shampoo-style optimizers approximate the gradient covariance matrix $C = E[\text{vec}(G)\text{vec}(G)^\top]$ using a Kronecker-factored structure $L \otimes R$ to reduce computational complexity. Recent work (Lin et al., 2026) unified various Shampoo variants by framing them as minimization problems of Bregman matrix divergences between $C$ and $L \otimes R$ . However, a critical gap remains: for generic dense covariance matrices, exact Kronecker factorization is impossible, meaning approximation error is unavoidable. It is unclear how different Bregman divergences distribute this error across the covariance spectrum and whether the resulting preconditioners trust all parts of the spectrum equally. Furthermore, since the gradient covariance often differs significantly from the true Hessian, not all spectral components are suitable for preconditioning. The central question addressed is: When exact Kronecker matching is impossible, which parts of the covariance spectrum should a structured optimizer approximate and trust?

Methodology and Analysis
The authors approach this problem through two main analytical lenses: the spectral geometry of Bregman approximations and the empirical reliability of covariance eigenspaces.

Spectral Geometry of Bregman Approximation:
The paper characterizes how Frobenius, von Neumann, and LogDet divergences penalize approximation errors differently. By decomposing the divergence into marginal spectrum terms and a cross-spectral coupling term $g_F(\lambda_i, \omega_j)$ , the authors show:
- Frobenius Divergence: The coupling term $-\lambda_i \omega_j$ imposes a bilinear magnitude penalty, strongly emphasizing the top (large eigenvalue) spectrum.
- von Neumann Divergence: The term $-\lambda_i \log \omega_j$ retains a preference for the top spectrum but dampens it logarithmically.
- LogDet Divergence: The term $\lambda_i / \omega_j$ creates a relative ratio distortion, making the approximation highly sensitive to small eigenvalues (the tail spectrum).
Theorem 2 establishes that stationary points of these minimization problems are not determined by the raw residual $(S^* - C)$ , but by a divergence-weighted residual $\nabla^2 F(S^*)[S^* - C]$ . Consequently, the resulting Kronecker factors align with the specific spectral preferences of the chosen divergence (e.g., LogDet factors are more sensitive to small eigenvalues, while Frobenius factors focus on dominant subspaces).
Spectral Reliability of Covariance Preconditioning:
The authors empirically investigate the alignment between the gradient covariance and the Hessian eigenspaces on a Transformer model trained on SST-2. Results (Table 3) indicate that the top covariance eigenspace is substantially better aligned with the Hessian than the lower-eigenvalue components. The tail spectrum is observed to be noisier and less reliable for preconditioning, consistent with prior findings (Kunstner et al., 2019; Thomas et al., 2020).
Proposed Solution: Subspace-Aware Kronecker Optimizer (BregTop):
Motivated by the mismatch between divergence-induced spectral preferences and the actual reliability of the spectrum, the authors propose BregTop. This optimizer combines the strengths of both observations:
- It selects Kronecker factors using a top-sensitive divergence (e.g., von Neumann or Frobenius).
- It applies eigenvalue-based preconditioning only within the top subspace (the largest $\lceil \rho mn \rceil$ entries of the joint Kronecker spectrum).
- It treats the complementary bottom subspace with an adaptive isotropic acceleration constant ( $\chi_{\rho, q}$ ), rather than attempting to precondition the noisy tail with potentially misleading curvature information.
Theoretically, this decomposes the preconditioner into a curvature-scaled top component and a uniform acceleration component for the bottom space (Theorem 4).

Key Contributions

Spectral Geometry Characterization: The paper provides a formal characterization of how Frobenius, von Neumann, and LogDet divergences induce different eigenvalue-dependent penalties, showing that stationary Kronecker factors are governed by divergence-weighted residuals rather than raw errors.
Empirical Spectral Reliability: The work demonstrates that while the full covariance spectrum is unreliable for preconditioning, the top eigenspace maintains strong alignment with the Hessian, justifying a subspace-aware approach.
Subspace-Aware Optimizer Design: The authors introduce BregTop, an algorithm that adaptively switches between eigenvalue-based scaling in the reliable top subspace and isotropic acceleration in the noisy bottom subspace.

Experimental Results
The authors evaluated BregTop against standard Shampoo, KL-Shampoo, VN-Shampoo, and F-Shampoo on a 4-layer Transformer trained on the full SST-2 dataset. The metric used was the number of steps required for the EMA-smoothed training loss to fall below 0.05.

BregTop-VN-v2 achieved the target loss in 1003 steps, outperforming the standard Shampoo (1119 steps) and the baseline VN-Shampoo (1127 steps).
The BregTop variants consistently improved upon their corresponding divergence baselines (e.g., BregTop-VN improved over VN-Shampoo; BregTop-F improved over F-Shampoo), validating the benefit of isolating the top subspace for preconditioning.

Significance and Claims
The paper claims that the design of structured optimizers should account for the unavoidable approximation error in Kronecker factorization and the varying reliability of the covariance spectrum. By recognizing that different Bregman divergences inherently prioritize different spectral regions, and that only the top spectrum is reliably aligned with the Hessian, the proposed subspace-aware approach offers a more robust preconditioning strategy. The authors note that their design is consistent with training dynamics observed in Song et al. (2025), where learning occurs in the bottom subspace; thus, applying isotropic acceleration there preserves necessary gradient flow while dampening noisy updates in high-curvature directions via the top subspace preconditioning. The results suggest that this hybrid approach improves step efficiency without requiring full second-order information.