Source-linked AI summary
Memory-Efficient Backpropagation Through Time
Audrūnas Gruslys, Remi Munos, Ivo Danihelka, Marc Lanctot, Alex Graves
TL;DR
BPTT training of RNNs is constrained by high memory consumption, especially on limited-memory devices. The paper uses dynamic programming to balance caching and recomputation under a fixed memory budget, finding an optimal execution policy. For sequences of length 1000, it saves 95% of memory while using 33% more time per iteration than standard BPTT.
Problem
BPTT requires large memory during RNN training, which is especially problematic on GPUs with limited memory.
Method
Dynamic programming balances caching intermediate results and recomputation to find an optimal policy for each fixed memory budget.
Results
95% memory savings for sequences of length 1000 required only 33% more time per training iteration than standard BPTT.
Takeaways & Limitations
The approach can tightly fit within almost any user-specified memory constraint while maximizing computational performance within that budget.
Abstract
from arXiv · showhide
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.
1 Introduction
RNNs support sequence tasks by propagating hidden-state information through time, but training them with BPTT can require substantial memory. The proposed dynamic-programming approach optimizes caching and recomputation for a fixed memory budget and works across most RNN architectures.
- RNNs propagate hidden-state information from earlier sequence positions to later ones, supporting sequence mapping problems.
- BPTT’s large memory consumption is especially troublesome on GPUs with limited memory.
- The algorithm uses dynamic programming to find an optimal memory-use policy that minimizes computational cost for any fixed memory budget.
- The approach is largely architecture agnostic and works with most recurrent neural networks.
2 Background and related work
BPTT unfolds an RNN into a time-ordered feed-forward computation whose gradients are evaluated in reverse. Memory-saving methods trade stored states for recomputation, with recursive strategies offering asymptotic savings but limited minimum memory control.
- An RNN core is a feed-forward network cloned repeatedly across time, with each clone representing a recurrence time point.
- For sequence length t, unfolding produces t RNN cores, while hidden states pass between consecutive cores.
- An internal RNN-core state stores the information needed to backpropagate through one time step given inputs and output gradients.
- A memory slot stores one hidden state or one internal state, depending on the strategy.
- BPTT treats the time-unfolded recurrent network as a feed-forward network with tied weights, then applies forward propagation and reverse backpropagation.
- Memory-saving BPTT recomputes intermediate results on demand instead of storing every hidden network state.
- Chen’s recursive algorithm uses O(k log_k+1(t)) memory and O(t log_k(t)) computational complexity.
3 Memory-efficient backpropagation through time
The paper develops dynamic-programming strategies that trade memory storage against recomputation for BPTT, including hidden-state, internal-state, and mixed-state memorization. These strategies fit constrained memory budgets while reducing computational cost, with especially strong savings for long sequences.
- BPTT-HSM: Constant-memory recomputation restores each hidden state by replaying the sequence from its beginning, substantially increasing forward computation.The minimal-memory approach repeatedly re-evaluates prior forward steps because no intermediate states can be stored.
- BPTT-HSM: Storing hidden states saves memory relative to internal states but requires an additional forward operation during backpropagation.Hidden states can be stored at time points and used to re-evaluate internal core states when gradients are propagated.
- BPTT-HSM: Dynamic programming chooses where to cache hidden states and reuses memory slots through divide-and-conquer recursion to minimize forward-operation cost.For a split at position y, the recurrence combines the cost of processing the right subsequence with m−1 slots and the left subsequence with m slots.
- BPTT-ISM: Saving internal states avoids one forward operation per divide-and-conquer step, but each stored state consumes more memory than a hidden state.The internal-state strategy modifies the boundary conditions and uses a corresponding divide-and-conquer recurrence.
- Results: 95% memory savings for sequences of length 1000 required only 33% more time per training iteration than standard BPTT.This measurement excludes input vectors and assumes each backward step costs twice as much as a forward step.
- BPTT-MSM: The mixed-state strategy optimizes over hidden- and internal-state storage and outperforms both single-state strategies for the compared settings.The general model measures internal-state storage as α times the memory of hidden-state storage, with α ≥ 2.
4 Discussion
The experiments compare memory-saving BPTT strategies under fixed computational cost or memory budgets. BPTT-MSM can fit almost arbitrary constant memory constraints, while retaining near-optimal computational cost in relevant regimes.
- 4.1 Optimality: BPTT-MSM is guaranteed to perform at least as well as Chen et al.’s strategies under any memory budget and sequence length.The guarantee follows because those strategies can be represented as potentially suboptimal policies under the same dynamic-programming assumptions.
- 4.1 Optimality: BPTT-MSM can fit within almost arbitrary constant memory constraints, unlike recursive approaches requiring at least log2 t hidden states.This provides finer control over memory usage rather than merely reducing it by a fixed amount.
- 4.2 Numerical comparison: Values below 1 in Figure 7 indicate memory savings at fixed execution speed of 2, with separate curves for β = 2, 5, 10.The left panel plots memory divided by t(1 + β) against sequence length.
- 4.2 Numerical comparison: At fixed computational cost, the approach saves a significant amount of memory as sequence length increases.The comparison fixes execution speed at 2 and evaluates memory usage across sequence lengths and β values.
- 4.2 Numerical comparison: Chen’s strategy is near optimal for one particular memory budget, so BPTT-MSM is not significantly faster there.Further memory reductions remain possible at nearly the same computational cost because that budget is already in a diminishing-returns regime.
5 Conclusion
The conclusion presents BPTT-HSM as an optimal-policy approach for training recurrent networks under fixed user-defined memory constraints. Its main advantage is tightly matching varied memory limits while maximizing computational performance.
- 5 Conclusion: The proposed approach finds optimal backpropagation strategies for recurrent neural networks under a fixed user-defined memory budget.The conclusion identifies this constrained optimization setting as the paper’s central contribution.
- 5 Conclusion: The most general algorithm performs at least as well as many commonly used heuristics.The conclusion frames this as a demonstrated comparison rather than a claim about every possible strategy.
- 5 Conclusion: The approach tightly fits almost any user-specified memory constraint while gaining maximal computational performance.This flexibility is stated as the main practical advantage of the method.
A Pseudocode for BPTT-HSM
The pseudocode first evaluates an optimal policy for a maximum sequence length and memory capacity, then executes that precomputed policy during recurrent-network backpropagation.
- A Pseudocode for BPTT-HSM: Policy evaluation uses dynamic-programming arrays sized by maximum sequence length and memory capacity.The evaluation algorithm has O(t^2 · m) complexity, reducible to O(t · m) by exploiting convexity, and is computed once before training.
- A Pseudocode for BPTT-HSM: Policy execution receives a precomputed policy, mutable RNN core, hidden-state stack, gradient, memory capacity, subsequence length, and start index.These inputs define the local subsequence computation and available stack memory.
- A Pseudocode for BPTT-HSM: Backward operations return the preceding hidden-state gradient and store the input gradient for the corresponding time position.The pseudocode assigns gradient outputs after retrieving the relevant input and hidden state.
- A Pseudocode for BPTT-HSM: During execution, forward operations generate hidden states that are pushed onto the stack for later backward computation.The pseudocode explicitly forwards through subsequence inputs and stores each resulting hidden state.
- A Pseudocode for BPTT-HSM: The execution recursively processes right and left subsequences using updated memory capacity and subsequence indices.The right call reduces available memory by one, while the left call reuses the original capacity.
B.1 General upper bound
The appendix derives a general upper bound for the dynamic program’s computational cost using nested induction and convexity. The resulting bound scales as m t^(1+1/m), with a separate linear form when memory is at least sequence length.
- B.1 General upper bound: The recurrence splits a sequence at position y and combines the costs of its two subsequences with the split’s forward-operation cost.The appendix introduces the dynamic program before optimizing the split position.
- B.1 General upper bound: The analysis defines C(t, m) as the dynamic program’s computational cost and uses induction over memory and sequence length to prove its bound.The proof establishes A(m) by induction on m and A(t, m + 1) by induction on t.
- B.1 General upper bound: Convexity of g(y) enables analysis of the optimal split and bounds the recurrence used in the induction.The function combines three convex terms, and its second derivative is bounded by 4.
- B.1 General upper bound: C(t, m + 1) is bounded by an expression involving the analytically selected split ỹ and the costs of both subproblems.The appendix derives this intermediate inequality before simplifying it.
- B.1 General upper bound: For t ≥ 4, the remaining inequality follows once ỹ ≥ 2; the cases t < 4 are verified numerically.The text identifies two small cases requiring numerical verification.
B.2 Upper bound for short sequences
The dynamic program yields optimal memory policies and supports analytical computational-cost bounds under fixed memory budgets. A recursive strategy can extend these bounds to longer sequences, although expressing the resulting computation explicitly in sequence length and memory is non-trivial.
- Dynamic programming finds an optimal policy by evaluating computational cost under a fixed memory budget.
- When sequence length equals memory, the computational cost is C(t, m) = 2t−1 < 2t.
- A recursively composed strategy achieves C(t, m) ≤ (a + 1)t for sequences bounded by the derived T(a, m) construction.
- The resulting bound is better than recursively subdividing intervals into equal sub-intervals at the same computational cost.
- Inverting the bound to express computation explicitly as a function of t and m is non-trivial.
B.3 Analytical upper bounds for BPTT-ISM and BPTT-MSM
The analytical bounds extend across the hidden-state and internal-state strategies, with internal-state memorization removing one forward operation per time-step in its bound. Because BPTT-MSM generalizes BPTT-HSM, HSM bounds also apply to the broader method.
- BPTT-ISM has computational cost C(t, m) ≤ at for sequences satisfying the stated length condition.
- Internal-state memorization removes one forward operation per time-step from the corresponding conservative upper-bound calculation.
- For internal-state strategies, the same memory-slot count can represent greater actual memory usage because internal states include hidden states.
- BPTT-HSM upper bounds also hold for BPTT-MSM because BPTT-MSM generalizes BPTT-HSM and uses an optimal policy.
C Generalizing to deep feed-forward networks
The BPTT-MSM framework can be generalized from recurrent networks to deep feed-forward networks with linear-chain computational graphs. The generalized dynamic program accounts for heterogeneous layer costs, memory requirements, and feasibility constraints, but its evaluation can become impractical for long sequences.
- Generalization: BPTT-MSM generalizes to deep architectures when their computational graph is a linear chain.
- Generalization: Different layer costs and memory requirements are handled by modifying the dynamic-program formulation.
- Dynamic program: The feasibility term assigns infinite cost when available memory cannot support the required backpropagation commitment.
- Dynamic program: The generalized cost recurrence combines forward propagation, memorization feasibility, and recursive costs for the two resulting sequence parts.
- Limitation: The generalized dynamic program requires a 3D rather than 2D array, making evaluation potentially impractical for sequences longer than a few hundred.