4-bit Shampoo for Memory-Efficient Network Training

二阶优化器通过维护一个被称为预处理器的矩阵，理论和实践上都优于一阶优化器。预处理器及其逆根的状态限制了二阶优化器训练模型的最大规模。为了解决这个问题，将32位优化器状态压缩到更低的位宽已经显示出降低内存使用的潜力。然而，目前的方法仅适用于一阶优化器。在本文中，我们提出了第一个4位二阶优化器，即4位Shampoo，其性能与32位优化器相似。我们发现，在4位Shampoo中，量化预处理器的特征向量矩阵比量化预处理器本身从理论和实验上都更好。通过矫正量化特征向量矩阵的正交性，我们增强了预处理器特征向量矩阵的逼近度，这也有利于计算其逆4次方根。此外，我们发现，在量化二阶优化器状态时，线性平方量化略优于动态树量化。在各种用于图像分类的网络上进行的评估表明，我们的4位Shampoo在更节省内存的同时，实现了与32位优化器相当的测试精度。源代码将提供。

本文旨在解决二阶优化器在训练大型模型时预处理矩阵状态过大的问题，提出了一种四位二阶优化器的解决方案。

本文提出了一种四位二阶优化器的解决方案，即将预处理矩阵的特征向量矩阵量化为四位，通过修正正交性来增强预处理矩阵的逼近性，同时提高其逆四次方的计算效率。

本文提出的四位Shampoo在多个图像分类网络上的实验中，与32位优化器相比，具有相似的测试精度，同时更加节省内存。此外，本文还探讨了量化二阶优化器状态时的线性平方量化和动态树量化，发现线性平方量化略优于动态树量化。

在最近的相关研究中，也有一些关于量化优化器状态的工作，如《Training Deep Neural Networks with 8-bit Floating Point Numbers》和《Training Quantized Nets: A Deeper Understanding》等。