4-Bit-Shampoo für speichereffizientes Netzwerktraining

4-bit Shampoo for Memory-Efficient Network Training

May 28, 2024

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

cs.AI

Zusammenfassung

Optimierer zweiter Ordnung, die eine Matrix namens Präkonditionierer beibehalten, sind sowohl in der Theorie als auch in der Praxis den Optimierern erster Ordnung überlegen. Die Zustände, die den Präkonditionierer bilden, und sein inverser Wurzelbeschränken die maximale Größe von Modellen, die von Optimierern zweiter Ordnung trainiert werden. Um dies zu lösen, hat sich die Komprimierung von 32-Bit-Optimiererzuständen auf niedrigere Bitbreiten als vielversprechend erwiesen, um den Speicherverbrauch zu reduzieren. Allerdings beziehen sich aktuelle Ansätze nur auf Optimierer erster Ordnung. In diesem Artikel schlagen wir die ersten 4-Bit-Optimierer zweiter Ordnung vor, exemplarisch dargestellt durch 4-Bit Shampoo, die eine ähnliche Leistung wie die 32-Bit-Optimierer beibehalten. Wir zeigen, dass die Quantisierung der Eigenvektormatrix des Präkonditionierers in 4-Bit Shampoo sowohl theoretisch als auch experimentell deutlich besser ist als die Quantisierung des Präkonditionierers selbst. Durch die Korrektur der Orthogonalität der quantisierten Eigenvektormatrix verbessern wir die Approximation der Eigenvektormatrix des Präkonditionierers, was auch die Berechnung seiner inversen 4. Wurzel begünstigt. Darüber hinaus stellen wir fest, dass die lineare Quadratquantisierung bei der Quantisierung von Zuständen von Optimierern zweiter Ordnung leicht besser abschneidet als die dynamische Baumquantisierung. Die Evaluation an verschiedenen Netzwerken für die Bildklassifizierung zeigt, dass unser 4-Bit Shampoo eine vergleichbare Testgenauigkeit wie sein 32-Bit Pendant erreicht, während es speichereffizienter ist. Der Quellcode wird verfügbar gemacht.

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.