Research map / Architectures and efficiency / Attention and transformers
Various Lengths, Constant Speed: Efficient Language Modeling with Lightning Attention
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Lightning Attention, an exact IO-aware linear attention implementation that splits causal attention into intra-block (left product) and inter-block (right product) computations to avoid cumsum, plus the TransNormerLLM (TNL) architecture, and reports competitive LLM benchmarks.
Problem
Linear attention has low theoretical advantage in causal language modeling because the cumsum operation makes training slow, and linear attention models often underperform softmax transformers.
Approach
Divide the sequence into blocks; compute intra-block attention with conventional masked attention and inter-block contributions with the linear kernel trick via a running KV state, eliminating cumsum and enabling constant training speed across sequence lengths with tiled forward/backward kernels in Triton.
Key results
- Lightning Attention shows linear runtime growth in sequence length while Vanilla and FlashAttention-2 grow quadratically, with lower memory than Flash2.
- TNL reaches 23.46 validation PPL on Wikitext-103 versus 24.40 for Transformer, the lowest among compared 40M models.
- Inference throughput of 7B TNL is reported as up to 11x higher than transformer models.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use TNL instead of TNN | 44M language model · WikiText-103 test | PPL: 24.67 -> 24.03 | Helped |
| Use LRPE-d instead of Mix positional encoding (ablation) | TNL · 385M, 100K updates | Loss: 2.248 -> 2.236 | Helped |
| Add decay temperature (ablation) | TNL · 385M, 100K updates | PPL: 4.804 -> 4.770 | Helped |
| Add gating mechanism (ablation) | TNL · 385M, 100K updates | Loss: 2.263 -> 2.248 | Helped |
| Replace SRMSNorm with RMSNorm (ablation) | TNL · 385M, 100K updates | Loss: 2.248 -> 2.247 | No clear effect |
| Replace SRMSNorm with LayerNorm (ablation) | TNL · 385M, 100K updates | Loss: 2.248 -> 2.247 | No clear effect |
| Replace Swish with 1+elu in GLA (ablation) | TNL · 385M, 100K updates | Loss: 2.248 -> 2.252 | No clear effect |
| Remove the GLA activation (ablation) | TNL · 385M, 100K updates | Loss: 2.248 -> 2.283 | Hurt |
| Remove the GLU activation (ablation) | TNL · 385M, 100K updates | Loss: 2.254 -> 2.248 | No clear effect |
Surprising findings
- Removing the activation from the GLU (SGLU) and using SRMSNorm caused no performance loss in ablations.
- 1+elu activation produced NaN issues in the 7B model, leading the authors to choose Swish.
Limitations
- Training corpus is self-collected and not released, limiting reproducibility of LLM benchmark comparisons.
- Baselines like FlashAttention-2 and LLaMA are reproduced by the authors under differing setups; official results mixed with reproduced ones in tables.
Experimental setup
Datasets: Wikitext-103 self-collected corpus (~2T tokens)
Benchmarks: Wikitext-103 BoolQ PIQA HellaSwag WinoGrande ARC-e ARC-c OpenBookQA MMLU C-Eval SCROLLS
Metrics: perplexity accuracy acc_norm training loss TGS (tokens per GPU per second) throughput runtime ms memory GB ROUGE F1
Baselines: LLaMA with FlashAttention-2 HGRN TNN FlashAttention-2 RWKV OPT Pythia Falcon Baichuan2 LLaMA2
Builds on: TransNormer FlashAttention-2 HGRN TNN LRPE
Closest papers by meaning
- Tiled Flash Linear Attention: More Efficient Linear RNN and xLSTM Kernels (NeurIPS 2025 · similarity 0.89)
- A Training-Free Sub-quadratic Cost Transformer Model Serving Framework with Hierarchically Pruned Attention (ICLR 2025 · similarity 0.88)
- ZeCO: Zero-Communication Overhead Sequence Parallelism for Linear Attention (NeurIPS 2025 · similarity 0.88)
- MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention (NeurIPS 2024 · similarity 0.88)
- Short-Long Convolutions Help Hardware-Efficient Linear Attention to Focus on Long Sequences (ICML 2024 · similarity 0.88)
- Beyond Linear Approximations: A Novel Pruning Approach for Attention Matrix (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.