[论文解读] Few-Bit Backward: Quantized Gradients of Activation Functions for Memory Footprint Reduction

Georgii Novikov, Daniel Bershatsky+4 more|arXiv (Cornell University)|Feb 1, 2022

Advanced Neural Network Applications被引用 5

一句话总结

该论文提出通过分段常数近似方法对逐点激活函数的梯度进行量化，以减少深度学习训练期间的内存占用。通过在反向传播中仅存储3位索引而非全精度输入，该方法在基准测试中实现了高达40%的内存减少，且准确率无损失，利用动态规划计算最优量化方案，作为ReLU、GELU和Swish等常见激活函数的即插即用替代方案。

ABSTRACT

Memory footprint is one of the main limiting factors for large neural network training. In backpropagation, one needs to store the input to each operation in the computational graph. Every modern neural network model has quite a few pointwise nonlinearities in its architecture, and such operation induces additional memory costs which -- as we show -- can be significantly reduced by quantization of the gradients. We propose a systematic approach to compute optimal quantization of the retained gradients of the pointwise nonlinear functions with only a few bits per each element. We show that such approximation can be achieved by computing optimal piecewise-constant approximation of the derivative of the activation function, which can be done by dynamic programming. The drop-in replacements are implemented for all popular nonlinearities and can be used in any existing pipeline. We confirm the memory reduction and the same convergence on several open benchmarks.

研究动机与目标

• 减少深度神经网络训练中的内存占用，特别是反向传播期间存储的激活张量内存占用。
• 解决现代架构中存储逐点非线性函数全精度输入所被忽视的内存成本问题。
• 开发一种与数据无关、即插即用的激活函数替代方法，实现低比特梯度存储，且无需修改训练流水线。
• 在多种架构和任务中实现显著的内存节省，同时保持模型收敛性和准确率。
• 提供一种可推广的框架，适用于任何模型（包括Transformer），无需重新训练或修改网络结构。

提出的方法

• 使用动态规划计算激活函数导数的最优分段常数近似，以最小化量化误差。
• 在反向传播中仅用分箱索引（例如3位）替代全精度激活输入，显著降低每个张量元素的内存占用。
• 仅存储量化后的导数值，而非原始激活输入，实现有损压缩，对梯度影响可忽略不计。
• 实现常见激活函数（ReLU、GELU、Swish、Sigmoid）的即插即用替代，无缝集成至现有PyTorch流水线。
• 在模型中普遍应用该方法，包括Transformer和EfficientNet，且不依赖数据分布或模型结构。
• 使用CUDA优化内核，实现高效推理与训练，性能开销极低。

实验结果

研究问题

• RQ1 将逐点激活函数的梯度量化至少数比特是否能减少反向传播期间的内存使用，且不降低模型准确率？
• RQ2 何种分段常数近似方案能最小化低比特表示下的量化误差，从而最优逼近激活函数的导数？
• RQ3 在实际应用中，对Transformer、EfficientNet和GPT-2等多样化模型，使用3位或4位量化激活梯度能实现多大程度的内存减少？
• RQ4 所提方法能否作为即插即用方案应用于不同模型和训练阶段，包括微调预训练模型？
• RQ5 鉴于其与数据无关的设计，该方法在不同数据分布和模型架构下是否仍保持有效性？

主要发现

• 通过每个激活元素仅存储3位索引，该方法在GPT2和RoBERTa-large等模型上实现了高达40%的内存减少。
• 在ResNet、EfficientNet和基于BERT的模型等多个基准测试中，3位量化平均实现20%的内存节省。
• 该方法在所有评估模型（包括ResNet-50、EfficientNet-B7和RoBERTa）上均保持与全精度训练相同的训练收敛性和测试准确率。
• 通过动态规划计算的激活函数导数最优分段常数近似在所有测试激活函数（包括ReLU、GELU、Swish和Sigmoid）上均有效。
• 该方法与现有训练流水线兼容，可应用于任意阶段（包括微调预训练模型），且无需修改网络结构。
• CUDA实现具有极低开销，支持在实际训练工作流中高效部署，性能影响可忽略。