# SUS backprop: linear backpropagation algorithm for long inputs in transformers

arXiv: https://arxiv.org/abs/2505.15080 · HTML: https://mlsanity.com/paper/2505.15080/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 29 of its 33 references that are in the map.
- First posted on arXiv 21 May 2025.

## Builds on (most cited first, 29 of 29)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Robust Speech Recognition via Large-Scale Weak Supervision](https://mlsanity.com/paper/2212.04356/) (2022)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Informer: Beyond Efficient Transformer for Long Sequence Time-Series Forecasting](https://mlsanity.com/paper/2012.07436/) (2020)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model](https://mlsanity.com/paper/2405.04434/) (2024)
- [What Does BERT Look At? An Analysis of BERT's Attention](https://mlsanity.com/paper/1906.04341/) (2019)
- [Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned](https://mlsanity.com/paper/1905.09418/) (2019)
- [Fast Transformer Decoding: One Write-Head is All You Need](https://mlsanity.com/paper/1911.02150/) (2019)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention](https://mlsanity.com/paper/2502.11089/) (2025)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)
- [A Comprehensive Survey on Applications of Transformers for Deep Learning Tasks](https://mlsanity.com/paper/2306.07303/) (2023)
- [Unbiased Gradient Estimation in Unrolled Computation Graphs with Persistent Evolution Strategies](https://mlsanity.com/paper/2112.13835/) (2021)
- [Stochastic Backpropagation: A Memory Efficient Strategy for Training Video Models](https://mlsanity.com/paper/2203.16755/) (2022)
- [An In-depth Study of Stochastic Backpropagation](https://mlsanity.com/paper/2210.00129/) (2022)

## Abstract (arXiv)

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(n2) 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 ~ 25-30) results in relative gradient variance increase of only about 1% for n ~ 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.