[论文解读] Backprop with Approximate Activations for Memory-efficient Network Training

Ayan Chakrabarti, Benjamin Moseley|arXiv (Cornell University)|Jan 23, 2019

Advanced Neural Network Applications参考文献 21被引用 17

一句话总结

本文提出一种内存高效的反向传播方法，通过使用近似激活值来减少训练期间的GPU内存占用，仅在反向传播中存储低精度近似值，同时在前向传播中重新计算精确激活值。该方法在对训练精度影响极小的情况下，将内存使用量最高降低8倍，从而可在深层网络（如1000多层的ResNets）上实现更大的批量大小和更高的GPU利用率，最低支持4位精度。

ABSTRACT

Training convolutional neural network models is memory intensive since back-propagation requires storing activations of all intermediate layers. This presents a practical concern when seeking to deploy very deep architectures in production, especially when models need to be frequently re-trained on updated datasets. In this paper, we propose a new implementation for back-propagation that significantly reduces memory usage, by enabling the use of approximations with negligible computational cost and minimal effect on training performance. The algorithm reuses common buffers to temporarily store full activations and compute the forward pass exactly. It also stores approximate per-layer copies of activations, at significant memory savings, that are used in the backward pass. Compared to simply approximating activations within standard back-propagation, our method limits accumulation of errors across layers. This allows the use of much lower-precision approximations without affecting training accuracy. Experiments on CIFAR-10, CIFAR-100, and ImageNet show that our method yields performance close to exact training, while storing activations compactly with as low as 4-bit precision.

研究动机与目标

• 为解决训练深度卷积神经网络时因存储全精度激活值用于反向传播而导致的高内存占用问题。
• 在不显著增加计算成本或降低模型性能的前提下减少内存使用。
• 通过最小化激活值存储，使在单张GPU上训练极深网络（如1001层ResNets）成为可能。
• 探究随机梯度下降在反向传播中对激活值近似误差的鲁棒性。
• 证明通过一种新型缓冲区复用策略，可有效控制近似误差，使其在各层之间不会显著累积。

提出的方法

• 该方法在前向传播期间使用可复用的缓冲区临时存储精确激活值，使用后立即覆盖以节省内存。
• 为每层保留一个低精度（如4位）的激活值近似副本，用于反向传播。
• 前向传播使用精确激活值计算梯度，确保准确性；反向传播则使用近似副本以减少内存占用。
• 通过复用同一缓冲区存储精确激活值并仅保存压缩后的近似副本，该方法避免了在所有层中存储完整激活值。
• 通过确保前向传播保持精确，该方法防止了误差在各层之间的累积，从而将反向传播中近似带来的影响限制在可接受范围内。
• 该方法利用SGD固有的噪声来吸收因激活值近似而引入的微小额外误差，从而具备良好的鲁棒性。

实验结果

研究问题

• RQ1 即使使用如4位这样的极低精度表示，是否可以在反向传播中使用近似激活值而不显著降低模型精度？
• RQ2 如何控制近似误差在各层之间的传播，以防止在反向传播过程中误差累积？
• RQ3 在保持训练效率和模型性能的前提下，训练期间内存使用量最多可降低多少？
• RQ4 激活值近似引入的误差是否与随机梯度下降的固有噪声相当？
• RQ5 该方法是否能够使在单张GPU上训练极深网络（如1001层ResNets）并使用大批次成为可能？

主要发现

• 与精确训练相比，该方法将内存使用量最高降低了8倍，使在单张1080Ti GPU上训练1001层ResNets成为可能。
• 使用4位近似时，该方法在CIFAR-10、CIFAR-100和ImageNet上的测试精度与精确训练相比仅相差0.5%以内。
• 参数梯度中的近似误差比SGD的固有噪声小1至2个数量级，解释了其对性能影响极小的原因。
• 该方法使更大的批量大小成为可能，提升了GPU利用率并减少了每样本的训练耗时。
• 该方法在每次迭代的运行时间与精确训练几乎完全相同，表明计算开销可忽略不计。
• 在相同单张1080Ti GPU设置下，该方法可在CIFAR-10上实现完整的128批次训练，而精确训练则需每GPU使用64批次并执行两次前向-反向传播。