Shampoo: Diagonal Preconditioning
Advanced
The problem, from first principles
Shampoo takes the inverse fourth root of the preconditioners, not the inverse. With a left and right preconditioner each taking a quarter power, the overall scaling matches an inverse square root of the full curvature.
From theory to code
Implement precondition_diag(L, G, R), the two-sided scaling under diagonal preconditioners.
Constraints
- Use the element-wise inverse fourth root of L and R.
Hints
Hint 1
Multiply each row of G by L to the minus one quarter, and each column by R to the minus one quarter.
Keywords
research-papersoptimizationpreconditioning
Python
Run your code to see results
No results yet
Run your code to check it against the examples, or submit to run every test.