ActNN：通过 2 位激活压缩训练降低训练内存占用
机器学习
2021-07-07 v2 计算机视觉与模式识别
机器学习
摘要
神经网络模型规模的不断增长对其精度提升至关重要，但设备内存的增长速度并未与之同步。这在有限内存环境下训练神经网络时带来了根本性挑战。本文提出 ActNN，一种高效内存训练框架，其通过存储随机量化的激活值用于反向传播。我们证明了 ActNN 在通用网络架构上的收敛性，并通过梯度方差的精确表达式刻画了量化对收敛的影响。基于该理论，我们提出了利用激活在特征维度、样本和层之间异质性的新型混合精度量化策略。这些技术只需替换层，便可轻松应用于 PyTorch 等现有动态图框架。我们在分类、检测和分割任务的主流计算机视觉模型上评估了 ActNN。在所有这些任务中，ActNN 将激活平均压缩至 2 位，且精度损失可忽略。ActNN 将激活的内存占用降低 12 倍，并支持使用大 6.6 倍至 14 倍的批量大小进行训练。
引用
@article{arxiv.2104.14129,
title = {ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training},
author = {Jianfei Chen and Lianmin Zheng and Zhewei Yao and Dequan Wang and Ion Stoica and Michael W. Mahoney and Joseph E. Gonzalez},
journal= {arXiv preprint arXiv:2104.14129},
year = {2021}
}
备注
to be published in ICML 2021