PSGD's Objective Shown Equivalent to KL-Shampoo, Sparking Optimizer Discussion

YouJiacheng discovered that the online PCA cost function in Appendix H.3 of the paper "Curvature-Informed SGD via General Purpose Lie-Group Preconditioners" (arXiv 2402.04553) is exactly identical to the KL objective in KL-Shampoo, and that PSGD's original objective combined with the QEQ update rule yields an update rule exactly the same as KL+AIRM, revealing a deep mathematical connection underlying preconditioned optimizers. The finding sparked discussion among researchers on PSD parameterization and eigendecomposition-free optimizer design.

Confirmed

• @YouJiacheng noted that the online PCA cost in Appendix H.3 of arXiv 2402.04553 is exactly equivalent to the KL objective in KL-Shampoo, i.e., using the KL divergence KL(N(0, Σ), N(0, P⁻²)) as the objective, though the original paper did not interpret it as a KL divergence.
• @YouJiacheng further found that PSGD's original objective combined with the QEQ update rule gives an update rule identical to KL+AIRM.
• @YouJiacheng also discussed the uniqueness of that gradient when requiring the Euclidean gradient to depend linearly on Σ.
• Building on this, @YouJiacheng proposed a new optimizer combining PSGD-Kron multiplicative updates with the KL-Shampoo objective, with the advantages of requiring no eigendecomposition and handling gauge degrees of freedom (see repost m1).
• @arohan documented discussion among researchers: @YouJiacheng asked whether this online PCA cost is known in the literature, with @HessianFree and others joining the exchange.
• @HessianFree recalled that the relevant optimizer formulas were derived on a whiteboard together with Xilin, and believes these formulas help handle high gradient noise or low-rank regimes.

Why it matters

• The equivalence suggests that seemingly different preconditioned SGD methods (PSGD, Shampoo, KL+AIRM) may share the same mathematical foundation, helping unify our understanding of second-order optimizers.
• If the eigendecomposition-free design holds up, it could cut the computational cost of matrix-type optimizers, making them more suitable for large-model training.

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

• 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
• [source] Optimizer Formulas Born on Whiteboard: Useful for Noisy or Low-Rank Gradients — HessianFree · 2026-09-01
• Researchers spot online PCA cost equals KL-shampoo objective in optimizer discussion — _arohan_ · 2026-09-01