SeerAttention: Learning Intrinsic Sparse Attention in Your LLMs

## Abstract

Abstract Attention is the cornerstone of modern Large Language Models (LLMs). Yet its quadratic complexity limits the efficiency and scalability of LLMs, especially for those with a long-context window. A promising approach addressing this limitation is to leverage the sparsity in attention. However, existing sparsity-based solutions predominantly rely on predefined patterns or heuristics to approximate sparsity. This practice falls short to fully capture the dynamic nature of attention sparsity in language-based tasks.
This paper argues that attention sparsity should be learned rather than predefined . To this end, we design SeerAttention, a new Attention mechanism that augments the conventional attention with a learnable gate that adaptively selects significant blocks in an attention map and deems the rest blocks sparse.
Such block-level sparsity effectively balances accuracy and speedup.
To enable efficient learning of the gating network, we develop a customized FlashAttention implementation that extracts the block-level ground truth of attention map with minimum overhead.
SeerAttention not only applies to post-training, but also excels in long-context fine-tuning.
Our results show that at post-training stages, SeerAttention significantly outperforms state-of-the-art static or heuristic-based sparse attention methods, while also being more versatile and flexible to adapt to varying context lengths and sparsity ratios.
When applied to long-context fine-tuning with YaRN, SeerAttention can achieve a remarkable 90% sparsity ratio at a 32k context length with minimal perplexity loss, offering a 5.67 × 5.67\times 5.67 × speedup over FlashAttention-2.
Code will be available at: https://github.com/microsoft/SeerAttention

## 1 Introduction

Attention is a fundamental mechanism in transformer-based LLMs (Vaswani, ). Despite its effectiveness, the quadratic complexity of attention requires substantial computation and memory resources, limiting the scalability and efficiency of LLMs, especially with a long-context window. This problem has become an active research topic in the community. One potential solution is to replace the quadratic attention with cheaper architectures like linear attention or recurrent networks (Katharopoulos et al., ; Peng et al., ) with subquadratic complexity. These solutions, although efficient, struggle to match the efficacy of full attention, especially when the scale is large.

A promising approach with increasing interests is to leverage sparsity in attention. Sparsity commonly exists in attention maps, and it becomes more prominent in longer contexts. In certain LLM attention heads, the sparsity ratio can reach 95% or even 99%, posing great opportunities for efficiency improvement. However, prior studies often rely on predefined sparsity patterns or heuristics to approximate the attention mechanism (Jiang et al., ; Fu et al., ; Lee et al., ; Zhu et al., ; Han et al., ).
The sparsity shown in attention maps varies significantly across different models, language inputs and attention heads, making predefined patterns or heuristics insufficient.

In this paper, we argue that attention sparsity should be learned rather than predefined. To achieve this, we introduce SeerAttention, a novel attention mechanism that enhances the standard attention with a learnable gate.
During the forward pass of SeerAttention, the Q and K inputs are pooled and processed by the learnable gate to adaptively identify the important blocks, allowing the downstream block-sparse attention kernel to effectively reduce the I/O and computation by skipping unimportant blocks.

In the training of SeerAttention, the gate learns the block-wise attention sparsity from the model itself, i.e., using the attention map generated from the standard attention.
However, FlashAttention (Dao et al., ), the state-of-the-art attention kernel used predominantly in LLMs, eliminates the explicit output of the intermediate attention maps via operation fusion to improve efficiency.
This posts great challenges in our training processes, especially in long-context scenarios as the naïve manual attention implementation is slow and memory hungry.
To address this challenge, SeerAttention customizes a FlashAttention kernel to extract our targeted block-wise attention map information without maintaining the original full size attention map. This new implementation achieves negligible overheads and significantly boosts the scalability of training process.

We evaluate SeerAttention in two settings: post-training, where only the gate parameters are learned using a small set of calibration data;
and fine-tuning, where both the gate and weights of the original model are jointly optimized during long context extension.
Our results demonstrate that SeerAttention surpasses state-of-the-art sparse attention methods like Minference (Jiang et al., ) and MoA (Fu et al., ).
Notably, in contrast to previous methods that require careful calibration of sparse configuration for different settings, SeerAttention offers strong capabilities of adaptation to arbitrary context lengths and sparsity ratios.
More importantly, the inherent learning capability of SeerAttention achieves near-lossless accuracy with 50% sparsity and minimal loss even with 90% sparsity during long-context fine-tuning (shown in Figure 1). The block-sparse kernel also demonstrates up to $5.67\times$ speedup over FlashAttention-2 dense baseline at 32k context size with 90% sparsity.
Remarkably, with block-sparse pattern, SeerAttention exhibits the ability to learn more diverse patterns, including A-shape and Vertical-Slash, further demonstrating its versatility and performance.

Our contributions can be summarized as follows:

- •
We propose SeerAttention, an innovative attention mechanism that learns and leverages the intrinsic sparsity in attention to enhance efficiency for long-context LLMs.
- •
We develop a customized FlashAttention kernel that effectively obtains block-level attention map ground truth, enabling scalable learning of sparse attention.
- •
Experiments show that SeerAttention outperforms previous approaches in post-training, offers adaptability to various context lengths and sparsity ratios, and excels in long-context fine-tuning, maintaining near-lossless accuracy even at high sparsity levels.

## 2 Background and Motivation

##### Powerful but Complex Attention in Transformer.

The advent of attention mechanisms, particularly within the Transformer architecture (Vaswani, ), marked a significant advancement in natural language processing.
Attention enables improved handling of long-range dependencies and a better understanding of context by
attending each token to every other token in the sequence, resulting in a quadratic memory and time complexity $O(n^{2})$, where $n$ is the sequence length.
This presents a significant challenge as the community moves towards LLMs that can process increasingly longer contexts.
Many studies explore alternative attention mechanisms to mitigate this complexity.
The Reformer architecture (Kitaev et al., ) reduces the complexity to $O(n\log n)$ and the linear attention mechanism (Katharopoulos et al., ) further decreases complexity to $O(n)$.
Recently, there has been a trend of revisiting recurrent neural networks, leading to the proposal of new architectural frameworks such as RWKV (Peng et al., ), RetNet (Sun et al., ), and Mamba (Gu & Dao, ).
Despite their promise of efficiency, these methods struggle to match the performance of full attention mechanisms, particularly with larger models and longer contexts.

##### Intrinsic but Dynamic Sparsity in Attention.

Attention mechanisms inherently exhibit sparsity, which arises from the attention map $\mathbf{A}$ generated by $\mathbf{Q}$ and $\mathbf{K}$: $\mathbf{A}=\mathrm{softmax}(\mathbf{QK^{T}}/\sqrt{d})$.
The softmax function often produces a multitude of negligible scores that can be treated as zeros without impacting model accuracy (Zaheer et al., ; Liu et al., ; Wang et al., ; Child et al., ).
Attention sparsity becomes more pronounced with longer contexts, presenting opportunities to optimize inference speed.
Unfortunately, this sparsity is dynamic, varying across different inputs and attention heads, each displaying distinct sparsity locations and ratios.
Prior research has attempted to approximate attention sparsity using predefined patterns and heuristics (Fu et al., ; Jiang et al., ).
Yet, these methods lack generality and often rely on handcrafted features, struggling to fully capture the sparsity behavior of attention mechanisms.
The dynamic and input-dependent nature of attention sparsity echoes the principles of Mixture of Experts (MoE) models  (Shazeer et al., ; Fedus et al., ) suggesting that sparsity should ideally be learned directly from data within the model itself. This approach would allow models to adaptively harness sparsity, improving efficiency while maintaining accuracy.

## 3 SeerAttention

SeerAttention adopts a fully learning-based approach to adaptively identify attention sparsity in LLMs and leverages the learned sparsity for efficient inference. To ensure efficiency on modern hardware like GPUs, we focus on learning block sparsity, which can seamlessly integrate with the tiling computation scheme of FlashAttention (Dao et al., ).
Figure 2 illustrates the overall model architecture of SeerAttention, which augments conventional attention with a learnable gating module, termed Attention Gate (AttnGate). This module contains learnable parameters that identify the locations of significant blocks in the attention maps. By utilizing these block indices, the subsequent attention computation can employ a block-sparse FlashAttention kernel, significantly enhancing performance by reducing I/O and computation overhead.

### 3.1 Attention Gate

The AttnGate module is designed to learn block-wise information with minimal overhead. It takes the original matrices $\mathbf{Q}$ and $\mathbf{K}$ as inputs and downsamples them using pooling along the sequence dimension. As shown in Figure 2(a), for a given attention head, the sizes of the downsampled $\mathbf{Q}$ and $\mathbf{K}$ become $[seq/B,d]$, where $B$ is the block size. The downsampled $\mathbf{Q}$ and $\mathbf{K}$ are then processed through a linear layer and multiplied together, similar to the standard attention operation. This results in a matrix of size $[seq/B,seq/B]$, where each element corresponds to one block in the full attention map. With a typical block size of 64, the output of the AttnGate module is only $\frac{1}{4096}$ the size of the original attention map. During inference, by selecting the Top-k blocks in each row, the block-sparse FlashAttention kernel can efficiently load and process only the active blocks.

##### Pooling Selection.

In SeerAttention, different pooling methods can be composed for the $\mathbf{Q}$ and $\mathbf{K}$ tensors, currently allowing for combinations of average, max, and min pooling. Multiple pooling operations can be applied to each matrix, with the resulting downsampled matrices concatenated before being fed into the linear layer. Experimental results indicate that the optimal combination is to use average pooling on $\mathbf{Q}$ and a combination of max and min pooling on $\mathbf{K}$ (details in Figure 10).

##### Additional RoPE in Attention Gate.

Modern LLMs typically employ RoPE (Su et al., ) to encode positional information.
If the AttnGate relies solely on the original RoPE in the model, i.e., feeding the AttnGate with $\mathbf{Q}$ and $\mathbf{K}$ after RoPE, the relative positional encoding properties will be lost because of the pooling operation.
This compromises the AttnGate’s ability to extrapolate to longer context lengths during training.
Specifically, if the AttnGate is trained on 8k sequences, it struggles with inputs longer than 16k. To address this issue, we introduce a separate RoPE within the AttnGate.
This RoPE can reuse the parameters from the original RoPE, but assigns position ids based on the starting positions of each block. This is equivalent to using a reduced rotational angle $\theta^{\prime}=\theta/B$, but encode the position of each block.

### 3.2 Block-Sparse FlashAttention Inference Kernel

Block sparsity is not officially supported in FlashAttention-2 (Dao, ), so we implement our own block-sparse FlashAttention kernel with Triton (Tillet et al., ) to speedup the inference of SeerAttention. It uses similar dataflow of FlashAttention-2 where $\mathbf{Q}$ is split across different warps. Each warp reads the sparse block indices generated by AttnGate and loads the corresponding $\mathbf{K}$ and $\mathbf{V}$ blocks on-chip for computation. This approach efficiently reduces both I/O and computation overhead by skipping non-activated blocks.

Figure: (a) SeerAttention at Inference Time
Refer to caption: x4.png

## 4 Training SeerAttention

While the introduced SeerAttention architecture is straightforward, training it is challenging. Jointly training the gate and attention from scratch, as in MoE, is costly and difficult.
Fortunately, unlike MoE, where gating network must learn expert selection from scratch, the AttnGate in SeerAttention has a ground truth in standard attention as guidance.

### 4.1 Training the Attention Gate

We train the AttnGate to learn block-level sparsity by using the 2D max-pooled attention map from full attention as ground truth, as illustrated in Figure 2(b).
To align distributions, the AttnGate’s output is scaled and passed through a softmax, similar to standard attention mechanisms. Additionally, the max-pooled attention map is row-normalized to sum to 1, consistent with the softmax output. Mean-Square-Error (MSE) loss is used in training.
This auto-regressive training scheme also enables flexible usage of SeerAttention, allowing users to adjust the Top-k ratio to balance accuracy and efficiency with a single model.

### 4.2 FlashAttenion with Max-Pooling: A Customized training kernel

Obtaining the max-pooled attention map for training is non-trivial especially in long-context scenarios.
Modern LLMs rely on FlashAttention, which fuses operations and doesn’t explicitly compute the attention map. The naïve mannul implementation is impractical due to quadratic memory complexity. To address this challenge, we customize an efficient kernel that directly outputs the max-pooled attention map by modifying FlashAttention but largely reuses its original computation flow. Figure 3 shows the pseudo code and diagram of this customized kernel.

Figure: Figure 3: Efficient FlashAttention kernel with pooling of attention map.
Refer to caption: x6.png

Normally, the softmax function ensures numerical stability by subtracting the maximum value before applying the exponential operation. FlashAttention computes the local row max of each block, and gradually updates the global maximum through iteration:

$$ $S_{ij}=Q_{i}K^{T}_{j};\ \,r_{ij}=\mathrm{rowmax}(S_{ij});\,\ m_{ij}=\mathrm{ max}(m_{i(j-1)},r_{ij})$ (1) $$

where $r_{ij}$ is typically treated as a temporary result. However, we store it in HBM and rescale it later with the final global max $m_{i}$ and sum of exp $l_{i}$ after the iteration:

$$ $a_{ij}=\mathrm{exp}(r_{ij}-m_{i})/l_{i}$ (2) $$

This $a_{ij}$ represents the correct row max of the original attention block. With that, 2D max-pooling is achieved by applying a column max over $a_{ij}$. This introduces only minor overhead (storing and rescaling $r_{ij}$) but significantly improves the efficiency of obtaining the ground truth. Detailed code is available in the Appendix A, and the overhead analysis is in Figure 8.

### 4.3 Apply SeerAttention in Post-training and Fine-tuning Stages

##### Post-training.

SeerAttention can be directly applied to a pre-trained model.
In this case, only the weights of AttnGate are learned and updated, leaving the original model weights unchanged.
This method is highly efficient and cost-effective, requiring gradients solely for the AttnGate, and quickly converges using minimal calibration data.
The learned gate also allows for adjustable Top-k ratios during inference, providing a flexible tradeoff between accuracy and efficiency.

##### Fine-tuning.

SeerAttention can also be applied to long-context extension fine-tuning, enabling improved model performance and higher sparsity ratios. In practice, to ensure stable training, the AttnGate is first initialized using the post-training method before fine-tuning the entire model. During fine-tuning, we fix the Top-k ratio and use both the original training loss and the attention map MSE loss.

## 5 Experiments

In this section, we evaluate both the accuracy and efficiency of SeerAttention.
The accuracy is evaluated under two distinct scenarios: (1) post-training stage and (2) long-context extension fine-tuning stage.
For the efficiency evaluation, we present kernel-level and end-to-end latency speedup results across various sparse configurations.
In our current experiments, block-size $B$ for the model and kernel is fixed at 64 and AttnGate currently solely applies in the prefill stage.

##### Models, Tasks and Baselines.

We apply SeerAttention to the pre-trained models Llama-3.1-8B (Dubey et al., ) and Mistral-7B-v0.3 (Jiang et al., ) to assess its impact on LLM perplexity across different AttnGate designs and sparsity configurations. For perplexity evaluation, we use the PG19 (Rae et al., ) and Proof-pile () test splits. Following YaRN (Peng et al., ), 10 documents over 128k tokens are sampled from Proof-pile, while all documents exceeding 128k tokens from PG19 are selected. The input sequences are truncated to evaluation context length before feeding into the model.
We also conduct experiments using an instruction-tuned model, Llama-3.1-8B-Instruct, and compare SeerAttention with two state-of-the-art sparse attention methods, MoA (Fu et al., ) and MInference (Jiang et al., ), on the LongBench (Bai et al., ) benchmark, perplexity and efficiency. MoA uses an offline search scheme to apply static sparse patterns across different attention heads, while MInference dynamically generates sparse indices using heuristic methods for each head based on pre-defined sparse patterns.

##### Post-training Setup.

We use the RedPajama (Computer, ) dataset for calibration, chunked into 64k and 32k segments for Llama-3.1 and Mistral, respectively.We use a learning rate of 1e-3 with cosine decay and the global batch size of 16. The AttnGate is only trained in 500 steps using the ground truths from our customized FlashAttention kernel, with DeepSpeed (Rasley et al., ) stage 2 optimization on 4 A100 GPUs. As only AttnGate parameters are learned and updated in post-training, this process can be completed with hours.

##### Long-context Extension Fine-tuning Setup.

We extend the context size of a Llama-3-8B model from 8K to 32K, following the setup from YaRN (Peng et al., ), while introducing attention sparsity via SeerAttention. The Top-k number in the AttnGate is fixed during the forward pass to allow the model to adapt to the sparsity.
We use a learning rate of 1e-5 with linear decay and a global batch size of 8 on RedPajama dataset. The entire model weights are fine-tuned on 4 A100 GPUs with DeepSpeed stage 3 optimization.

### 5.1 Accuracy of Post-Training

##### Perplexity on Pre-trained Models.

Figure 4 shows the perplexity results on the Proof-pile dataset for both Llama-3.1-8B and Mistral-7B-v0.3 across different context lengths and sparsity ratios.
It should be noted that the results for each model come from the same checkpoint with trained AttnGates, and different sparsity ratios are achieved by adjusting the value of $k$ in the Top-k. The results show that SeerAttention only slightly increases perplexity as the sparsity ratio increases, compared to full attention. For instance, with the Mistral-7B model at a 32k context size, SeerAttention achieves a perplexity of 2.45, compared to the baseline of 2.29, despite introducing a significant 90% attention sparsity. Figure 4 also demonstrates that longer context lengths allow for greater sparsity with minimal accuracy degradation.

Figure: (a) Llama-3.1-8B
Refer to caption: x7.png

##### Perplexity Comparison with Related Works.

Table 1 compares the perplexity of SeerAttention at post-training with MoA and MInference, using the Llama-3.1-8B-Instruct model on the PG19 dataset.
For MoA, we adopt their ”KV Sparsity” in 0.5 which means ”Attention Sparsity” in 0.35.
For MInference, we use their official setup, where all attention heads choose the ”Vertical-Slash” sparsity pattern for Llama-3.1-8B-Instruct.
Since MInference dynamically generates sparse indices for each input, we record their average attention sparsity across different context lengths for comparison.
SeerAttention outperforms MoA and MInference even with higher sparisty in most cases, except at the 128k context length.
This is likely due to MInference applies varying sparsity per head, whereas the fixed sparsity ratio across all heads in SeerAttention.
Varying sparsity per head could be applied to SeerAttention for enhancement, which remains a topic for future work.

**Table 1: Comparing the perplexity of SeerAttention at post-training with MoA and MInference, using the Llama-3.1-8B-Instruct model on the PG19 dataset.**
|  |  |  | Evaluation Context Length |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | Sparsity $s$ | 8k | 16k | 32k | 64k | 128k |
| Original | 0.0 | 10.03 | 9.88 | 9.92 | 9.97 | 10.03 |
| MoA | 0.35 | 10.07 | 9.97 | 10.02 | 10.13 | OOM |
| MInference |  | 10.12 | 10.06 | 10.24 | 10.43 | 10.89 |
|  | $s=0.37$ | $s=0.55$ | $s=0.69$ | $s=0.80$ | $s=0.9$ |  |
| SeerAttention | 0.4 | 10.06 | 9.92 | 9.96 | 10.10 | 10.29 |
| 0.5 | 10.08 | 9.94 | 9.99 | 10.15 | 10.38 |  |
| 0.6 | 10.12 | 9.96 | 10.04 | 10.21 | 10.50 |  |
| 0.7 | 10.18 | 10.01 | 10.10 | 10.29 | 10.71 |  |
| 0.8 | 10.30 | 10.07 | 10.18 | 10.39 | 11.18 |  |
| 0.9 | 10.75 | 10.24 | 10.30 | 10.56 | 13.20 |  |

##### LongBench Evaluation.

To evaluate performance on instruction-following tasks, we conduct experiments on LongBench, a long-context understanding benchmark, and compare the results with MoA and MInference using the Llama-3.1-8B-Instruct model.
As shown in Table 2, SeerAttention consistently outperforms both MoA and MInference under similar or higher sparsity ratios.

**Table 2: Comparing the accuracy of SeerAttention at post-training with MoA and MInference on LongBench.**
|  |  |  | LongBench |  |  |
| --- | --- | --- | --- | --- | --- |
| Model | Attention | Sparsity $s$ | 0-4k | 4-8k | 8k+ |
| Llama-3.1-8B-Instruct | Original | 0.0 | 55.32 | 53.98 | 52.90 |
| MoA | 0.35 | 50.74 | 49.84 | 51.89 |  |
| MInference |  | 55.23 | 53.87 | 52.18 |  |
|  | $s=0.06$ | $s=0.25$ | $s=0.45$ |  |  |
| SeerAttention | 0.1 | 55.91 | 54.32 | 53.28 |  |
| 0.25 | 55.00 | 54.09 | 52.22 |  |  |
| 0.5 | 52.40 | 52.85 | 52.43 |  |  |

### 5.2 Accuracy of Long-Context Extension Fine-Tuning

We follow YaRN (Peng et al., ) to extend the context size of a Llama-3-8B model from 8k to 32k. We integrate SeerAttention into YaRN and compare the performance against the YaRN dense baseline and the post-training SeerAttention applied after YaRN.
Figure 1(a) presents the loss curves of the YaRN dense baseline and SeerAttention at 50% and 90% sparsity. The curve at 50% sparsity nearly overlaps with the baseline, while the curve at 90% sparsity shows slightly higher loss.
Table 3 displays the test perplexity on the PG19 and ProofPile datasets evaluated at a 32k context length. The YaRN dense baseline achieves perplexity scores of 8.79 and 2.46, respectively. Post-training SeerAttention results in increased perplexity.
When applying SeerAttention during the YaRN extension fine-tuning, it maintains near-lossless performance at 50% sparsity (with scores of 8.81 and 2.47), and even at 90% sparsity, the loss remains minimal.

**Table 3: Perplexity of YaRN baseline, SeerAttention after YaRN and YaRN with SeerAttention.**
|  | YaRN | Post-training SeerAttention after YaRN | YaRN with SeerAttention |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Sparsity | 0.0 | 0.5 | 0.6 | 0.7 | 0.8 | 0.9 | 0.5 | 0.6 | 0.7 | 0.8 | 0.9 |
| PG19 | 8.79 | 9.16 | 9.30 | 9.48 | 9.73 | 10.18 | 8.81 | 8.82 | 8.85 | 8.93 | 9.16 |
| Proof-pile | 2.46 | 2.53 | 2.57 | 2.61 | 2.68 | 2.85 | 2.47 | 2.47 | 2.48 | 2.51 | 2.60 |

### 5.3 Efficiency evaluation

We evaluate the efficiency of SeerAttention using our implementation of Triton (Tillet et al., ) kernels. We evaluate the kernel-level as well as end-to-end speedup using a Llama-3.1-8B-Instruct on a single A100 GPU. Results are compared to FlashAttention-2 (dense baseline), MoA and MInference.

#### 5.3.1 Kernel evaluation

##### Negligible Overhead in AttnGate and Top-k.

Figure 5 shows the kernel-level latency breakdown of SeerAttention. It demonstrates that the overhead introduced by the AttnGate and Top-k operations during inference is minimal. For instance, at a context length of 32k and a sparsity of 0.5, the AttnGate and Top-k contribute only 1% and 2% to the total latency, respectively. In the cases of 128k sequence length, the relative overhead almost diminishes.

##### Block-Sparse FlashAttention Kernel Speedup.

Figure 5 also shows that our kernel demonstrates linear speedup over various sparsity levels. At a sequence length of 128k with 90% sparsity, SeerAttention achieves a speedup of 5.47× compared with FlashAttention-2 on a single A100 GPU. While the current implementation is based on Triton, further performance gains are possible by optimizing the kernel using CUDA in future work.

Figure: (a) Sequence length = 8k
Refer to caption: x9.png

##### Compared to Related Works.

We compare the speedup of SeerAttention with MoA and MInference. MInference uses offline calibration to identify a pre-defined sparse pattern for each layer. For Llama-3.1-8B-Instruct model, MInference consistently uses ”Vertical-slash” pattern across all layers. During runtime, MInference will dynamically generate non-zero indices based-on their approximation algorithm. On the other hand, MoA uses ”A-shape” blocks as their sparse pattern and calibrate the shape parameters offline under given sparsity constraint.

Figure 6 shows the sparsity v.s. speedup plots of different methods on 8k, 32k, 128k sequences length, where the speedup baseline is FlashAttention-2.
The sparsity statistics were collected on PG19 datasets. For MoA, we generated the sparse configurations under their 0.5 overall ”KV-sparsity” constraints, which corresponds to an average of 0.35 sparsity in attention. The results demonstrates that SeerAttention outperforms both MoA and MInference in most cases. At 128k, the performance of all three methods converges, where the benefits of sparsity significantly outweigh the associated overhead.

Figure: (a) Sequence length = 8k
Refer to caption: x12.png

#### 5.3.2 End-to-End Speedup

To assess the end-to-end speedup of our method, we measured the average prefilling time, or time-to-first-token (TTFT), using the Llama-3.1-8B-Instruct model. Following the experimental setup used in MoA, we also recorded the average sparsity statistics for each method. The results show that SeerAttention consistently achieves lower latency compared to MInference, even with lower sparsity ratios. As for MoA, it requires an exhaustive search for different sparse configurations under varying sparsity constraints, which is time-consuming. Therefore, we only compared against its default configuration.

**Table 4: Time to First Token results (s).**
| Latency (Sparsity) | Evaluation Context Length |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
| 8k | 16k | 32k | 64k | 128k |  |
| FlashAttn-2 | 0.90 (0) | 1.95 (0) | 4.63 (0) | 10.09 (0) | 35.54 (0) |
| MoA | 1.29 (0.35) | 3.44 (0.35) | 10.34 (0.35) | 36.34 (0.35) | OOM |
| MInference | 2.33 (0.37) | 3.10 (0.65) | 4.68 (0.77) | 8.21 (0.86) | 14.38 (0.95) |
| SeerAttention | 0.78 (0.50) | 1.65 (0.60) | 3.60 (0.70) | 7.69 (0.80) | 13.37 (0.95) |

## 6 Analysis and Ablation

##### Visualization of Learned Attention Maps.

The AttnGate module automatically learns diverse sparse patterns without any prior knowledge or heuristics. Figure 7 shows several example outputs from AttnGate, including (a) ”A-shape,” (b) ”Vertical,” (c) ”Slash” with empty vertical spaces, (d) block sparsity along the diagonal, and (e) random patterns. These patterns not only encompass but also extend beyond those observed in previous works such as MoA and MInference, showcasing the versatillty of our learning based methods.

Figure: (a)
Refer to caption: x15.png

##### Analysis of FlashAttention with Max-Pooling Training Kernel.

We evaluate our customized FlashAttention kernel with maxpooling attention map for scalable training of SeerAttention by comparing against with PyTorch naïve manual attention implementation and FlashAttention-2. As shown in Figure 8(b), the PyTorch kernel runs out of memory (OOM) when the sequence length exceeds 4k, while our customized kernel costs similar peak memory usage compared to FlashAttention-2. Regarding latency, since PyTorch encounters OOM for sequences longer than 8K, the attention operations per head into a loop to assess kernel-level latency. Figure 8(b) shows that the latency overhead introduced by the additional pooling operation is minimal compared to the FlashAttention-2, while the PyTorch implementation suffers from a significant slowdown.

Figure: (a) Memory
Refer to caption: x20.png

##### RoPE Ablation.

In our experiments, we found that incorporating an additional RoPE module in the AttnGates, shown in Figure 2, significantly enhances its ability to extrapolate context length when training AttnGates.
Figure 9 shows the results with and without RoPE in AttnGate on PG19 with Llama-3.1-8B model using 8k length training data.
The AttnGate with RoPE shows very consistent performances on larger context lengths despite only trained with 8k length data.

Figure: (a) Sparsity = 0.5
Refer to caption: x22.png

##### Pooling Ablation.

Different pooling methods can be composed in the AttnGates. To study the best configuration, we test all the possible combinations of pooling in $\mathbf{Q}$ and $\mathbf{K}$ choosing from average, max and min pooling. There are in total 49 combinations. We train each configuration in post-training setting using Llama-3.1-8B model with 32k length data and test the perplexity on PG19 dataset with 128k evaluation context length. Figure 10 shows the Top-12 configurations. The one with average pooling on $\mathbf{Q}$ and max plus minpooling on $\mathbf{K}$ performs the best.
This may relate to the observations in LLM quantization that $\mathbf{K}$ often exhibits more outliers.

Figure: Figure 10: Perplexity of SeerAttention with different pooling methods.
Refer to caption: x25.png

## 7 Conclusion and Future Work

This paper presents SeerAttention, a new attention mechanism that learns and leverages the intrinsic sparsity in attention to boost long-context LLMs.
Our experiments demonstrate that SeerAttention not only outperforms previous approaches in post-training scenarios, but also excels in long-context fine-tuning, maintaining near-lossless accuracy even at high sparsity levels.
For future work, there are several promising directions to explore for improving and expanding the capabilities of SeerAttention.
One key area is enhancing the training methodologies for SeerAttention, such as applying SeerAttention in long-context continued pre-training with more training tokens to achieve higher sparsity without sacrificing accuracy.
Another important avenue is applying SeerAttention in the decoding stage of LLMs.
While this work primarily focuses on the prefill phase, it remains an open question whether the learned attention sparsity can similarly benefit the efficiency and accuracy of the decoding process.

## Appendix A Appendix

Figure: Algorithm 1 Customized FlashAttention with Max-pooling Kernel

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: * What if the pooling is done on the RoPE'd QK embeddings? Do we still need the RoPE in the AttnGate?
* I assume that using MSE loss to train AttnGate requires fewer training steps and samples. However, what if we train AttnGate using end-to-end LM loss?
* What is the sparsity pattern? Does the max-pooled attention pool prefer recent tokens (like sliding window attention) or does it selectively preserve "important tokens" (like needles in the haystack)?
* Will SeerAttention impact general benchmarks like MMLU?
soundness: 3
strengths: * SeerAttention operates at the block level, leading to high efficiency potential.
* SeerAttention learns the sparsity pattern during fine-tuning, which is more flexible than heuristic sparsity.
* An efficient SeerAttention kernel is provided and leads to a 5+x speedup over FlashAttention.
summary: This work introduces SeerAttention, which uses pooled embedding to compute attention block mask and uses the mask to do block sparse attention for reduced latency. To enable efficient learning of the gating network, a customized FlashAttention kernel is implemented that extracts the block-level ground
truth of attention map with minimum overhead. Experiments show that SeerAttention can achieve 90% sparsity ratio at 32k context length and offere 5.67x speedup over FlashAttention-2.
weaknesses: * I understand that SeerAttention allows users to adjust the balance between sparsity and accuracy. However, how to determine the sparsity (or Top-k) in practice is not clear to me. I suggest the authors include a discussion on how to choose sparsity in order to maintain high accuracy.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: NA
soundness: 2
strengths: - The proposed SeerAttention mechanism learns attention sparsity instead of relying on predefined patterns. This allows for better adaptation to different language tasks and models, as demonstrated by its performance across various experiments.
- The development of a customized FlashAttention implementation enables efficient learning of the gating network by extracting the block-level ground truth of the attention map with minimum overhead. This not only improves the training process but also contributes to the efficiency of the model.
summary: This paper presents SeerAttention, a attention mechanism learning intrinsic sparse attention in LLMs. It outperforms existing methods in post-training and fine-tuning, with efficient implementation and good adaptability to different context lengths and sparsity ratios.
weaknesses: - In section 3.1, the description is insufficient. It is hard to understand the proposed method only with the description in Section 3.1, but I found Figure 2 is easy to understand. I suggrest that more details around Figure 2 should be added in Section 3.1 to enhance understanding. For instance, the operations and significance of each component in the AttnGate module need to be elaborated.
- In section 4.2, many symbols related to FlashAttention are used without prior explanation, making it difficult for readers unfamiliar with FlashAttention to follow.
- The meaning of "max and min pooling on K matrix" is not clear. It should be precisely defined whether it is applied to different dimensions or in a specific way.
- When seq cannot be evenly divided by B (as the size of Q and K after downsampling is [seq/B,d]) during the growth of seq in token generation, the proposed method does not adequately discuss how to handle these situations. It is also unclear if the method is only applicable to the Prefilling stage and not suitable for the token decoding stage.
- After reading the paper, I found that only the blocks with larger attention scores from Atten Gate are used for calculation with flash attention during infernece. If an inference stage Triton kernel is implemented, its pseudo code should be included in the paper.
- There are inconsistencies in algorithm symbols, such as the use of A in Figure 3 and D in the appendix to represent Attention Score, which is confusing.
- Using only a part of the blocks for attention calculation raises questions about information retention. In very long sequences, discarding a large number of tokens may affect the model's ability to perform various text tasks. Since most experimental results report PPL metric, which may not be sensitive to token information sufficiency, additional experiments like Needle-In-A-Haystack or PassKey Retrivel could be added for a more comprehensive evaluation.
- It is not clear if the results in Figure 8 report to Prefilling or decoding stage delays.
- The data requirements of SeerAttention during Post-training and Fine-tuning are not specified.
- The inference delay of the model could be reported in Table 1 for a more complete comparison.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: - **Theoretical or Intuitive Explanation for Pooling Selection**: In the pooling selection section, the proposed approach utilizes average pooling on Q and a combination of max and min pooling on K. Could the authors provide a more theoretical or intuitive explanation for this choice? Specifically, given that many existing studies have highlighted the presence of excessively high attention scores on tokens with limited semantic information [1,2,3], could this phenomenon support the rationale behind the pooling selection?

- **Row-Normalization of Max-Pooled Attention Map**: In Section 4.1, the authors mention scaling the max-pooled attention map by row-normalizing it to sum to 1. Could the authors further elaborate on this step? As previous studies [4] suggest, it is critical to maintain the values of the attention distribution and avoid out-of-distribution attention scores (i.e., scores that are too large or too small) for informative tokens. Given this, might the proposed normalization inadvertently produce out-of-distribution attention scores? Have the authors experimented with alternative approaches, such as leaving the attention map unaltered or redirecting unused attention to the first attention sink token?

- **Applicability to Other LLM Architectures**: Can the proposed approach be adapted for other types of large language models, such as models based on the Mixture of Experts (MoE) architecture?

[1] Yu, Zhongzhi, et al. "Unveiling and harnessing hidden attention sinks: Enhancing large language models without training through attention calibration." arXiv preprint arXiv:2406.15765 (2024).

[2] Sun, Mingjie, et al. "Massive activations in large language models." arXiv preprint arXiv:2402.17762 (2024).

[3] Ge, Suyu, et al. "Model tells you what to discard: Adaptive kv cache compression for llms." arXiv preprint arXiv:2310.01801 (2023).

[4] Xiao, Guangxuan, et al. "Efficient streaming language models with attention sinks." arXiv preprint arXiv:2309.17453 (2023).
soundness: 3
strengths: - Addressing the accuracy-efficiency trade-off of LLMs during long-context inference is a critical challenge, especially given the recent trend of employing LLMs to tackle increasingly complex problems with sophisticated inference processes.

- The proposed approach, which learns a sparse attention distribution rather than relying on pre-defined attention patterns or heuristics to approximate sparsity, is intuitively effective in achieving a superior accuracy-efficiency trade-off.

- The customized FlashAttention modification enhances real-device efficiency, providing tangible performance improvements over the naïve implementation and reinforcing the feasibility of applying this approach in real-world applications.

- The paper is well-organized, making it accessible and straightforward to read and comprehend.
summary: This paper introduces a learnable method to leverage attention sparsity, enhancing the efficiency of large language models (LLMs) in long-context inference. Specifically, the authors developed a learnable gate that adaptively selects significant blocks within the attention map, treating non-significant blocks as sparse. Furthermore, to facilitate efficient training of this gate, the paper incorporates a customized FlashAttention implementation that efficiently extracts block-level ground truth from the attention map with minimal overhead. Experimental results demonstrate that the proposed technique achieves a superior accuracy-efficiency trade-off when processing long-context inputs compared to baseline approaches.
weaknesses: - **Limitations in Related Work Discussion**: Alleviating attention sparsity to enable long-context inference through attention optimization has been a significant area of research, with numerous discussions surrounding it. To provide a more comprehensive background, it would be beneficial for the authors to include additional related works on attention sparsity, particularly approaches that leverage KV cache eviction, encompassing both pre-defined patterns [1] and dynamic pattern attention sparsity [2,3,4]. 

- **Limited Context Length Evaluated**: Recent works utilizing attention optimization to enhance long-context inference capabilities have successfully extended the context length to settings up to 1M, significantly exceeding the pretrained context limits of LLMs, even without training adjustments [1,2]. However, this paper evaluates only up to a 128K context length. Given that the pretraining context length of the target Llama-3.1 model is also 128K, further experiments with longer context windows would help elucidate the proposed approach’s potential in handling extended contexts.

- **Deployment Cost of the Proposed Approach**: While existing approaches aim to improve the long-context capabilities of LLMs without requiring additional training by exploiting attention map sparsity, this work entails additional model tuning. Benchmarking both the post-training and fine-tuning variants of the proposed approach against existing training-free techniques would offer greater insight into the effectiveness of the required additional training.

- **Setting Mismatch Between Training and Memory Profiling**: In the pooling selection section, the authors state that they use average pooling for Q and a combination of max and min pooling for K. However, in the memory profiling results presented in Figure 8, only max pooling is considered in the training kernel. Would incorporating average pooling, as implemented in the proposed approach, significantly increase memory overhead? Additionally, what would the memory overhead look like during the inference process? How does the additional training introduced by this approach impact memory usage?

[1] Xiao, Guangxuan, et al. "Efficient streaming language models with attention sinks." arXiv preprint arXiv:2309.17453 (2023).

[2] Zhang, Zhenyu, et al. "H2o: Heavy-hitter oracle for efficient generative inference of large language models." Advances in Neural Information Processing Systems 36 (2024).

[3] Ge, Suyu, et al. "Model tells you what to discard: Adaptive kv cache compression for llms." arXiv preprint arXiv:2310.01801 (2023).

[4] Wang, Zheng, et al. "Model tells you where to merge: Adaptive kv cache merging for llms on long-context tasks." arXiv preprint arXiv:2407.08454 (2024).

---

confidence: 5
contribution: 3
initial_rating: 5
presentation: 4
questions: - Can you show how long the training took with 400 steps? I have a huge concern about the cost, and it may be unrealistic with a long context model. I want to ask you to add FLOPs measures or wall clock GPU-hour measures of training. And the training context length is quite short compared to up-dated long context LLMs (GLM4 = 1M, Llama3.1 = 128k). Therefore, I cannot agree that training will always end in a few hours because the training time grows quadratically.

- Can you show the generalizability of the attention gating network by extending the context window further than the trained (attention gate trained context, 64k~32k)? In Table 1, attention gating tends to fail to estimate exceeding trained context length (64k), which is still inside of LLM pretrained context length.

- Can you add a hyperparameter study about block size? I think the block size of each query and key dimension will show different patterns and have different sensitivity.

---

Therefore, while I agree that the paper is well-written and organized, I think it requires more revisions. I hope we can fix the problems and improve the paper partially during the discussion period.
soundness: 3
strengths: 1. Simple yet effective.

The overall paper is crystal clear and intuitive. I have no concerns about the presentation and methodology explanation.
summary: SeerAttention proposes an efficient attention implementation that reduces attention computation. It uses the learned attention-gating network, which is the same attention mechanism but with a pooled sequence.
weaknesses: 1. Require of training.

Even this paper claims that sparsity should be learned, but in most cases, it is quite expensive to do so with a large-scale language model such as llama3 70B or 405B.
I think the requirement of learning sparsity patterns is the extreme downside of this methodology, and this direction has also been tried in the past years and always not chosen in practice scenarios due to required training and additional engineering (TDSA: https://arxiv.org/abs/2110.11299, Reformer: https://arxiv.org/abs/2001.04451)

I have concerns that the efficiency of this approach is really better than MInference with training free context extension methods (Self-Extend) during the whole life cycle of LLM. I think if the method should be trained every time a new LoRA or new model is coming, its effectiveness will be quite limited and will have the same limitations as previous methods (TDSA, Reformer, Mamba: https://arxiv.org/abs/2312.00752, etc. (training requires methods).) Therefore, I suggest you add how the trained SeerAttention can adapt to new tasks or unseen weights (LoRA).

Furthermore, is there really no way to pool tokens and approximate attention patterns without training or minimal training (a few steps)? I think there should be some great way to do something like QUEST (https://arxiv.org/abs/2406.10774). I strongly suggest authors try to investigate or at least compare with simple heuristics (averaging, absolute maximum, QUEST-style reducing)

Additionally, the requirement for training may lead to bad generalizability, as shown by the perplexity exploding in Table 1's 128k context length, even if 128k is inside of Llama 3.1 's pretrained context window. 

2. Asymptotic memory complexity is now quadratic.

One of the most important parts of FlashAttention was reducing the memory complexity of the attention mechanism to linear from quadratic by fusing attention score computation. However, because of the requirement of top-k for each row, this method has to store an O (T^2) buffer before converting it to a sparse mask (indices). This peak memory consumption of the temporary buffer will be around 15.625GB for a 1M (Gemini's context length) sequence with 64 block sizes. And if we have to allocate additional buffers such as indices and sorted values for top-k operation (let's think the backend of top-k is merge-sort), we have to store at most 93.72GB (15.62GB for input values, 15.62GB for sorted values, 62.48 for indices). Therefore, I think authors should investigate the memory consumption of the algorithm more carefully and should investigate the way to reduce peak memory (e.g., increasing the block size of gating)

3. Limited methodological contribution.

I think the overall approach is quite similar with TDSA, but just pooling dimension is different. (hidden dim -> sequence dim). I think if there is some kind of justification or comparision between various pooling dim (hidden, head, sequence) and showing why pooling sequence dim is the most effective should be better than current form.

The reducing sequence dim approaches for generating attention masks already exist (QUEST, MInfernece), so I think we need some more justification for additional training in attention-gating networks.

4. Limited downstream task evaluation.

The only downstream task evaluation is the LongBench in Table 2. (I do not think PPL of PG19 can demonstrate actual downstream task performance because perplexity is well-known to have a weak correlation between instruction tasks.)

The context length of Table 2 is around 8 to 32k, which is within the range of their speedup claim (x5.67 on 32k). However, I think if we can show that the performance degradation is minimal or negligible in a longer context, such as 128k, the impact of this work can be more powerful (InfiniteBench: https://github.com/OpenBMB/InfiniteBench, LOFT: https://github.com/google-deepmind/loft, RULER: https://github.com/NVIDIA/RULER, etc.).

Meta-review: **Summary:** The paper introduces SeerAttention, which exploits block-level sparsity in attention maps by learning a gating function. After employing a gating network with block-wise pooling and a customized FlashAttention kernel, the approach can achieve up to 5.67x speedups (for 32k contexts) while maintaining competitive perplexity scores.

**Strength:** 

1. The idea of learning sparsity patterns using gating networks is straightforward yet effective.

2. It demonstrates strong speedups while maintaining long-context performance for up to 128k contexts.

3. It provides a customized FlashAttention kernel, enhancing both real-world applicability and efficiency.


**Weakness:**

1. Training cost concerns: Learning sparse attention requires additional training resources, making it less favorable compared to training-free sparse attention methods. Furthermore, a comprehensive comparison with both training-free sparse attention techniques and training-free context length extension methods is needed.

2. Limited novelty and lack of thorough analysis: The concept of learnable gating for pooling has been explored in prior works (e.g., TDSA), and adaptive sequence dimension compression has also been explored by methods such as QUEST and MInference. A more detailed analysis is necessary to clearly demonstrate the advantages of blockwise attention gating.

3. Clarity and writing issues: Certain methodological descriptions lack clarity, and there are issues with undefined or inconsistent symbols throughout the text.


**Reasons for the decision:**

While SeerAttention introduces an efficient mechanism for adaptive sparse attention and demonstrates promising results, it lacks a thorough analysis of its key advantages and contributions compared to training-free sparse attention methods and other sequence dimension compression techniques. Consequently, I am inclined to reject this paper.

Decision: Reject