Extreme Tensoring for Low-Memory Preconditioning

Xinyi Chen, Naman Agarwal, Elad Hazan, Cyril Zhang, Yi Zhang

Keywords: memory, optimization

Abstract: State-of-the-art models are now trained with billions of parameters, reaching hardware limits in terms of memory consumption. This has created a recent demand for memory-efficient optimizers. To this end, we investigate the limits and performance tradeoffs of memory-efficient adaptively preconditioned gradient methods. We propose \emph{extreme tensoring} for high-dimensional stochastic optimization, showing that an optimizer needs very little memory to benefit from adaptive preconditioning. Our technique applies to arbitrary models (not necessarily with tensor-shaped parameters), and is accompanied by regret and convergence guarantees, which shed light on the tradeoffs between preconditioner quality and expressivity. On a large-scale NLP model, we reduce the optimizer memory overhead by three orders of magnitude, without degrading performance.

Similar Papers

Neural Stored-program Memory
Hung Le, Truyen Tran, Svetha Venkatesh,
Large Batch Optimization for Deep Learning: Training BERT in 76 minutes
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, Cho-Jui Hsieh,
Efficient and Information-Preserving Future Frame Prediction and Beyond
Wei Yu, Yichao Lu, Steve Easterbrook, Sanja Fidler,