메모리 효율적인 네트워크 학습을 위한 4비트 샴푸

4-bit Shampoo for Memory-Efficient Network Training

May 28, 2024

저자: Sike Wang, Jia Li, Pan Zhou, Hua Huang

cs.AI

초록

2차 최적화 기법은 프리컨디셔너(preconditioner)라는 행렬을 유지하며, 이론과 실제 모두에서 1차 최적화 기법보다 우수합니다. 그러나 프리컨디셔너와 그 역행렬의 제곱근을 구성하는 상태들은 2차 최적화 기법으로 훈련할 수 있는 모델의 최대 크기를 제한합니다. 이를 해결하기 위해 32비트 최적화 상태를 더 낮은 비트폭으로 압축하는 방법이 메모리 사용량을 줄이는 데 유망한 결과를 보여주었습니다. 하지만 현재의 접근법들은 1차 최적화 기법에만 적용되고 있습니다. 본 논문에서는 4비트 Shampoo를 예시로, 32비트 버전과 유사한 성능을 유지하면서도 4비트 2차 최적화 기법을 최초로 제안합니다. 우리는 4비트 Shampoo에서 프리컨디셔너의 고유벡터 행렬을 양자화하는 것이 프리컨디셔너 자체를 양자화하는 것보다 이론적 및 실험적으로 훨씬 우수함을 보여줍니다. 양자화된 고유벡터 행렬의 직교성을 보정함으로써, 프리컨디셔너의 고유벡터 행렬 근사를 개선하고, 이는 역 4제곱근 계산에도 이점을 제공합니다. 또한, 2차 최적화 상태를 양자화할 때 선형 제곱 양자화가 동적 트리 양자화보다 약간 더 나은 성능을 보임을 발견했습니다. 이미지 분류를 위한 다양한 네트워크에서의 평가 결과, 우리의 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.