[论文解读] ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training

Jianfei Chen, Lianmin Zheng+5 more|arXiv (Cornell University)|Apr 29, 2021

Advanced Memory and Neural Computing参考文献 43被引用 8

一句话总结

ActNN 提出了一种内存高效的训练框架，通过使用随机量化将激活值压缩至 2 位精度，实现了显著的内存减少，且精度损失可忽略不计。该框架提供了理论上的收敛保证，并采用感知异质性的混合精度量化方法，将激活内存减少 12 倍，同时在 PyTorch 中仅需极少代码修改即可实现高达 14 倍的批量大小提升。

ABSTRACT

The increasing size of neural network models has been critical for improvements in their accuracy, but device memory is not growing at the same rate. This creates fundamental challenges for training neural networks within limited memory environments. In this work, we propose ActNN, a memory-efficient training framework that stores randomly quantized activations for back propagation. We prove the convergence of ActNN for general network architectures, and we characterize the impact of quantization on the convergence via an exact expression for the gradient variance. Using our theory, we propose novel mixed-precision quantization strategies that exploit the activation's heterogeneity across feature dimensions, samples, and layers. These techniques can be readily applied to existing dynamic graph frameworks, such as PyTorch, simply by substituting the layers. We evaluate ActNN on mainstream computer vision models for classification, detection, and segmentation tasks. On all these tasks, ActNN compresses the activation to 2 bits on average, with negligible accuracy loss. ActNN reduces the memory footprint of the activation by 12x, and it enables training with a 6.6x to 14x larger batch size.

研究动机与目标

• 解决尽管 GPU 显存增长有限，但大规模神经网络训练中显存瓶颈日益加剧的问题。
• 为各类神经网络架构的激活压缩训练（ACT）提供通用的理论收敛保证。
• 设计利用特征维度、样本和层之间激活异质性的量化策略，以最小化梯度方差。
• 在现有深度学习框架（如 PyTorch）中实现 ACT 的实际部署，且仅需极少代码修改。
• 在视觉任务中实现高倍率压缩（例如 2 位）的同时，精度损失极小。

提出的方法

• 将 ActNN 中的量化梯度视为全精度梯度的无偏估计器，确保理论收敛性。
• 推导出梯度方差的精确表达式，以量化量化对训练稳定性的影响。
• 引入基于每组特征激活统计信息自适应调整量化级别的分组量化器。
• 设计细粒度的混合精度量化策略，在固定内存预算下最小化梯度方差。
• 将 ActNN 实现为 PyTorch 库，提供即插即用的层替换（如 actnn.Conv2d），实现无缝集成。
• 应用动态量化策略调优，以在内存节省与训练效率之间取得平衡。

实验结果

研究问题

• RQ1 激活压缩训练是否可被理论保证收敛于通用神经网络架构？
• RQ2 量化如何影响梯度方差？是否可精确建模以指导量化设计？
• RQ3 利用激活异质性的混合精度量化策略能否实现更高压缩率并降低精度损失？
• RQ4 2 位激活压缩在包括分割与检测在内的多样化视觉任务中是否有效？
• RQ5 在固定内存约束下，ActNN 能在模型深度、宽度或分辨率方面实现多大程度的扩展？

主要发现

• ActNN 平均将激活值压缩至 2 位，且在图像分类、目标检测与分割任务中精度损失低于 0.5%。
• 在相同 GPU 上，该框架使批量大小提升 6.6 倍至 14 倍，显著提高训练吞吐量。
• 在相同内存预算下，ActNN 可将 ResNet-152 模型扩展至 6.4 倍更深、3.7 倍更宽或 3.1 倍更高分辨率，同时保持原始训练吞吐量的 64% 至 155%。
• ActNN 实现了 12 倍的激活内存占用减少，使单张 GPU 上实现大批次训练成为可能，适用于高分辨率分割与检测任务。
• 在自监督学习中，ActNN 使用 2 位激活在 1024 批次大小下达到 72.65% 的 ImageNet 顶级-1 准确率，与全精度基线模型相当或更优。
• 即使仅使用 1.25 位激活，ActNN 仍能收敛并生成合理结果，优于以往仅限于 4 位压缩的工作。