PaperlayerView exact revision on arXiv

Source-linked AI summary

SeerAttention: Learning Intrinsic Sparse Attention in Your LLMs

Yizhao Gao, Zhichen Zeng, Dayou Du, Shijie Cao, Peiyuan Zhou, Jiaxing Qi, Junjie Lai, Hayden Kwok-Hay So, Ting Cao, Fan Yang, Mao Yang

arXiv:2410.13276v4cs.CL

TL;DR

Full attention’s quadratic complexity constrains long-context efficiency, while existing sparse-attention patterns and heuristics have limited adaptability to varying contexts. SeerAttention learns block-level sparsity from the LLM itself, and the authors report better long-context accuracy and prefilling latency than prior methods, including a 7.3× speedup at 90% sparsity on 128k sequences.

  • Problem

    Full attention has quadratic complexity, while attention sparsity varies across models, contexts, and heads, limiting the adaptability of predefined patterns and heuristics.

    Source: §1 · p. 1

  • Method

    SeerAttention learns block-level attention sparsity from the LLM using a learnable gate that pools and transforms Q and K to score attention blocks.

    Source: §1 · p. 1, §3 · p. 3, §3.1 · p. 3

  • Results

    SeerAttention surpasses prior sparse-attention methods in long-context accuracy and prefilling latency, with a 7.3× speedup over dense configurations at 90% sparsity on 128k sequences.

    Source: §1 · p. 2

  • Takeaways & Limitations

    SeerAttention adapts sparsity to different heads and contexts, unlike prior methods that require careful calibration of sparse configurations.

    Source: §1 · p. 2

Abstract

from arXiv · show

Attention is the cornerstone of modern Large Language Models (LLMs). Yet its quadratic complexity hinders efficiency and scalability, especially for long-context processing. A promising approach is to leverage sparsity in attention. However, existing sparsity-based solutions predominantly rely on predefined patterns or heuristics at the attention head level, struggling to adapt dynamically to different contexts efficiently. We propose SeerAttention, a simple yet effective attention mechanism that directly learns the block-level attention sparsity from the LLM itself. Inspired by the gating mechanism in Mixture of Experts (MoE), SeerAttention augments the conventional attention with a learnable gate that selectively activates important blocks within the attention map. Specifically, the gate first pools the query (Q) and key (K) tensors along the sequence dimension and processes them through learnable linear layers. The resulting matrices are then multiplied together to produce the gating scores, which are used to predict block-level attention sparsity. Combined with our block-sparse FlashAttention kernel, SeerAttention can achieve significant speedup on GPUs. When applied to pre-trained LLMs, SeerAttention only requires training the gate parameters in a lightweight self-distillation manner, allowing rapid convergence. Our evaluation results demonstrate that SeerAttention achieves better model accuracy and lower latency for long-context pre-filling compared to prior methods. Code is available at: https://github.com/microsoft/SeerAttention

1. Introduction

Full attention’s quadratic cost limits long-context efficiency, while attention sparsity varies across models, contexts, and heads. SeerAttention learns block-level sparsity from the LLM and reports better long-context accuracy and prefilling latency than prior sparse methods.

1. Introduction

  • Attention sparsity varies across models, contexts, and heads, limiting predefined patterns and heuristics despite sparsity reaching 95% or 99% in some heads.The paper notes that sparsity becomes more prominent in longer contexts.Source: §1 · p. 1
  • SeerAttention adds a learnable gate that pools Q and K, transforms them with linear layers, and computes scores to select important attention blocks.Skipping unimportant blocks reduces I/O overhead and computational cost.Source: §1 · pp. 1–2
  • A lightweight self-distillation scheme trains the gate while keeping other pretrained-model parameters fixed.The pooled standard-attention map serves as teacher for block-level sparsity learning.Source: §1 · p. 2
  • 7.3× speedup over dense configurations at 90% sparsity on 128k sequences, alongside higher long-context accuracy and lower prefilling latency than prior sparse methods.The reported comparison includes Minference, MoA, and DuoAttention.Source: §1 · p. 2

2. Background and Related Works

Attention supports long-range context but has quadratic time and memory costs, motivating more efficient alternatives. The section reviews sparse attention’s dynamic behavior and other approaches to long-context optimization.

Powerful but Complex Attention in Transformer

  • Transformer attention has O(n^2) time and memory complexity; alternatives such as Reformer and linear attention reduce complexity, but can struggle to match full attention’s performance.The passage also describes recurrent frameworks including RWKV, RetNet, and Mamba.Source: §2 · p. 2

Intrinsic but Dynamic Sparsity in Attention

  • Attention sparsity becomes more pronounced in longer contexts but varies by input and head, making predefined patterns and handcrafted heuristics less general.The passage connects this dynamic behavior to the motivation for learning sparsity from data.Source: §2 · p. 2

Long-Context LLM Optimizations

  • Long-context optimization includes prefill approaches such as prompt compression and sparse attention, plus decoding optimizations and KV-cache sharing, eviction, and quantization.These strategies target different parts of long-context processing.Source: §2 · pp. 2–3

3. SeerAttention

SeerAttention learns block-level attention sparsity with a gate that pools Q and K, then uses predicted sparsity with block-sparse FlashAttention for inference. The design combines self-distillation, pooling choices, and block-level positional encoding to support efficient attention across context lengths.

3.1. Attention Gate Design

  • AttnGate pools Q and K along the sequence dimension, applies linear layers, and multiplies the outputs to form one score per attention block.With a typical block size of 64, the gate output is 1/4096 the size of the original attention map.Source: §3.1 · p. 3
  • On PG19, average pooling on Q and combined max, min, and average pooling on K achieve the best test perplexity across sparsity ratios.When multiple pooling methods are used for a tensor, their outputs are concatenated in the hidden dimension before the linear layer.Source: §3.1 · p. 3
  • Block-level RoPE lets AttnGate extrapolate across context lengths: without it, training at 8k performs inadequately beyond 8k, and training at 64k does not perform well at 128k.The gate receives Q and K before RoPE and applies an additional block-level RoPE, with θ′ = θ/B.Source: §3.1 · p. 3

3.2. AttnGate Distillation in SeerAttention

  • AttnGate is distilled against the 2D-max-pooled full-attention map, whose block values are small only when all attention scores in that block are small.The gate uses KL-divergence to mimic the attention distribution; the authors modify FlashAttention to produce the max-pooled training target efficiently.Source: §3.2 · pp. 3–4

3.3. Inference with SeerAttention

  • At inference, TopK or thresholding converts gate scores to a binary block mask, and block-sparse FlashAttention skips I/O and computation for unactivated blocks.Users can adjust the TopK ratio or threshold at test time; the gate block size aligns with FlashAttention tiling, typically 64 or 128.Source: §3.3 · p. 5

4. Experiments

SeerAttention is evaluated for accuracy and efficiency on long- and short-context tasks, with learned sparsity tested against prior methods and at multiple sparsity levels. The experiments also examine kernel overhead, end-to-end prefill speed, and the attention patterns learned by its gate.

4. Experiments

  • The evaluation fixes the block size at 64 and applies AttnGate only during prefill; the main model and baseline comparisons use Llama-3.1-8B-Instruct.Distillation uses RedPajama chunks of 64k tokens and updates only AttnGate parameters for 500 steps.Source: §4 · p. 5

4.1. Accuracy of Evaluation

  • On PG19, a single trained AttnGate supports test-time TopK or threshold changes, generally giving lower perplexity than MoA and MInference at similar sparsity ratios.MoA's 128k result is missing because of an out-of-memory issue on a single A100 GPU.Source: §4.1 · p. 6, Figure 5 · p. 6
  • Across LongBench test lengths, SeerAttention consistently outperforms other methods and achieves the highest average score and average sparsity.Its score also exceeds the dense baseline on the 0–4k and 4–8k tests.Source: §4.1 · p. 6, Table 1 · p. 6
  • On RULER, SeerAttention has the highest average accuracy among compared models, 0.41% below the dense baseline, and the highest average prefill speedup at 1.41×.At 128k, it maintains over 80% sparsity while DuoAttention uses less than 50%, accompanying slightly lower performance.Source: §4.1 · p. 6, Table 2 · p. 7
  • At average sequence length 872 on GSM-8K, SeerAttention has 0.1% accuracy degradation at 52% average sparsity; sparse attention offers little latency benefit on short inputs.Attention contributes a smaller share of total runtime for short-context inputs.Source: §4.1 · p. 6, Table 3 · p. 7

4.2. Efficiency Evaluation

  • At 32K context and 0.5 sparsity, AttnGate contributes 1% of attention-layer latency, and its relative overhead almost disappears at 128K.Efficiency tests use Llama-3.1-8B-Instruct on a single A100 and compare against FlashAttention-2, MoA, MInference, and DuoAttention.Source: §4.2.1 · p. 7, §4.2 · p. 7

4.2.1. KERNEL EVALUATION

  • At 128k sequence length and 90% sparsity, SeerAttention's block-sparse kernel is 7.3× faster than full-attention FlashAttention-2 on one A100 GPU.The kernel outperforms MoA and MInference in most tested sparsity and sequence-length cases; end-to-end prefill speedup reaches 2.43× at 128k.Source: §4.2.1 · pp. 7–8

4.3. Visualization of Learned Attention Maps.

  • AttnGate learns varied patterns without prior heuristics, including A-shape, vertical, slash, diagonal-block, and random patterns that extend beyond examples in prior work.The paper notes that these learned patterns encompass and extend those observed in MoA and MInference.Source: §4.3 · p. 8

5. Conclusion and Future Work

The conclusion presents SeerAttention as a way to learn and use intrinsic attention sparsity for long-context LLMs, reporting better accuracy and prefill latency than prior approaches. Future work includes improving training through long-context continued pretraining and applying the method during decoding, especially for long-CoT.

5. Conclusion and Future Work

  • The paper concludes that SeerAttention learns intrinsic attention sparsity and outperforms prior approaches in long-context model accuracy and prefill latency.It identifies long-context continued pretraining with more training tokens as a direction for seeking higher sparsity without sacrificing accuracy.Source: §5 · pp. 8–9
  • Future directions include improved training methods, including long-context continued pretraining, and applying SeerAttention during decoding, especially for long-CoT.The paper notes preliminary experiments on continued pretraining in Appendix A.2.Source: §5 · p. 9

A.1. Training SeerAttention with Customized GPU Kernel

The appendix describes a customized FlashAttention kernel that extracts a 2D-max-pooled attention map for SeerAttention training. It reports low memory and latency overhead compared with FlashAttention-2.

A.1. Training SeerAttention with Customized GPU Kernel

  • A Triton kernel modifies FlashAttention to extract the 2D-max-pooled full-attention map used as the gate's distillation target.Applying a column max to a_ij produces the 2D pooling; the authors say this adds only minor overhead while improving ground-truth extraction efficiency.Source: §A.1 · p. 12, Figure 10 · p. 12
  • The kernel stores intermediate r_ij values in HBM and rescales them with the final global maximum m_i and exponential sum l_i before pooling.FlashAttention updates the global maximum iteratively from local row maxima; the resulting a_ij represents the correct row maximum of the original attention block.Source: §A.1 · p. 12
  • During AttnGate training, the customized kernel adds minimal memory and latency overhead versus FlashAttention-2, while PyTorch manual attention runs out of memory above 4k tokens.The customized kernel's peak memory is similar to FlashAttention-2; PyTorch also encounters OOM above 8k in the latency comparison.Source: §A.1 · pp. 12–13, Figure 11 · p. 13

A.2. Preliminary Experiments of Fine-tuning with SeerAttention

The preliminary experiment integrates SeerAttention into YaRN context-extension fine-tuning for Llama-3-8B. At 50% sparsity it reports near-baseline perplexity, with slightly higher loss at 90% sparsity.

A.2. Preliminary Experiments of Fine-tuning with SeerAttention

  • Fine-tuning loss at 50% sparsity nearly overlaps the dense YaRN baseline, while 90% sparsity has slightly higher loss; the 50%-to-90% curves are described as nearly identical.The experiment compares the dense baseline with SeerAttention at 50% and 90% sparsity.Source: Figure 12 · p. 13, §A.2 · p. 13
  • At 50% sparsity, YaRN fine-tuning with SeerAttention yields PG19 perplexity 8.81 and ProofPile perplexity 2.47, versus dense YaRN's 8.79 and 2.46.These test perplexities are evaluated at 32k context length.Source: Table 4 · p. 13, §A.2 · p. 13
  • The experiment extends Llama-3-8B from 8k to 32k context length and combines cross-entropy loss with AttnGate loss during YaRN fine-tuning.AttnGates are first initialized using post-training self-distillation before the entire model is fine-tuned.Source: §A.2 · p. 13
Loading 2410.13276v4…