Shampoo de 4 bits para el entrenamiento eficiente en memoria de redes neuronales

4-bit Shampoo for Memory-Efficient Network Training

May 28, 2024

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

cs.AI

Resumen

Los optimizadores de segundo orden, que mantienen una matriz denominada precondicionador, son superiores a los optimizadores de primer orden tanto en teoría como en la práctica. Los estados que forman el precondicionador y su raíz inversa limitan el tamaño máximo de los modelos entrenados por optimizadores de segundo orden. Para abordar esto, la compresión de los estados del optimizador de 32 bits a anchos de bits más bajos ha mostrado ser prometedora para reducir el uso de memoria. Sin embargo, los enfoques actuales solo se aplican a optimizadores de primer orden. En este artículo, proponemos los primeros optimizadores de segundo orden de 4 bits, ejemplificados por Shampoo de 4 bits, que mantienen un rendimiento similar al de los de 32 bits. Demostramos que cuantizar la matriz de vectores propios del precondicionador en Shampoo de 4 bits es notablemente mejor que cuantizar el precondicionador en sí, tanto teórica como experimentalmente. Al rectificar la ortogonalidad de la matriz de vectores propios cuantizada, mejoramos la aproximación de la matriz de vectores propios del precondicionador, lo que también beneficia el cálculo de su raíz inversa de cuarto orden. Además, encontramos que la cuantización lineal cuadrática supera ligeramente a la cuantización dinámica en árbol al cuantizar los estados de los optimizadores de segundo orden. La evaluación en varias redes para clasificación de imágenes demuestra que nuestro Shampoo de 4 bits logra una precisión en pruebas comparable a su contraparte de 32 bits, siendo más eficiente en memoria. El código fuente estará 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.