What is Shampoo and when does second-order optimization pay off?
hardAnswer
- Shampoo (Gupta et al., 2018 / Anil et al., 2020) is a second-order-ish optimizer that maintains a per-parameter matrix preconditioner factorized along tensor axes — computes the inverse Kronecker-factored covariance of gradients.
- Uses more compute per step than Adam but converges in far fewer steps.
- Practical for very large-scale training where compute is limited by memory bandwidth rather than FLOPs (T5, PaLM experiments).
Check yourself — multiple choice
- Shampoo uses no gradient information
- Second-order-ish preconditioner (Kronecker-factored covariance) — fewer steps, more per-step compute
- Same as SGD
- Only works on GPUs
Shampoo: Kronecker-factored preconditioner → fewer steps, worthwhile at scale.
#optimizers#optimization
Practise Deep Learning
214 interview questions in this topic.