Skip to content

Change initial eigen values in KLShampoo to 1 #310

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
skyw merged 2 commits into from
Oct 16, 2026
Merged

Change initial eigen values in KLShampoo to 1 #310

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ac6a500
change init value
skyw Oct 14, 2026
070a043
grammar fix
skyw Oct 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions emerging_optimizers/shampoo/shampoo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -195,8 +195,10 @@ def init_state(
{
"Q_L": torch.eye(m, device=device),
"Q_R": torch.eye(n, device=device),
"eigvals_L": torch.zeros(m, device=device),
"eigvals_R": torch.zeros(n, device=device),
# Initialize eigenvalues to one to skip KL correction for the first step.
# They will be overwritten immediately afterward.
"eigvals_L": torch.ones(m, device=device),
"eigvals_R": torch.ones(n, device=device),
Comment on lines +200 to +201

@greptile-apps greptile-apps Bot Oct 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 First-step fix lacks a test

The tests do not protect the behavior changed here. test_init_state_layout checks eigenvalue shapes but not their values. test_update_kronecker_factors_matches_legacy checks only a later update using the same helper. Returning to zero initialization could therefore pass both tests.

Add a first-step test with a rectangular gradient G, shampoo_beta=0, and small positive eps. Check that init_step produces L = G @ G.T / n and R = G.T @ G / m.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

}
)
return state
Expand Down
Loading