ICML 2024 | 梯度检查点太慢？不降速、省显存，LowMemoryBP大幅提升反向传播显存效率
内容提要
南开大学研究团队提出了两种反向传播改进策略，Approx-BP和MS-BP，旨在降低Transformer模型微调时的显存占用。实验表明，这些方法能有效减少20%至30%的显存，同时保持训练速度和效果不变，为优化神经网络的显存使用提供了新思路。
延伸解读
与梯度检查点的本质区别
梯度检查点通过牺牲训练速度来节省显存，而LowMemoryBP则从反向传播原理入手，不增加计算量。文章指出，LowMemoryBP在降低20%~30%峰值显存的同时，完全保持了训练速度，这与梯度检查点形成鲜明对比。对于需要频繁微调大模型的场景，这种不降速的显存优化更具实用价值。
两种策略的适用场景
Approx-BP针对激活函数，用分段线性函数逼近GELU/SiLU，使激活存储仅需2bit；MS-BP则针对LayerNorm/RMSNorm，通过共享激活张量消除冗余。两者分别从激活函数和标准化层入手，可独立使用也可结合。读者可根据模型结构中显存瓶颈的具体位置，选择或组合这两种策略。
与8-bit压缩方法的对比
文章将LowMemoryBP与Mesa（一种8-bit激活压缩训练方法）对比，Mesa会使训练速度降低约20%，而LowMemoryBP则完全保持训练速度。这表明LowMemoryBP在显存节省的同时，避免了压缩方法常见的速度损失，对于训练效率敏感的任务更具优势。
理论突破与开源实践
Approx-BP理论突破了传统神经网络优化框架，为使用非配对导数提供了理论可行性，其导出的ReGELU2和ReSiLU2展示了实践价值。相关模块已在GitHub开源，便于研究者复现和集成到现有训练流程中，降低了应用门槛。
Q&A
LowMemoryBP方法的主要目标是什么？
LowMemoryBP方法旨在降低Transformer模型微调时的显存占用，同时保持训练速度和效果不变。
Approx-BP和MS-BP分别是什么？
Approx-BP通过分段线性函数逼近激活函数来减少显存占用，MS-BP则通过共享激活张量来降低冗余。
使用LowMemoryBP方法可以降低多少显存？
使用LowMemoryBP方法可以将Transformer模型的显存占用降低20%至30%。
LowMemoryBP方法对训练速度有影响吗？
LowMemoryBP方法不会影响训练速度，保持了训练吞吐量。
Approx-BP理论的创新之处是什么？
Approx-BP理论突破了传统神经网络的优化框架，为使用非配对导数提供了理论可行性。
实验结果显示了什么？
实验结果表明，LowMemoryBP方法在多个模型微调中显著降低了显存占用，同时没有损失训练效果。
