🤖 AI Summary
This work addresses the prohibitively high memory overhead of second-order optimizers, which hinders their scalability to ultra-large models. To this end, we propose 4-bit Shampoo—the first 4-bit precision second-order optimizer. Our method avoids direct quantization of high-dimensional preconditioning matrices; instead, it quantizes the eigenvector matrix and introduces an orthogonality correction to preserve numerical accuracy during inverse fourth-root computation. We employ a 4-bit linear-squared quantization scheme tailored for second-order statistics. Experiments on image classification and language modeling tasks demonstrate that 4-bit Shampoo matches the convergence and accuracy of full-precision (32-bit) Shampoo while reducing GPU memory consumption by 8× and significantly improving training memory efficiency. This represents the first practical breakthrough enabling high-accuracy, low-bitwidth second-order optimization at scale.
📝 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 and natural language modeling demonstrates that our 4-bit Shampoo achieves comparable performance to its 32-bit counterpart while being more memory-efficient.