4 位 Shampoo 用于内存高效网络训练

二阶优化器维持一种称为预条件器的矩阵，相较于一阶优化器在理论和实践中表现更佳。构成预条件器及其逆根的状态限制了二阶优化器训练的模型最大规模。为此，将 32 位优化器状态压缩到更低的位宽以减少内存占用已显示出前景。然而，现有方法仅适用于一阶优化器。在本文中，我们提出了首个 4 位二阶优化器，以 4 位 Shampoo 为例，保持与 32 位优化器相似的性能。我们表明，对预条件器特征向量矩阵进行 4 位量化，比对预条件器本身进行量化在理论和实验上都显著更好。通过纠正量化特征向量矩阵的正交性，我们增强了对预条件器特征向量矩阵的逼近，这也有利于其逆 4 次根的计算。此外，我们发现线性平方量化在量化二阶优化器状态时略优于动态树量化。针对 various 用于图像分类和自然语言建模的网络进行评估，证明我们的 4 位 Shampoo 在保持与 32 位对应物相似性能的同时，具有更高的内存效率。