Backpropagation Through Time › Computational Cost of BPTT

Understanding the computational cost of BPTT matters because it directly constrains what models you can train on what hardware. The cost includes more than FLOPs; memory bandwidth, cache utilization, and the need to store intermediate activations all play important roles.

Let's be concrete. Suppose you have a sequence of length $T$, hidden state dimension $d$, input dimension $k$, and batch size $B$.

Forward pass cost (full sequence):

The dominant operation at each step is the matrix-vector product $\mathbf{W}_{hh}\mathbf{h}_{t-1}$, which costs $O(d^2)$ operations (since $\mathbf{W}_{hh}$ is $d \times d$). Across all $T$ steps and $B$ batches, the forward pass costs $O(T \cdot d^2 \cdot B)$.

Backward pass cost (full BPTT):

The backward pass has roughly the same cost as the forward pass, $O(T \cdot d^2 \cdot B)$. Computing gradient contributions to $\mathbf{W}_{hh}$ involves outer products at each step costing $O(d^2)$, and contributions to $\mathbf{W}_{xh}$ cost $O(d \cdot k)$ per step.

Total cost (full BPTT): $O(T \cdot d^2 \cdot B)$ for the complete forward and backward passes.

Total cost (truncated BPTT with chunk size $k_{\text{chunk}}$): The same asymptotic cost, but practical benefits arise from fitting the computation within GPU cache, allowing larger batches, and parallelism across chunks. Each of the $T / k_{\text{chunk}}$ chunks runs a self-contained forward/backward pass of depth $k_{\text{chunk}}$.

The hidden dimension $d$ dominates the cost in practice. Doubling the hidden size quadruples the computation for the recurrent matrix products. This is why RNNs with large hidden states are expensive to train on long sequences, and why the community moved toward transformer architectures that compute attention in $O(T^2 \cdot d)$ rather than $O(T \cdot d^2)$, a trade-off that favors larger $d$ at the cost of longer sequences.

The sequential nature of the RNN forward pass is a critical bottleneck that the computational cost analysis understates. To compute $\mathbf{h}_t$, you must first have computed $\mathbf{h}_{t-1}$. This strict temporal dependency means that no matter how many parallel processors you have, the forward pass through a sequence of length $T$ requires $T$ sequential steps. Modern GPUs are optimized for massively parallel operations, and sequential dependencies waste this parallelism. A transformer processes all tokens in parallel (though attention still has quadratic cost in sequence length), which is one key reason transformers train dramatically faster than RNNs on modern hardware despite having similar theoretical FLOPs.

For applications that require full BPTT on long sequences, gradient checkpointing offers a practical solution. Instead of storing all intermediate activations during the forward pass, you store only a subset of checkpoint activations and recompute the missing ones during the backward pass when they are needed.

The trade-off is straightforward: checkpointing reduces memory by a factor of roughly $\sqrt{T}$ (with the right checkpointing strategy) at the cost of one additional forward pass through the checkpointed segments. If the sequence is very long and memory is the bottleneck, this trade is usually worth it. PyTorch provides checkpointing support through torch.utils.checkpoint.checkpoint_sequential, which can be applied to RNN segments without modifying the gradient computation.

Gradient checkpointing is especially useful during fine-tuning of pretrained language models on long documents, where full sequences may contain thousands of tokens that would exhaust GPU memory under standard full-BPTT training. The additional compute overhead of recomputation is typically a 30-50% increase in training time, which is often acceptable when the alternative is running out of memory or reducing batch size to a single example.