Research map / Architectures and efficiency / Quantization
4-bit Shampoo for Memory-Efficient Network Training
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes 4-bit Shampoo, the first low-bit second-order optimizer, by quantizing preconditioner eigenvector matrices (with Bjorck orthogonal rectification and linear-square quantization), matching 32-bit Shampoo accuracy with lower memory.
Problem
Second-order optimizers like Shampoo store large preconditioner and inverse-root states, limiting trainable model size; prior low-bit optimizer state compression only covered first-order optimizers.
Approach
Quantize the eigenvector matrix of the preconditioner rather than the preconditioner itself (so small singular values, which control the inverse 4th root, are preserved), rectify the quantized eigenvectors' orthogonality with one Bjorck iteration, and use linear-square quantization instead of dynamic tree quantization.
Key results
- 4-bit Shampoo test accuracy differs from 32-bit Shampoo by -0.7% to 0.5% across CIFAR-100, Tiny-ImageNet and ImageNet-1k models.
- 4-bit Shampoo saves 4.5% to 41% total GPU memory relative to 32-bit Shampoo, with wall-clock time increases from -0.2% to 9.5%.
- Quantizing the eigenvector matrix with Linear-2 and orthogonal rectification reduces 4-bit NRE in A^-1/4 from 0.6243 (quantizing A) to 0.0343 on the real-world matrix.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use 4-bit Shampoo instead of 32-bit Shampoo | Shampoo · Image classification benchmarks in Table 1 | test accuracy: differences ranging from -0.7% to 0.5% | Mixed |
| Use 4-bit Shampoo instead of 32-bit Shampoo | Shampoo · Image classification benchmarks in Table 1 | total GPU memory cost: memory savings of 4.5% to 41% | Helped |
| Quantize eigenvector matrix U rather than preconditioner A (ablation) | 4-bit Shampoo · CIFAR-100, Swin-Tiny | test accuracy: 76.95 -> 78.63 | Helped |
| Add orthogonal rectification (ablation) | 4-bit Shampoo with Linear-2 and quantized U · CIFAR-100, Swin-Tiny | test accuracy: 78.22 -> 78.63 | Helped |
| Replace DT with Linear-2 quantization (ablation) | 4-bit Shampoo with quantized U and no OR · CIFAR-100, Swin-Tiny | test accuracy: 78.70 -> 78.22 | No clear effect |
| Use 4-bit instead of 3-bit quantization (ablation) | Shampoo with Linear-2 quantization of U and OR · CIFAR-100, Swin-Tiny | test accuracy: 78.63 and 78.53 | No clear effect |
| Quantize U without orthogonal rectification at 3-bit (ablation) | Shampoo · CIFAR-100, Swin-Tiny | training loss: NaN | Hurt |
Surprising findings
- Linear-2 quantization is only marginally better than dynamic tree quantization for second-order states, and the 4-bit eigenvector quantization error is smaller than 8-bit quantization of the preconditioner itself.
- Quantizing the preconditioner directly at 4-bit causes very large errors (NRE about 0.62) in its inverse 4th root.
Limitations
- Evaluation limited to image classification and natural language modeling; not tested on billion-parameter models.
- Eigenvector matrices are non-symmetric so they cannot exploit the upper-triangular storage trick, making 8-bit preconditioner memory equal to 4-bit eigenvector memory.
- Theoretical results rely on strong assumptions (e.g., two distinct singular values, large condition number).
Experimental setup
Datasets: CIFAR-100 Tiny-ImageNet ImageNet-1k
Benchmarks: CIFAR-100 Tiny-ImageNet ImageNet-1k
Metrics: test accuracy wall-clock time total GPU memory cost normwise relative error angle error
Baselines: 32-bit Shampoo SGDM AdamW naive 4-bit Shampoo (quantizing preconditioner) 32-bit K-FAC 32-bit AdaBK 32-bit CASPR
Builds on: Shampoo block-wise 8-bit optimizers 4-bit Adam states Bjorck orthonormalization K-FAC AdaBK CASPR
Closest papers by meaning
- Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner (NeurIPS 2025 · similarity 0.88)
- A New Perspective on Shampoo's Preconditioner (ICLR 2025 · similarity 0.86)
- A Stable Whitening Optimizer for Efficient Neural Network Training (NeurIPS 2025 · similarity 0.85)
- ParetoQ: Improving Scaling Laws in Extremely Low-bit LLM Quantization (NeurIPS 2025 · similarity 0.84)
- SDP4Bit: Toward 4-bit Communication Quantization in Sharded Data Parallelism for LLM Training (NeurIPS 2024 · similarity 0.84)
- QBB: Quantization with Binary Bases for LLMs (NeurIPS 2024 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.