4-bit Shampoo voor geheugenefficiënte netwerktraining

4-bit Shampoo for Memory-Efficient Network Training

May 28, 2024

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

cs.AI

Samenvatting

Tweedegraads-optimalisatoren, die een matrix onderhouden die een preconditioner wordt genoemd, zijn zowel in theorie als praktijk superieur aan eerstegraads-optimalisatoren. De toestanden die de preconditioner en zijn inverse wortel vormen, beperken de maximale grootte van modellen die door tweedegraads-optimalisatoren getraind kunnen worden. Om dit aan te pakken, heeft het comprimeren van 32-bits optimalisatietoestanden naar lagere bitbreedtes veelbelovende resultaten laten zien in het verminderen van geheugengebruik. Huidige benaderingen zijn echter alleen van toepassing op eerstegraads-optimalisatoren. In dit artikel stellen we de eerste 4-bits tweedegraads-optimalisatoren voor, geïllustreerd door 4-bits Shampoo, die een vergelijkbare prestaties behouden als hun 32-bits tegenhangers. We tonen aan dat het kwantiseren van de eigenvectormatrix van de preconditioner in 4-bits Shampoo zowel theoretisch als experimenteel aanzienlijk beter is dan het kwantiseren van de preconditioner zelf. Door de orthogonaliteit van de gekwantiseerde eigenvectormatrix te corrigeren, verbeteren we de benadering van de eigenvectormatrix van de preconditioner, wat ook ten goede komt aan de berekening van zijn inverse 4-de wortel. Daarnaast vinden we dat lineaire kwadratische kwantisering lichtelijk beter presteert dan dynamische boomkwantisering bij het kwantiseren van tweedegraads-optimalisatietoestanden. Evaluatie op verschillende netwerken voor beeldclassificatie toont aan dat onze 4-bits Shampoo een vergelijkbare testnauwkeurigheid bereikt als zijn 32-bits tegenhanger, terwijl het geheugenefficiënter is. De broncode zal beschikbaar worden gesteld.

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.