Use eigenvalue EMA over per-step outer products for per-factor correction#264

Open

runame wants to merge 7 commits into

facebookresearch:mainfrom

runame:pr4/eigenvalue-ema-v2

Commits on May 7, 2026

Refactor: extract shared EigendecompositionBasedShampooKroneckerFactorsUnwrapped base class

runame
and
claude
committed
Refactor: eliminate _compute_outer_product_list via _transform_grad_for_outer_product hook

runame
and
claude
committed
Add per-factor eigenvalue correction for Distributed Shampoo

runame
and
claude
committed
Add tests for PerFactorEigenvalueCorrectedKLShampooPreconditionerList

runame
and
claude
committed
Add full inverse_exponent_override docstrings to PerFactor config classes

runame
and
claude
committed
Export PerFactor config classes from distributed_shampoo __init__.py

runame
and
claude
committed
Use eigenvalue EMA over per-step outer products for per-factor correction

runame
and
claude
committed