Shampoo: Preconditioned Stochastic Tensor Optimization

Abstract

Preconditioned gradient methods are among the most general and powerful tools in optimization. However, preconditioning requires storing and manipulating prohibitively large matrices. We describe and analyze a new structure-aware preconditioning algorithm, called Shampoo, for stochastic optimization over tensor spaces. Shampoo maintains a set of preconditioning matrices, each of which operates on a single dimension, contracting over the remaining dimensions. We establish convergence guarantees in the stochastic convex setting, the proof of which builds upon matrix trace inequalities. Our experiments with state-of-the-art deep learning models show that Shampoo is capable of converging considerably faster than commonly used optimizers. Although it involves a more complex update rule, Shampoo's runtime per step is comparable to that of simple gradient methods such as SGD, AdaGrad, and Adam.

References in corpus (6)

• Adam: A Method for Stochastic Optimization
• Optimizing Neural Networks with Kronecker-factored Approximate Curvature
• Path-SGD: Path-Normalized Optimization in Deep Neural Networks
• Convergence rates of sub-sampled Newton methods
• Second-Order Stochastic Optimization for Machine Learning in Linear Time
• Faster SGD Using Sketched Conditioning

Cited by in corpus (19)

• Massively Multilingual Neural Machine Translation in the Wild: Findings and Challenges
• Applications and Techniques for Fast Machine Learning in Science
• Practical Quasi-Newton Methods for Training Deep Neural Networks
• Descending through a Crowded Valley - Benchmarking Deep Learning Optimizers
• Scalable Second Order Optimization for Deep Learning
• Riemannian adaptive stochastic gradient algorithms on matrix manifolds
• Does Knowledge Distillation Really Work?
• Sketchy Empirical Natural Gradient Methods for Deep Learning
• LocoProp: Enhancing BackProp via Local Loss Optimization
• Matrix-Free Preconditioning in Online Learning
• Physics-informed Transformers for Electronic Quantum States
• VecHGrad for Solving Accurately Complex Tensor Decomposition
• Structured second-order methods via natural gradient descent
• Whitening and second order optimization both make information in the dataset unusable during training, and can reduce or prevent generalization
• Tensor Normal Training for Deep Learning Models
• Bridging the Gap between Deep Learning and Frustrated Quantum Spin System for Extreme-scale Simulations on New Generation of Sunway Supercomputer
• Invexifying Regularization of Non-Linear Least-Squares Problems
• mL-BFGS: A Momentum-based L-BFGS for Distributed Large-Scale Neural Network Optimization
• Adaptive Online Learning with Varying Norms