Reproduce “Shampoo: Preconditioned Stochastic Tensor Optimization” on OpenResearch

About this paper

Vineet Gupta, Tomer Koren, Yoram Singer

Google Brain researchers developed Shampoo, a preconditioned stochastic tensor optimization algorithm, to efficiently leverage the tensor structure of deep learning model parameters. The algorithm achieved faster convergence and improved performance on state-of-the-art deep neural networks across vision and language tasks while maintaining a competitive runtime per step compared to standard optimizers.

Key insights

• Exploiting the inherent tensor structure of deep learning parameters provides a principled way to achieve more powerful preconditioning than diagonal methods without incurring the full computational cost of explicit full-matrix approaches.
• The dimension-wise preconditioners, when combined through Kronecker-related structures, can implicitly lower-bound a full-matrix AdaGrad-like preconditioner, thereby capturing essential curvature information efficiently.
• Efficient implementation using modern tensor operations and strategic heuristics (e.g., infrequent preconditioner updates) allows for a sophisticated preconditioning scheme to operate with per-step runtimes comparable to simpler first-order optimizers.

Summary by alphaXiv.