Memory-Efficient Backpropagation Through Time

TL;DR

This paper introduces a memory-efficient algorithm for backpropagation through time in RNNs that balances caching and recomputation, significantly reducing memory use while maintaining computational efficiency, especially for long sequences.

Contribution

A novel dynamic programming-based method that optimally manages memory and computation trade-offs in BPTT, adaptable to various memory constraints.

Findings

Reduces memory usage by 95% for sequences of length 1000

Increases computational time by only one third compared to standard BPTT

Effective for training long sequence RNNs

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.