PSGD Objective Found Mathematically Equivalent to KL-Shampoo

YouJiacheng discovered that the online PCA cost in arXiv 2402.04553 Appendix H.3 is exactly equivalent to the KL objective in KL-Shampoo, and that PSGD's original objective combined with QEQ update rules yields updates identical to KL+AIRM, revealing deep mathematical connections between preconditioned optimizers. The finding sparked discussion among researchers on PSD parameterization and eigendecomposition-free optimizer design.

Confirmed

@YouJiacheng showed the Appendix H.3 online PCA cost equals the KL divergence KL(N(0, Σ), N(0, P⁻²)) used in KL-Shampoo, though the original paper did not interpret it as a KL divergence.
@YouJiacheng further found that PSGD's original objective with QEQ updates produces exactly the same update rule as KL+AIRM.
@YouJiacheng also discussed the uniqueness of the Euclidean gradient when it is required to depend linearly on Σ.
Building on this, @YouJiacheng proposed a new optimizer combining PSGD-Kron multiplicative updates with the KL-Shampoo objective, which avoids eigendecomposition and handles gauge freedoms (see repost m1).
@arohan documented researcher discussion: @YouJiacheng asked whether this online PCA cost is known in the literature, with @HessianFree and others joining.
@HessianFree recalled the optimizer formulas were derived on a whiteboard with Xilin and believes they help with noisy or low-rank gradients.

Why it matters

The equivalence suggests that seemingly different preconditioned SGD methods (PSGD, Shampoo, KL+AIRM) may share a common mathematical foundation, aiding a unified understanding of second-order optimizers.
An eigendecomposition-free design, if it holds, could cut the computational cost of matrix-type optimizers, making them more practical for large-model training.

2026-08-31 ~ 2026-09-01 · 6 related posts

Primary sources

PSGD cost function matches KL-shampoo objective exactly, author notes — YouJiacheng · 2026-08-31 15:01 (UTC+08:00)
PSGD's QEQ update rule yields same rule as KL+AIRM — YouJiacheng · 2026-08-31 15:03 (UTC+08:00)
Researchers spot online PCA cost equals KL-shampoo objective in optimizer discussion — _arohan_ · 2026-09-01 00:57 (UTC+08:00)

Optimizer combining PSGD and KL-Shampoo avoids eig/inverse/solve — YouJiacheng · 2026-08-31
[source] PSGD cost function matches KL-shampoo objective exactly, author notes — YouJiacheng · 2026-08-31
[source] PSGD's QEQ update rule yields same rule as KL+AIRM — YouJiacheng · 2026-08-31
Euclidean gradient uniqueness under Σ dependency — YouJiacheng · 2026-08-31
Optimizer Formulas Born on Whiteboard: Useful for Noisy or Low-Rank Gradients — HessianFree · 2026-09-01
[source] Researchers spot online PCA cost equals KL-shampoo objective in optimizer discussion — _arohan_ · 2026-09-01