Shampoo 4 bits pour un entraînement de réseau économe en mémoire

4-bit Shampoo for Memory-Efficient Network Training

May 28, 2024

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

cs.AI

Résumé

Les optimiseurs de second ordre, qui maintiennent une matrice appelée préconditionneur, surpassent les optimiseurs de premier ordre à la fois en théorie et en pratique. Les états formant le préconditionneur et sa racine inverse limitent la taille maximale des modèles entraînés par les optimiseurs de second ordre. Pour remédier à cela, la compression des états de l'optimiseur de 32 bits à des largeurs de bits inférieures a montré des résultats prometteurs pour réduire l'utilisation de la mémoire. Cependant, les approches actuelles ne concernent que les optimiseurs de premier ordre. Dans cet article, nous proposons les premiers optimiseurs de second ordre en 4 bits, illustrés par le Shampoo 4 bits, qui maintiennent des performances similaires à celles des versions 32 bits. Nous montrons que la quantification de la matrice des vecteurs propres du préconditionneur dans le Shampoo 4 bits est nettement meilleure que la quantification du préconditionneur lui-même, tant sur le plan théorique qu'expérimental. En rectifiant l'orthogonalité de la matrice des vecteurs propres quantifiée, nous améliorons l'approximation de la matrice des vecteurs propres du préconditionneur, ce qui bénéficie également au calcul de sa racine inverse 4-ième. Par ailleurs, nous constatons que la quantification linéaire carrée surpasse légèrement la quantification dynamique en arbre lors de la quantification des états des optimiseurs de second ordre. L'évaluation sur divers réseaux pour la classification d'images démontre que notre Shampoo 4 bits atteint une précision de test comparable à celle de sa version 32 bits tout en étant plus économe en mémoire. Le code source sera rendu disponible.

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.