Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Memory-Efficient 4-bit Preconditioned Stochastic Optimization

Jingyang Li, Kuangyu Ding, Kim-Chuan Toh, Pan Zhou

TL;DR

The paper tackles the memory bottleneck of full-matrix preconditioners in Shampoo by introducing a memory-efficient 4-bit quantization strategy based on Cholesky factorization, complemented with an error-feedback mechanism. It provides rigorous convergence guarantees for both smooth and nonsmooth stochastic nonconvex objectives and demonstrates substantial memory savings with minimal performance loss across CNNs, vision transformers, and large language model pre-training. Cholesky quantization better preserves spectral properties than direct quantization, enabling performance close to the 32-bit baseline while using only 4-bit storage. Overall, the work significantly enhances the practicality of non-diagonal preconditioned optimization for large-scale deep learning under realistic hardware constraints.

Abstract

Preconditioned stochastic optimization algorithms, exemplified by Shampoo, outperform first-order optimizers by offering theoretical convergence benefits and practical gains in large-scale neural network training. However, they incur substantial memory overhead due to the storage demands of non-diagonal preconditioning matrices. To address this, we introduce 4-bit quantization for Shampoo's preconditioners. We introduce two key methods: First, we apply Cholesky decomposition followed by quantization of the Cholesky factors, reducing memory usage by leveraging their lower triangular structure while better preserving spectral properties to minimize information loss. To our knowledge, this is the first quantization approach applied to Cholesky factors of preconditioners. Second, we incorporate error feedback in the quantization process, efficiently storing Cholesky factor and error state in the lower and upper triangular parts of the same matrix. Through extensive experiments, we demonstrate that combining Cholesky quantization with error feedback enhances memory efficiency and algorithm performance in large-scale deep-learning tasks. Theoretically, we also provide convergence proofs for quantized Shampoo under both smooth and non-smooth stochastic optimization settings.

Memory-Efficient 4-bit Preconditioned Stochastic Optimization

TL;DR

The paper tackles the memory bottleneck of full-matrix preconditioners in Shampoo by introducing a memory-efficient 4-bit quantization strategy based on Cholesky factorization, complemented with an error-feedback mechanism. It provides rigorous convergence guarantees for both smooth and nonsmooth stochastic nonconvex objectives and demonstrates substantial memory savings with minimal performance loss across CNNs, vision transformers, and large language model pre-training. Cholesky quantization better preserves spectral properties than direct quantization, enabling performance close to the 32-bit baseline while using only 4-bit storage. Overall, the work significantly enhances the practicality of non-diagonal preconditioned optimization for large-scale deep learning under realistic hardware constraints.

Abstract

Preconditioned stochastic optimization algorithms, exemplified by Shampoo, outperform first-order optimizers by offering theoretical convergence benefits and practical gains in large-scale neural network training. However, they incur substantial memory overhead due to the storage demands of non-diagonal preconditioning matrices. To address this, we introduce 4-bit quantization for Shampoo's preconditioners. We introduce two key methods: First, we apply Cholesky decomposition followed by quantization of the Cholesky factors, reducing memory usage by leveraging their lower triangular structure while better preserving spectral properties to minimize information loss. To our knowledge, this is the first quantization approach applied to Cholesky factors of preconditioners. Second, we incorporate error feedback in the quantization process, efficiently storing Cholesky factor and error state in the lower and upper triangular parts of the same matrix. Through extensive experiments, we demonstrate that combining Cholesky quantization with error feedback enhances memory efficiency and algorithm performance in large-scale deep-learning tasks. Theoretically, we also provide convergence proofs for quantized Shampoo under both smooth and non-smooth stochastic optimization settings.

Paper Structure

This paper contains 26 sections, 7 theorems, 50 equations, 4 figures, 11 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminaries
  4. Practical Shampoo
  5. Linear Square Quantization for Compression
  6. Memory-Efficient Shampoo Via Compensated Cholesky Quantization
  7. Quantization for Shampoo Compression
  8. Efficient and Stable Cholesky Quantization
  9. Compensated Cholesky Quantization
  10. Theoretical Analysis
  11. Smooth Nonconvex Training Loss
  12. Nonsmooth Nonconvex Training Loss
  13. Experiments
  14. Test Performance
  15. Memory and Computational Efficiency
  16. ...and 11 more sections

Key Result

Proposition 5.1

For the 4-bit Shampoo in alg:4bit, let $M_k:= (\mathcal{D}(\Bar{C}_k^L) \mathcal{D}(\Bar{C}_k^L)^T + \lambda_{\max}^L \epsilon I_m)^{-1/4}$, if $\left\|M_k\right\|_{{\rm off},\max}\leq C_B$, then its preconditioners hold that where $\left\|\cdot\right\|_{{\rm off},\max}$ is the maximal absolute value of all off-diagonal entries and $n_k$ is the number of rows in $W_k$. Furthermore, if for every r

Figures (4)

  • Figure 1: Comparison of test accuracy and peak memory usage for training ResNet-34 on CIFAR-100 dataset.
  • Figure 2: Efficient storage for Cholesky factor and error state.
  • Figure 3: Eigenvalue frequency of the dequantized preconditioners $\mathcal{D}(\hat{L})$ and $\mathcal{D}(\hat{R})$ of VGG-19 on CIFAR-100 at 50, 100, 150, and 200 training epochs, all eigenvalues are greater than 0.
  • Figure 4: Comparison of training loss and test accuracy (%) for training ResNet-34 on CIFAR-100 and ViT-Small on Tiny-ImageNet. The left figure shows ResNet-34 results, and the right figure shows ViT-Small results.

Theorems & Definitions (13)

  • Proposition 5.1
  • Theorem 5.1
  • Theorem 5.2
  • Proposition B.1
  • proof
  • Proposition B.2
  • proof
  • Theorem B.1
  • proof
  • Definition B.1
  • ...and 3 more