Shampoo a 4 bit per un Addestramento Efficiente in Memoria delle Reti

4-bit Shampoo for Memory-Efficient Network Training

May 28, 2024

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

cs.AI

Abstract

Gli ottimizzatori del secondo ordine, che mantengono una matrice denominata precondizionatore, sono superiori agli ottimizzatori del primo ordine sia in teoria che in pratica. Gli stati che formano il precondizionatore e la sua radice inversa limitano la dimensione massima dei modelli addestrati dagli ottimizzatori del secondo ordine. Per affrontare questo problema, la compressione degli stati dell'ottimizzatore da 32 bit a larghezze di bit inferiori ha dimostrato di ridurre l'uso della memoria. Tuttavia, gli approcci attuali riguardano solo gli ottimizzatori del primo ordine. In questo articolo, proponiamo i primi ottimizzatori del secondo ordine a 4 bit, rappresentati da Shampoo a 4 bit, che mantengono prestazioni simili a quelle delle versioni a 32 bit. Dimostriamo che quantizzare la matrice degli autovettori del precondizionatore in Shampoo a 4 bit è notevolmente migliore rispetto a quantizzare il precondizionatore stesso, sia teoricamente che sperimentalmente. Correggendo l'ortogonalità della matrice degli autovettori quantizzata, miglioriamo l'approssimazione della matrice degli autovettori del precondizionatore, il che beneficia anche il calcolo della sua radice inversa alla quarta potenza. Inoltre, scopriamo che la quantizzazione lineare quadrata supera leggermente la quantizzazione dinamica ad albero quando si quantizzano gli stati degli ottimizzatori del secondo ordine. La valutazione su varie reti per la classificazione delle immagini dimostra che il nostro Shampoo a 4 bit raggiunge un'accuratezza di test comparabile alla sua controparte a 32 bit, pur essendo più efficiente in termini di memoria. Il codice sorgente sarà reso disponibile.

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.