# Multi-Layer Transformers Gradient Can be Approximated in Almost Linear Time

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

## Key facts

- Cited by 22 papers in the map (data through 14 Oct 2026).
- Builds on 109 of its 144 references that are in the map.
- First posted on arXiv 23 Aug 2024.

## Builds on (most cited first, 30 of 109)

- [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)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks](https://mlsanity.com/paper/1803.03635/) (2018)
- [Retrieval-Augmented Generation for Large Language Models: A Survey](https://mlsanity.com/paper/2312.10997/) (2023)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Fast Inference from Transformers via Speculative Decoding](https://mlsanity.com/paper/2211.17192/) (2022)
- [The Rise and Potential of Large Language Model Based Agents: A Survey](https://mlsanity.com/paper/2309.07864/) (2023)
- [SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models](https://mlsanity.com/paper/2211.10438/) (2022)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shot](https://mlsanity.com/paper/2301.00774/) (2023)
- [Effective Approaches to Attention-based Neural Machine Translation](https://mlsanity.com/paper/1508.04025/) (2015)
- [Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads](https://mlsanity.com/paper/2401.10774/) (2024)
- [Evaluating Very Long-Term Conversational Memory of LLM Agents](https://mlsanity.com/paper/2402.17753/) (2024)
- [RWKV: Reinventing RNNs for the Transformer Era](https://mlsanity.com/paper/2305.13048/) (2023)
- [SNIP: Single-shot Network Pruning based on Connection Sensitivity](https://mlsanity.com/paper/1810.02340/) (2018)
- [A Simple and Effective Pruning Approach for Large Language Models](https://mlsanity.com/paper/2306.11695/) (2023)
- [Extending Context Window of Large Language Models via Positional Interpolation](https://mlsanity.com/paper/2306.15595/) (2023)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [What is the State of Neural Network Pruning?](https://mlsanity.com/paper/2003.03033/) (2020)

## Built on by (most cited first, 22 of 22)

- [LazyDiT: Lazy Learning for the Acceleration of Diffusion Transformers](https://mlsanity.com/paper/2412.12444/) (2024)
- [Circuit Complexity Bounds for RoPE-based Transformer Architecture](https://mlsanity.com/paper/2411.07602/) (2024)
- [Bypassing the Exponential Dependency: Looped Transformers Efficiently Learn In-context by Multi-step Gradient Descent](https://mlsanity.com/paper/2410.11268/) (2024)
- [The Computational Limits of State-Space Models and Mamba via the Lens of Circuit Complexity](https://mlsanity.com/paper/2412.06148/) (2024)
- [Looped ReLU MLPs May Be All You Need as Practical Programmable Computers](https://mlsanity.com/paper/2410.09375/) (2024)
- [HSR-Enhanced Sparse Attention Acceleration](https://mlsanity.com/paper/2410.10165/) (2024)
- [High-Order Matching for One-Step Shortcut Diffusion Models](https://mlsanity.com/paper/2502.00688/) (2025)
- [RoPE Attention Can Be Trained in Almost Linear Time](https://mlsanity.com/paper/2412.17316/) (2024)
- [Theoretical Constraints on the Expressive Power of mathsfRoPE-based Tensor Attention Transformers](https://mlsanity.com/paper/2412.18040/) (2024)
- [Video Latent Flow Matching: Optimal Polynomial Projections for Video Interpolation and Extrapolation](https://mlsanity.com/paper/2502.00500/) (2025)
- [On Computational Limits of FlowAR Models: Expressivity and Efficiency](https://mlsanity.com/paper/2502.16490/) (2025)
- [Modern Hopfield Networks Require Chain-of-Thought to Solve mathsfNC1-Hard Problems](https://mlsanity.com/paper/2412.05562/) (2024)
- [Fast RoPE Attention: Combining the Polynomial Method and Fast Fourier Transform](https://mlsanity.com/paper/2505.11892/) (2025)
- [Only Large Weights (And Not Skip Connections) Can Prevent the Perils of Rank Collapse](https://mlsanity.com/paper/2505.16284/) (2025)
- [Advancing the Understanding of Fixed Point Iterations in Deep Neural Networks: A Detailed Analytical Study](https://mlsanity.com/paper/2410.11279/) (2024)
- [Scaling Law Phenomena Across Regression Paradigms: Multiple and Kernel Approaches](https://mlsanity.com/paper/2503.01314/) (2025)
- [HOFAR: High-Order Augmentation of Flow Autoregressive Transformers](https://mlsanity.com/paper/2503.08032/) (2025)
- [When Can We Solve the Weighted Low Rank Approximation Problem in Truly Subquadratic Time?](https://mlsanity.com/paper/2502.16912/) (2025)
- [Time and Memory Trade-off of KV-Cache Compression in Tensor Transformer Decoding](https://mlsanity.com/paper/2503.11108/) (2025)
- [Theoretical Foundation of Flow-Based Time Series Generation: Provable Approximation, Generalization, and Efficiency](https://mlsanity.com/paper/2503.14076/) (2025)
- [GRACE: A Dynamic Coreset Selection Framework for Large Language Model Optimization](https://mlsanity.com/paper/2604.11810/) (2026)
- [Fundamental Limits of Crystalline Equivariant Graph Neural Networks: A Circuit Complexity Perspective](https://mlsanity.com/paper/2510.05494/) (2025)

## Abstract (arXiv)

The computational complexity of the self-attention mechanism in popular transformer architectures poses significant challenges for training and inference, and becomes the bottleneck for long inputs. Is it possible to significantly reduce the quadratic time complexity of computing the gradients in multi-layer transformer models? This paper proves that a novel fast approximation method can calculate the gradients in almost linear time n1+o(1) where n is the input sequence length, while it maintains a polynomially small approximation error 1 / poly(n) across the entire model. Our theory holds for general loss functions and when the multi-layer transformer model contains many practical sub-modules, such as residual connection, casual mask, and multi-head attention. By improving the efficiency of gradient computation, we hope that this work will facilitate more effective training and deployment of long-context language models based on our theoretical results.