用池化压缩反向传播激活图，降低显存占用，让小GPU也能训练大模型。
Less Memory Means smaller GPUs: Backpropagation with Compressed Activations
- 反向传播时对激活图进行池化压缩，前向计算保持不变
- 峰值显存减少29%，训练时间稍长但精度不变
- 适合资源有限却想训练大模型的研究者或开发者
深度神经网络规模持续增长，导致计算资源需求急剧上升。当前主流的大语言模型等架构需依赖数千个加速器（如GPU/TPU）的超算系统进行训练。除了海量浮点运算，模型的内存占用也迅速膨胀。而GPU架构本身显存有限，甚至一些EfficientNet变体也无法在单张消费级显卡上以合理小批量大小训练。训练过程中，中间输入激活值必须存储至反向传播阶段用于梯度计算，占据绝大部分内存。本文提出在反向传播中通过池化压缩激活图，从而减少内存占用和数据移动量。前向计算仍保持原始精度。我们以经典视觉架构ResNet为例，实证验证了收敛性，并研究其对特征检测的影响。结果表明，该方法可将峰值内存消耗降低29%，代价是训练周期略长，但预测精度与未压缩基线相当。
原文摘要 · Abstract (English)
The ever-growing scale of deep neural networks (DNNs) has lead to an equally rapid growth in computational resource requirements. Many recent architectures, most prominently Large Language Models, have to be trained using supercomputers with thousands of accelerators, such as GPUs or TPUs. Next to the vast number of floating point operations the memory footprint of DNNs is also exploding. In contrast, GPU architectures are notoriously short on memory. Even comparatively small architectures like some EfficientNet variants cannot be trained on a single consumer-grade GPU at reasonable mini-batch sizes. During training, intermediate input activations have to be stored until backpropagation for gradient calculation. These make up the vast majority of the memory footprint. In this work we therefore consider compressing activation maps for the backward pass using pooling, which can reduce both the memory footprint and amount of data movement. The forward computation remains uncompressed. We empirically show convergence and study effects on feature detection at the example of the common vision architecture ResNet. With this approach we are able to reduce the peak memory consumption by 29% at the cost of a longer training schedule, while maintaining prediction accuracy compared to an uncompressed baseline.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。