# 4-bit Shampoo for Memory-Efficient Network Training

arXiv: https://arxiv.org/abs/2405.18144 · HTML: https://mlsanity.com/paper/2405.18144/

## Key facts

- Cited by 18 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 42 references that are in the map.
- First posted on arXiv 28 May 2024.

## Builds on (most cited first, 24 of 24)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [QLoRA: Efficient Finetuning of Quantized LLMs](https://mlsanity.com/paper/2305.14314/) (2023)
- [Conformer: Convolution-augmented Transformer for Speech Recognition](https://mlsanity.com/paper/2005.08100/) (2020)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [8-bit Optimizers via Block-wise Quantization](https://mlsanity.com/paper/2110.02861/) (2021)
- [Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training](https://mlsanity.com/paper/2305.14342/) (2023)
- [ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning](https://mlsanity.com/paper/2006.00719/) (2020)
- [The Road Less Scheduled](https://mlsanity.com/paper/2405.15682/) (2024)
- [Vision Transformer for Small-Size Datasets](https://mlsanity.com/paper/2112.13492/) (2021)
- [8-Bit Approximations for Parallelism in Deep Learning](https://mlsanity.com/paper/1511.04561/) (2015)
- [M-FAC: Efficient Matrix-Free Approximations of Second-Order Information](https://mlsanity.com/paper/2107.03356/) (2021)
- [Memory Efficient Optimizers with 4-bit States](https://mlsanity.com/paper/2309.01507/) (2023)
- [Disentangling Adaptive Gradient Methods from Learning Rates](https://mlsanity.com/paper/2002.11803/) (2020)
- [Sketchy: Memory-efficient Adaptive Regularization with Frequent Directions](https://mlsanity.com/paper/2302.03764/) (2023)

## Built on by (most cited first, 18 of 18)

- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)
- [Dion: Distributed Orthonormalized Updates](https://mlsanity.com/paper/2504.05295/) (2025)
- [COSMOS: A Hybrid Adaptive Optimizer for Memory-Efficient Training of LLMs](https://mlsanity.com/paper/2502.17410/) (2025)
- [Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner](https://mlsanity.com/paper/2506.03595/) (2025)
- [SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training](https://mlsanity.com/paper/2412.13148/) (2024)
- [Low-Precision Training of Large Language Models: Methods, Challenges, and Opportunities](https://mlsanity.com/paper/2505.01043/) (2025)
- [AbbIE: Autoregressive Block-Based Iterative Encoder for Efficient Sequence Modeling](https://mlsanity.com/paper/2507.08567/) (2025)
- [Understanding Quantization of Optimizer States in LLM Pre-training: Dynamics of State Staleness and Effectiveness of State Resets](https://mlsanity.com/paper/2603.16731/) (2026)
- [Scalable Parameter and Memory Efficient Pretraining for LLM: Recent Algorithmic Advances and Benchmarking](https://mlsanity.com/paper/2505.22922/) (2025)
- [Conda: Column-Normalized Adam for Training Large Language Models Faster](https://mlsanity.com/paper/2509.24218/) (2025)
- [Memory-Efficient 4-bit Preconditioned Stochastic Optimization](https://mlsanity.com/paper/2412.10663/) (2024)
- [DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers](https://mlsanity.com/paper/2602.02016/) (2026)
- [Mitigating Staleness in Asynchronous Pipeline Parallelism via Basis Rotation](https://mlsanity.com/paper/2602.03515/) (2026)
- [Budget-aware Auto Optimizer Configurator](https://mlsanity.com/paper/2605.04711/) (2026)
- [MuonQ: Enhancing Low-Bit Muon Quantization via Directional Fidelity Optimization](https://mlsanity.com/paper/2605.11396/) (2026)
- [DHO2: Accelerating Distributed Hybrid Order Optimization via Model Parallelism and ADMM](https://mlsanity.com/paper/2505.00982/) (2025)
- [Clean: Second-order LLM Training at Linear Memory Cost via Nystr"om Sketching](https://mlsanity.com/paper/2610.04204/) (2026)
- [Minimax Optimality of Shampoo for Nonstationary Stochastic Optimization](https://mlsanity.com/paper/2610.16923/) (2026)

## Abstract (arXiv)

Second-order optimizers, maintaining a matrix termed a preconditioner, are superior to first-order optimizers in both theory and practice. The states forming the preconditioner and its inverse root restrict the maximum size of models trained by second-order optimizers. To address this, compressing 32-bit optimizer states to lower bitwidths has shown promise in reducing memory usage. However, current approaches only pertain to first-order optimizers. In this paper, we propose the first 4-bit second-order optimizers, exemplified by 4-bit Shampoo, maintaining performance similar to that of 32-bit ones. We show that quantizing the eigenvector matrix of the preconditioner in 4-bit Shampoo is remarkably better than quantizing the preconditioner itself both theoretically and experimentally. By rectifying the orthogonality of the quantized eigenvector matrix, we enhance the approximation of the preconditioner's eigenvector matrix, which also benefits the computation of its inverse 4-th root. Besides, we find that linear square quantization slightly outperforms dynamic tree quantization when quantizing second-order optimizer states. Evaluation on various networks for image classification and natural language modeling demonstrates that our 4-bit Shampoo achieves comparable performance to its 32-bit counterpart while being more memory-efficient.