[论文解读] Memory Efficient Optimizers with 4-bit States

Bingrui Li, Jianfei Chen+1 more|arXiv (Cornell University)|Sep 4, 2023

Machine Learning and Data ClassificationComputer Science被引用 3

一句话总结

本文通过为Adam类优化器的一阶和二阶矩引入新型量化技术，提出了一种用于内存高效深度学习训练的4位优化器。通过采用更小的块大小、无零点的线性量化器处理二阶矩，以及秩-1归一化，该方法在多种自然语言处理、视觉和指令微调任务中，相较于8位基线模型，实现了50%的内存减少，同时保持了全精度收敛的准确性。

ABSTRACT

Optimizer states are a major source of memory consumption for training neural networks, limiting the maximum trainable model within given memory budget. Compressing the optimizer states from 32-bit floating points to lower bitwidth is promising to reduce the training memory footprint, while the current lowest achievable bitwidth is 8-bit. In this work, we push optimizer states bitwidth down to 4-bit through a detailed empirical analysis of first and second moments. Specifically, we find that moments have complicated outlier patterns, that current block-wise quantization cannot accurately approximate. We use a smaller block size and propose to utilize both row-wise and column-wise information for better quantization. We further identify a zero point problem of quantizing the second moment, and solve this problem with a linear quantizer that excludes the zero point. Our 4-bit optimizers are evaluated on a wide variety of benchmarks including natural language understanding, machine translation, image classification, and instruction tuning. On all the tasks our optimizers can achieve comparable accuracy with their full-precision counterparts, while enjoying better memory efficiency.

研究动机与目标

• 解决由高精度优化器状态（如32位Adam状态）引起的大型神经网络训练中的内存瓶颈问题。
• 将优化器状态的位宽从当前的8位最低值进一步降低至4位，以进一步压缩内存使用。
• 克服现有分块量化方法在捕捉一阶和二阶矩中复杂异常值模式方面的局限性。
• 解决二阶矩量化中的零点问题，该问题会导致参数更新方向出现显著偏差。
• 开发一种鲁棒的4位优化器，无需依赖稳定嵌入层即可保持收敛速度和准确性。

提出的方法

• 提出更小的量化块大小，以更好地捕捉一阶矩中的局部异常值模式。
• 引入无零点的线性量化器处理二阶矩，防止非零值被量化为零，从而避免更新方向的失真。
• 设计秩-1归一化，通过更有效地建模异常值结构，提升二阶矩的近似效果。
• 将4位量化的一阶矩与分解形式的二阶矩结合，构建一种混合优化器以增强内存效率。
• 将所提出的量化器集成到基于压缩的优化框架中，实现每轮迭代对优化器状态的压缩与解压缩。
• 通过实证分析验证量化器设计，并确保在4位精度下实现无损收敛。

Figure 1 : Visualization of the first moment in the layers.3.blocks.1.mlp.fc1 layer in a Swin-T model. (a): Magnitude of the first moment. (b): Histogram of the first moment. (c): Moment approximated by B128/DE . (d): Moment approximated by B2048/DE .

实验结果

研究问题

• RQ1 优化器状态能否被压缩至4位而不损失模型收敛性或准确性？
• RQ2 在4位压缩一阶和二阶矩时，特别是针对异常值模式，会面临哪些量化挑战？
• RQ3 二阶矩量化中的零点问题如何影响优化稳定性，是否可以被缓解？
• RQ4 所提出的量化技术能否在不依赖稳定嵌入等架构改进的前提下，实现与全精度优化器相当的性能？
• RQ5 与现有8位基线相比，4位优化器在多样化任务中的内存和训练效率增益如何？

主要发现

• 所提出的4位优化器在所有评估基准（包括自然语言处理、机器翻译、图像分类和指令微调）上，实现了与全精度Adam相当的收敛速度和最终准确性。
• 与8位优化器相比，该方法将内存消耗减少了50%，在现有低精度优化器中实现了最低的内存占用。
• 无零点的线性量化器在处理二阶矩时，有效防止了参数更新中的大幅偏差，从而实现了4位精度下的稳定训练。
• 秩-1归一化显著提升了二阶矩的近似效果，尤其在存在复杂异常值模式的情况下。
• 4位优化器在无需稳定嵌入层的情况下实现了无损收敛，展现出在不同模型架构上的鲁棒性与泛化能力。
• 通过优化器卸载（optimizer offloading），4位优化器因低精度状态带来的通信成本降低，在语言模型微调中提升了吞吐量。

Figure 2 : Outlier patterns vary across two first moment tensors. (a): outliers lie in fixed rows (dimension 0). (b): outliers lie in fixed columns (dimension 1).