Extreme Tensoring for Low-Memory Preconditioning

2019/02/12 by Xinyi Chen, Chen, Xinyi, Naman Agarwal +7
Computer Science · Mathematics · Physics and Astronomy · #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Model Reduction and Neural Networks #Stochastic Gradient Optimization Techniques #Tensor decomposition and applications

paper · pdf · doi:10.48550/arxiv.1902.04620

openalex publication_date 2019/02/12 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

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 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.

Citations

• ADADELTA: An Adaptive Learning Rate Method
• Neural Optimizer Search with Reinforcement Learning
• Learning Multiple Layers of Features from Tiny Images

Related

• Shampoo: Preconditioned Stochastic Tensor Optimization
• Tensor Programs IVb: Adaptive Optimization in the Infinite-Width Limit
• Hardware Acceleration of Sparse and Irregular Tensor Computations of ML Models: A Survey and Insights
• On The Concurrence of Layer-wise Preconditioning Methods and Provable Feature Learning
• Memory-Efficient LLM Pretraining via Minimalist Optimizer Design