[论文解读] Memory-Efficient Backpropagation Through Time

Audrūnas Gruslys, Rémi Munos+3 more|arXiv (Cornell University)|Jun 10, 2016

Neural Networks and Applications参考文献 11被引用 51

一句话总结

本文提出了一种基于动态规划的算法 BPTT-MSM，可在循环神经网络的反向传播时间（BPTT）中最优平衡内存使用与计算成本。通过智能地缓存和重新计算中间状态，该方法在序列长度为1000时，将内存使用量最多减少95%，同时计算时间仅比标准BPTT增加三分之一，从而在严格内存限制下实现高效训练。

ABSTRACT

We propose a novel approach to reduce memory consumption of the backpropagation through time (BPTT) algorithm when training recurrent neural networks (RNNs). Our approach uses dynamic programming to balance a trade-off between caching of intermediate results and recomputation. The algorithm is capable of tightly fitting within almost any user-set memory budget while finding an optimal execution policy minimizing the computational cost. Computational devices have limited memory capacity and maximizing a computational performance given a fixed memory budget is a practical use-case. We provide asymptotic computational upper bounds for various regimes. The algorithm is particularly effective for long sequences. For sequences of length 1000, our algorithm saves 95\% of memory usage while using only one third more time per iteration than the standard BPTT.

研究动机与目标

• 为解决标准反向传播时间（BPTT）在循环神经网络中高内存消耗的问题，尤其是在 GPU 等内存受限设备上。
• 开发一种方法，为任意固定的用户指定内存预算，找到内存使用与计算成本之间的最优权衡。
• 通过紧密适配任意内存约束，实现对长序列的高效训练，而无需依赖启发式内存节省策略。
• 提供一种通用且与架构无关的解决方案，其内存效率优于现有启发式方法（如 Chen 的 √t 算法）。

提出的方法

• 该方法使用动态规划计算在 BPTT 过程中缓存和重新计算中间状态的最优策略，以在固定内存预算下最小化总计算成本。
• 将问题建模为一系列时间步，其中每个状态要么被缓存，要么被重新计算，成本由前向传播操作定义。
• 算法定义了一个代价函数 Q_i(t,m)，用于计算在拥有 m 个内存槽的情况下，从时间 i 回溯到起点的最小代价。
• 采用递归公式，考虑所有可能的序列分割方式，并在各段之间平衡内存使用，确保全局最优性。
• 通过允许对缓存的中间状态数量进行细粒度控制，支持任意内存预算。
• 设计上与 RNN 架构无关，适用于标准 RNN、LSTM 及其他循环模型。

实验结果

研究问题

• RQ1 动态规划方法能否在 RNN 的 BPTT 中找到内存使用与计算成本之间的全局最优权衡？
• RQ2 与 Chen 的 √t 算法等启发式方法相比，所提方法在内存效率和计算成本方面表现如何？
• RQ3 在长序列中，内存使用量可被减少到何种程度，同时仅带来计算成本的轻微增加？
• RQ4 该方法是否可调节以适配任何用户定义的内存预算，包括现有启发式方法不支持的预算？

主要发现

• 对于长度为 1000 的序列，所提出的 BPTT-MSM 算法相比标准 BPTT 将内存使用量减少了 95%，同时计算时间仅增加三分之一。
• 该算法在类似 Chen 的 √t 方法所用内存预算下实现了接近最优的性能，但内存效率显著更高。
• 当计算成本固定为每时间步 2 次前向传播（与 Chen 的 √t 算法匹配）时，所提方法在长序列中使用的内存少于 Chen 方法所需内存的 50%。
• 由于能够针对任意内存预算进行优化，该方法在递减收益区域（diminishing returns）中表现出更优的内存节省效果，优于 Chen 的 √t 算法。
• 在给定假设下，动态规划公式保证了最优性，使其至少与任何启发式策略（包括 Chen 的方法）一样优秀。
• 该算法在广泛的序列长度和内存约束范围内均表现有效，即使在极低内存预算下，计算成本也仅适度增加。