[论文解读] Training Transformers with 4-bit Integers

Haocheng Xi, Changhao Li+2 more|arXiv (Cornell University)|Jun 21, 2023

Advanced Neural Network Applications被引用 8

一句话总结

论文提出了一种基于 INT4 的变换器训练方法，使用 Hadamard 量化进行前向传播，并通过位分割结合提升分数采样进行反向传播，在多种任务上实现有竞争力的精度和对硬件友好的加速。

ABSTRACT

Quantizing the activation, weight, and gradient to 4-bit is promising to accelerate neural network training. However, existing 4-bit training methods require custom numerical formats which are not supported by contemporary hardware. In this work, we propose a training method for transformers with all matrix multiplications implemented with the INT4 arithmetic. Training with an ultra-low INT4 precision is challenging. To achieve this, we carefully analyze the specific structures of activation and gradients in transformers to propose dedicated quantizers for them. For forward propagation, we identify the challenge of outliers and propose a Hadamard quantizer to suppress the outliers. For backpropagation, we leverage the structural sparsity of gradients by proposing bit splitting and leverage score sampling techniques to quantize gradients accurately. Our algorithm achieves competitive accuracy on a wide range of tasks including natural language understanding, machine translation, and image classification. Unlike previous 4-bit training methods, our algorithm can be implemented on the current generation of GPUs. Our prototypical linear operator implementation is up to 2.2 times faster than the FP16 counterparts and speeds up the training by up to 35.1%.

研究动机与目标

• 以超低精度(INT4)训练变换器以加速计算并降低内存占用的动机。
• 开发前向量化以通过 Hadamard 变换减小激活异常值。
• 利用梯度的结构稀疏性通过位分割和提升分数采样高效量化反向传播。
• 提供与现有 GPU 兼容的实用 INT4 MM 基训练流程。
• 在 NLP、翻译和视觉基准上展示具有竞争力的精度与速度提升。

提出的方法

• 将变换器中的所有线性运算形式化为矩阵乘法(MM)，并应用带有 LSQ 基量化的 INT4 运算。
• 引入 Hadamard 量化(HQ)，通过在量化前用分块对角 Hadamard 矩阵变换激活来抑制前向激活异常值。
• 通过对梯度在反向传播中应用位分割(BS)将梯度分成更高/更低的4 位部分，并使用提升分数采样(LSS)来选择对 MM 计算信息量大的行。
• 在每个 MM 的梯度动态量化的同时保持激活和权重表示为 INT4，必要时对最终结果进行解量化以进行 FP16 累积。
• 提供两种 MM 过程：前向的 HQ-MM 和反向的 LSS-MM，二者结合实现端到端的 INT4 MM 基训练。
• 在 NLP 基准（GLUE、SQuAD、CoNLL、MT）和图像分类（ViT、ImageNet）上评估，并与 FP、INT8、FP4、LSQ+LUQ 基线进行对比。

实验结果

研究问题

• RQ1 是否可以仅使用标准硬件在端到端训练中使用四比特整数运算，而无需自定义数值格式？
• RQ2 如何解决前向激活异常值与反向梯度稀疏性以在超低精度训练中保持精度？
• RQ3 Hadamard 基前向量化和基于提升分数采样的反向传播在 INT4 训练中的精度和效率权衡是什么？

主要发现

• HQ+LSS 在 NLP、MT 和视觉任务上与 FP 和现有的4-bit 训练方法相比具有竞争力的精度。
• 在 GLUE 和 SQuAD 风格评估中，HQ+LSS 相较于 LSQ+LUQ 有所提升，尤其在 Bert-base 与 Bert-large 变体上表现更优。
• 机器翻译（WMT）在 HQ+LSS 下的 BLEU 下降约 1.0%，优于某些超低位设置且与其他设置相当，同时对硬件友好。
• 使用 ViT 模型的图像分类中，与全精度基线相比精度下降很小（通常低于 1-2%），在若干情形下 HQ+LSS 超过 LSQ+LUQ。
• 该方法使当前 GPU 上实现 INT4 MM 成为可能，提供加速（比 FP16 MM 高达 2.2x）和训练速度提升最高可达 35.1%。
• 对 CoLA 的消融研究证实前向 Hadamard 量化在维持精度方面的重要性，而反向传播则从 BS+LSS 中受益。