メモリ効率の良いネットワーク訓練のための4ビットShampoo

4-bit Shampoo for Memory-Efficient Network Training

May 28, 2024

著者: Sike Wang, Jia Li, Pan Zhou, Hua Huang

cs.AI

要旨

第二階最適化手法は、プレコンディショナーと呼ばれる行列を保持するため、理論的にも実践的にも第一階最適化手法を凌駕しています。しかし、プレコンディショナーとその逆平方根を形成する状態量は、第二階最適化手法で訓練可能なモデルの最大サイズを制限します。この問題に対処するため、32ビットの最適化状態量を低ビット幅に圧縮することでメモリ使用量を削減する手法が有望視されています。しかし、現在のアプローチは第一階最適化手法に限定されています。本論文では、4ビットShampooを例とした初の4ビット第二階最適化手法を提案し、32ビット版と同等の性能を維持します。理論的にも実験的にも、プレコンディショナー自体を量子化するよりも、4ビットShampooにおけるプレコンディショナーの固有ベクトル行列を量子化する方が著しく優れていることを示します。量子化された固有ベクトル行列の直交性を補正することで、プレコンディショナーの固有ベクトル行列の近似精度を向上させ、その逆4乗根の計算にも寄与します。さらに、第二階最適化手法の状態量を量子化する際、線形二乗量子化が動的木量子化をわずかに上回ることを発見しました。画像分類のための様々なネットワークでの評価により、我々の4ビットShampooが32ビット版と同等のテスト精度を達成しつつ、よりメモリ効率が高いことが実証されました。ソースコードは公開予定です。

English

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.