Шампунь на 4 бита для эффективного обучения нейронных сетей.

4-bit Shampoo for Memory-Efficient Network Training

May 28, 2024

Авторы: Sike Wang, Jia Li, Pan Zhou, Hua Huang

cs.AI

Аннотация

Оптимизаторы второго порядка, поддерживающие матрицу, называемую предобуславливателем, превосходят оптимизаторы первого порядка как в теории, так и на практике. Состояния, образующие предобуславливатель и его обратный корень, ограничивают максимальный размер моделей, обучаемых оптимизаторами второго порядка. Для решения этой проблемы сжатие состояний оптимизатора с 32-битной точностью до меньшей разрядности показало перспективы в снижении использования памяти. Однако текущие подходы касаются только оптимизаторов первого порядка. В данной статье мы предлагаем первые 4-битные оптимизаторы второго порядка, в частности 4-битный Shampoo, обеспечивающие производительность, сходную с 32-битными. Мы показываем, что квантование матрицы собственных векторов предобуславливателя в 4-битном Shampoo значительно лучше, чем квантование самого предобуславливателя как с теоретической, так и с экспериментальной точек зрения. Путем исправления ортогональности квантованной матрицы собственных векторов мы улучшаем приближение матрицы собственных векторов предобуславливателя, что также положительно сказывается на вычислении его обратного четвертого корня. Кроме того, мы обнаружили, что линейное квантование незначительно превосходит динамическое дерево квантование при квантовании состояний оптимизатора второго порядка. Оценка на различных сетях для классификации изображений показывает, что наш 4-битный Shampoo достигает сравнимой точности тестирования с его 32-битным аналогом, при этом более эффективен с точки зрения использования памяти. Исходный код будет доступен.

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.