SUS backprop: linear backpropagation algorithm for long inputs in transformers

2025-05-21

AI Summary

To address the O(n²) computational bottleneck of attention backpropagation in Transformer training for long sequences, this work proposes a controllable sparse backpropagation mechanism: during backward pass, gradients are retained for at most c attention connections per token per head, achieved via unbiased probabilistic pruning. We theoretically prove that the induced gradient variance increase decays with sequence length. Experiments show that, at n ≈ 2000, pruning rates reach 99% (c = 20–30), increasing gradient variance by only ~1% while reducing backward computation by ~100×—making backward latency comparable to forward latency. This is the first method to reduce attention backpropagation complexity to O(nc), significantly improving training efficiency for long sequences without modifying the forward pass or model architecture.

Technology Category

• Natural Language Processing: Learning & Optimization for NLP
• Computer Vision: Large Vision Models
• Machine Learning: Mixture of Experts (MoE)

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Security and Privacy: Data transparency and provenance
• Web Mining and Content Analysis: Large pretrained models with web data

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 back-propagation 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 backpropagation 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 $99%$ of the attention gradient flow (i.e. choosing $c sim 20-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.

Problem

• Reduces quadratic backpropagation complexity in transformers to linear.
• Cuts insignificant attention weights to save computation efficiently.
• Minimizes gradient variance increase while significantly reducing compute cost.

Innovation

• Stochastic gradient cutting for backpropagation efficiency
• Linear complexity via probabilistic attention weight pruning
• Sparse matrix implementation for negligible backward pass cost