4-bit Shampoo for Memory-Efficient Network Training
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.
665d3bf0957df09a072cd3fa
Sike Wang
665d3bf0957df09a072cd3fb
Jia Li
665d3bf0957df09a072cd3fc
Pan Zhou
665d3bf0957df09a072cd3fd
Hua Huang
4-bit Shampoo for Memory-Efficient Network Training
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.
4-bit quantized Shampoo, a second-order optimizer, achieves comparable performance to 32-bit Shampoo while using less memory, by optimizing the eigenvector matrix quantization and enhancing orthogonality.
second-order optimizers preconditioner orthogonal eigenvector matrix quantization Shampoo linear square quantization dynamic tree quantization image classification