SUS backprop: linear backpropagation algorithm for long inputs in transformers

TL;DR

This paper introduces SUS backprop, a probabilistic method to cut most attention gradient flow in transformers, reducing backpropagation complexity from quadratic to linear for long sequences with minimal impact on gradient variance.

Contribution

It proposes a simple probabilistic rule to sparsify attention backpropagation, significantly reducing computational cost for long sequences in transformers.

Findings

Reduces attention backpropagation complexity from O(n^2) to O(nc).

Maintains minimal gradient variance increase (~1%) with high sparsity.

Effective for long sequence training with sparse matrix implementation.

Abstract

It is straightforward to design an unbiased gradient estimator that stochastically cuts the backpropagation flow through any part of a computational graph. By cutting the parts that have little effect on the computation, one can potentially save a significant amount of backpropagation computation in exchange for a minimal increase in the stochastic gradient variance, in some situations. Such a situation occurs in the attention mechanism of the transformer architecture. For long sequences, attention becomes the limiting factor, as its compute requirements increase quadratically with sequence length $n$. At the same time, most attention weights become very small, as most attention heads tend to connect a given token with only a small fraction of other tokens in the sequence. These weights become promising targets for cutting backpropagation. We propose a simple probabilistic rule controlled by a single parameter $c$ that cuts back-propagation through most attention weights, leaving at most $c$ interactions per token per attention head. This brings a factor of $c/n$ reduction in the compute required for the attention backpropagation, turning it from quadratic $O(n^2)$ to linear complexity $O(nc)$. We have empirically verified that, for a typical transformer model, cutting about $99\%$ of the attention gradient flow (i.e. choosing $c \sim 25-30$) results in relative gradient variance increase of only about $1\%$ for $n \sim 2000$, and it decreases with $n$. This approach is amenable to efficient sparse matrix implementation, thus being promising for making the cost of a backward pass negligible relative to the cost of a forward pass when training a transformer model on long sequences.