Shampoo de 4 bits para treinamento de redes com eficiência de memória

4-bit Shampoo for Memory-Efficient Network Training

May 28, 2024

Autores: Sike Wang, Jia Li, Pan Zhou, Hua Huang

cs.AI

Resumo

Os otimizadores de segunda ordem, que mantêm uma matriz denominada pré-condicionador, são superiores aos otimizadores de primeira ordem tanto na teoria quanto na prática. Os estados que formam o pré-condicionador e sua raiz inversa limitam o tamanho máximo dos modelos treinados por otimizadores de segunda ordem. Para resolver isso, a compressão dos estados do otimizador de 32 bits para larguras de bits menores tem mostrado potencial na redução do uso de memória. No entanto, as abordagens atuais se aplicam apenas a otimizadores de primeira ordem. Neste artigo, propomos os primeiros otimizadores de segunda ordem de 4 bits, exemplificados pelo Shampoo de 4 bits, mantendo um desempenho semelhante ao das versões de 32 bits. Mostramos que a quantização da matriz de autovetores do pré-condicionador no Shampoo de 4 bits é notavelmente melhor do que a quantização do próprio pré-condicionador, tanto teoricamente quanto experimentalmente. Ao corrigir a ortogonalidade da matriz de autovetores quantizada, aprimoramos a aproximação da matriz de autovetores do pré-condicionador, o que também beneficia o cálculo de sua raiz inversa de 4ª ordem. Além disso, descobrimos que a quantização linear quadrada supera ligeiramente a quantização dinâmica em árvore ao quantizar os estados dos otimizadores de segunda ordem. A avaliação em várias redes para classificação de imagens demonstra que nosso Shampoo de 4 bits alcança uma precisão de teste comparável à sua contraparte de 32 bits, sendo mais eficiente em termos de memória. O código-fonte será disponibilizado.

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.