4-bit Shampoo for Memory-Efficient Network Training
Authors: name: Sike Wang; name: Jia Li; name: Pan Zhou; name: Hua Huang
Published: 2024-05-28 13:02:56

Abstract: 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 demonstrates that our
4-bit Shampoo achieves comparable test accuracy to its 32-bit counterpart while
being more memory-efficient. The source code will be made available.