# More Value per Key: Asymmetric Sparse Attention for Faster LLM Decoding

> SAGA decouples key and value head counts to exploit sparse attention's shifted bottleneck, achieving over 2x decoding speedup at 128K context with near-baseline quality.

- **Source:** [arXiv](https://arxiv.org/abs/2610.04753)
- **Published:** 2026-10-10
- **Permalink:** https://picx.dev/p/DRD92V
- **Whiteboard:** https://picx.dev/p/DRD92V/image

## Summary

## Summary (Overview)

- **Key Contribution**: The paper introduces **Sparse Asymmetric Group-Query Attention (SAGA)**, an architectural modification that decouples the number of key heads ($h_k$) from value heads ($h_v$) in attention mechanisms, exploiting the observation that under sparse attention, the computational bottleneck shifts from probability-value multiplication to query-key multiplication.
- **Atop-N Attention**: The authors propose **Approximate Top-N (Atop-N) attention**, a simple, portable sparse attention method for LLM decoding that uses threshold-based selection (updated periodically) rather than expensive sorting, reducing the probability-value cost from $O(h l_{kv} d)$ to $O(h N d)$.
- **Speedup Results**: SAGA combined with Atop-N achieves **end-to-end decoding speedups exceeding 2×** over full-attention GQA baselines at long contexts (128K tokens) on Llama 3.2-1B and Qwen2.5-1.5B models.
- **Quality Retention**: Models trained from scratch with SAGA nearly match the quality of comparable GQA variants, and distilled pretrained models retain **96.3%** of the teacher's aggregate RULER score.
- **Practical Conversion**: The paper presents a **closed-form initialization method** (based on minimum-weight perfect matching) for converting pretrained GQA models to SAGA, followed by short distillation, avoiding costly retraining.

## Introduction and Theoretical Foundation

### Background

Autoregressive generation in Large Language Models (LLMs) is constrained by the memory and computational demands of attention mechanisms. The quadratic scaling of attention with sequence length remains a significant bottleneck, limiting context window sizes. Key-Value (KV) caching provides speedups by reusing key/value computations, and Group-Query Attention (GQA) reduces the number of key-value heads to accelerate inference.

### Core Theoretical Insight

The paper's central observation is that **under sparse attention, the computational bottleneck shifts**:

> "Under sparse attention, the bottleneck shifts from the probability-value multiplication to the query-key multiplication."

This asymmetry arises because:
- Sparse attention (like top-N) reduces the probability-value cost to accessing only $N$ values per query head
- The number of key-value heads $h_{kv}$ affects complexity **only** through the query-key multiplication: $\mathbf{A} = \frac{\mathbf{QK}^\top}{\sqrt{d}}$

### Attention Mechanism Formalization

The standard multi-head attention is defined as:

$$
\mathbf{Q} = \mathbf{XW}_Q, \quad \mathbf{K} = \mathbf{XW}_K, \quad \mathbf{V} = \mathbf{XW}_V, \tag{1}
$$

$$
\mathbf{A} = \frac{\mathbf{QK}^\top}{\sqrt{d}}, \quad \mathbf{P} = \operatorname{softmax}(\mathbf{A}), \quad \mathbf{O} = \mathbf{PV}, \tag{2-4}
$$

where $h$ is the number of heads, $D$ and $d$ are hidden and head dimensions, and $l_q$, $l_{kv}$ are query and key-value sequence lengths.

### Top-N Attention

Top-N attention identifies indices of the $N$ largest values in each row of $\mathbf{A}$:

$$
\mathbf{i}_{j,k,1:N} = \operatorname{TopN}(\mathbf{A}_{j,k,1:l_{kv}}), \tag{6}
$$

masking all other entries before softmax. However, naive top-N requires sorting, which negates computational benefits.

## Methodology

### Atop-N Attention

The proposed **Atop-N** method uses an **approximate threshold approach**:
1. Given the $N$-th largest attention value as a threshold, top-N entries are identified via simple per-element comparison (parallel and portable, much faster than sorting)
2. Since attention distributions shift marginally between consecutive decoding steps, exact thresholds are computed **once every $n$ tokens** and reused, amortizing the cost

The sparse attention computation is:

$$
\mathbf{O} = \mathrm{softmax}(\mathbf{A}[\mathbf{i}])\mathbf{V}[\mathbf{i}], \tag{7}
$$

reducing probability-value cost from $O(h l_{kv} d)$ to $O(h N d)$ when $N \ll l_{kv}$.

### SAGA Architecture

SAGA decouples key and value head counts:
- **Reduced $h_k$**: Directly reduces memory bandwidth for loading keys during query-key multiplication (the dominant latency bottleneck under Atop-N)
- **Large $h_v$**: Preserves model capacity since only top-N values are accessed regardless of $h_v$
- **Additional benefits**: Fewer key heads yield smaller key projection $\mathbf{W}_K$ (less per-token compute) and smaller KV cache footprint

### GQA-to-SAGA Conversion

#### Stage 1: Closed-Form Initialization

For converting a GQA model with $h_k^T$ key heads to SAGA with $h_k^S$ key heads (where $h_k^T = m \cdot h_k^S$), the objective is to minimize squared logit reconstruction error:

$$
\min_{\sigma} \min_{\mathbf{M}^{[1:h_k^S]}} \sum_{i=1}^{h_k^S} \mathcal{L}_{\sigma_{2i-1},\sigma_{2i}}(\mathbf{M}^i), \tag{14}
$$

where the pairwise loss for merging key heads $a$ and $b$ is:

$$
\mathcal{L}_{a,b}(\mathbf{M}) := \sum_{j \in \mathcal{Q}_a} \left\| \mathbf{Q}^j(\mathbf{XM})^\top - \mathbf{Q}^j(\mathbf{XW}_K^a)^\top \right\|_F^2 + \sum_{j \in \mathcal{Q}_b} \left\| \mathbf{Q}^j(\mathbf{XM})^\top - \mathbf{Q}^j(\mathbf{XW}_K^b)^\top \right\|_F^2, \tag{13}
$$

**Lemma 4.1** provides the closed-form optimal merged key projection:

$$
\hat{\mathbf{W}}_K(a,b) = \left(\mathbf{W}_K^a \mathbf{S}_a + \mathbf{W}_K^b \mathbf{S}_b\right)(\mathbf{S}_a + \mathbf{S}_b)^{\dagger}, \tag{15}
$$

where $\mathbf{S}_r := \sum_{j \in \mathcal{Q}_r} (\mathbf{Q}^j)^\top(\mathbf{Q}^j)$ and † denotes the Moore–Penrose pseudoinverse.

**Theorem 4.2** establishes that the optimal pairing of key heads is found via **minimum-weight perfect matching** on a complete weighted graph $G = (V, E, w)$ where edge weights are $w(a,b) := \mathcal{L}_{a,b}(\hat{\mathbf{W}}_K(a,b))$.

#### Stage 2: Distillation Fine-Tuning

The student model is trained with a combination of KL divergence from the teacher and hard-label cross-entropy, exposed to increasing context lengths over ~23K iterations.

## Empirical Validation / Results

### Inference Speedups

**Figure 2** shows end-to-end speedups over baseline full attention:

| Configuration | Speedup at 128K (batch=1) | Speedup at 128K (batch=16) |
|---|---|---|
| Llama 3.2-1B (H200) | 2.1× | 2.6× |
| Qwen2.5-1.5B (B200) | 2.2× | 2.2× |

Key findings:
- Merging key heads alone (SAGA full attention) provides consistent speedup (up to 1.47× at batch 16, 128K context)
- Atop-N adds overhead at short sequences but substantial gains at long contexts
- Flash attention provides little benefit during autoregressive decode (matrix-vector product, no quadratic matrix to tile)

### Quality Retention with Atop-N

On Qwen2.5-1.5B across RULER benchmark (8K, 16K, 32K contexts):
- Atop-N retains **83–97%** of full attention quality
- Degradation concentrated in multi-key retrieval tasks (require attending to many dispersed keys)
- Alternative eviction-based methods (SW+S, H2O, KIVI) fall below 50% aggregate RULER accuracy at 8K context

### Training from Scratch (SmolLM2-360M)

SAGA with $(h_k, h_v) = (1, 8)$:
- **Closely matches** the high-capacity baseline $(8, 8)$ quality
- **Outperforms** both matched-key-budget $(1, 1)$ and parameter-matched $(4, 4)$ baselines
- Uses a little over half the KV cache of $(8, 8)$
- A Sigma-style variant reducing key-head dimension $(h_k, d_k) = (2, 32)$ reaches 41.7% validation accuracy vs. 42.7% for SAGA, confirming that reducing key head count is preferable to reducing key dimension

### Distillation Results

**Table 1: General-capability benchmarks**

| Benchmark | Qwen2.5-1.5B Baseline | Qwen2.5-1.5B SAGA | Llama 3.2-1B Baseline | Llama 3.2-1B SAGA |
|---|---|---|---|---|
| HellaSwag (acc_norm ↑) | 0.678 | 0.644 | 0.643 | 0.583 |
| WikiText (word ppl ↓) | 12.10 | 12.43 | 11.98 | 13.09 |

The distilled SAGA student retains **96.3%** of the teacher's aggregate RULER score across 8K, 16K, and 32K contexts.

### Large-Scale Validation (Llama3.1-70B-Instruct, Atop-N only)

**Table 2: GPQA accuracy with varying N**

| N | GPQA Diamond @16K (↑) | GPQA Main @8K (↑) |
|---|---|---|
| Baseline | 0.4343 | 0.4241 |
| 2048 | 0.3939 | 0.4129 |
| 1024 | 0.3838 | 0.4330 |
| 512 | 0.3586 | 0.4241 |

GPQA Main accuracy is unharmed even at $N = 512$, while GPQA Diamond incurs some degradation.

## Theoretical and Practical Implications

### Theoretical Significance

1. **Computational Bottleneck Analysis**: The paper provides a formal analysis showing that sparse attention fundamentally changes the computational profile of attention, shifting the dominant cost from probability-value to query-key multiplication.

2. **Architecture-Sparsity Co-design**: The work demonstrates that attention architectures should be co-designed with their inference-time sparsity patterns—a principle that could guide future LLM design.

3. **Optimal Key Head Merging**: The minimum-weight perfect matching formulation (Theorem 4.2) provides a principled, closed-form approach to reducing key heads while minimizing attention logit distortion.

### Practical Implications

1. **Immediate Deployment**: The conversion method enables practitioners to benefit from SAGA without costly retraining, requiring only a short distillation phase (~23K iterations).

2. **Complementary to Existing Methods**: SAGA is orthogonal to KV cache quantization, eviction methods, and other compression techniques, allowing combination for further gains.

3. **Hardware Efficiency**: The approach works across GPU architectures (H200, B200) and batch sizes, with benefits growing at longer contexts—the regime where LLMs are most constrained.

4. **Parameter Budget Reallocation**: The findings suggest that for sparse attention inference, parameter capacity is better allocated to value heads than key heads, informing future architecture design choices.

## Conclusion

The paper introduces **SAGA (Sparse Asymmetric Group-Query Attention)**, an attention architecture that decouples key and value head counts to exploit the shifted computational bottleneck under sparse attention. Combined with **Atop-N**, a simple approximate top-N attention method, SAGA achieves:

- **>2× end-to-end decoding speedups** at long contexts (128K)
- **Near-baseline quality** when trained from scratch or distilled from pretrained models
- **Practical deployability** via closed-form initialization and short distillation

### Limitations and Future Directions

1. **Kernel Integration**: Current implementation relies on simple Atop-N; hardware-aware sparse attention kernels could yield further speedups
2. **Capacity Compensation**: Distilled models could compensate for reduced key heads by increasing value head count or dimension
3. **Scale Validation**: Quality experiments limited to ≤1.5B parameters; full-scale pretraining validation remains open
4. **Broader Co-design**: The authors advocate for co-designing attention architectures with inference-time sparsity patterns as a standard consideration in LLM design

> "We believe that co-designing attention architectures with their inference-time sparsity patterns is a fertile direction, and hope that SAGA serves as a step toward making this a standard consideration in LLM design."

---

_Markdown view of https://picx.dev/p/DRD92V, served by PicX — AI-generated visual whiteboard summaries of research papers._