Why Attend to Everything? Focus is the Key

Hengshuai Yao; Xing Chen; Ahmed Murtadha; Jin Li; Yasin Abbasi Yadkori; Shuai Shao; Changling Liu; Guan Wang

cs.CLMar 12, 2026

Why Attend to Everything? Focus is the Key

Authors: Hengshuai Yao, Xing Chen, Ahmed Murtadha, Jin Li, Yasin Abbasi Yadkori, Shuai Shao, Changling Liu, Guan Wang, +3 more

Organizations: Sapient Intelligence · University of Alberta · Tsinghua University

Abstract

Standard attention scales quadratically with sequence length. Efficient attention methods reduce this O(n^2) cost, but when retrofitted into pretrained models, they often degrade perplexity, downstream accuracy, or both. We introduce Focus, a method that learns which token pairs matter. Focus adds a small set of learnable centroids--as few as 148K parameters per layer--that act as gates: only token pairs belonging to the same centroid group attend to each other over long ranges. Focus is composable: it can be added to any pretrained model by training only the centroids while keeping all original weights frozen. Experiments show that composing Focus onto pretrained models yields zero degradation on downstream benchmarks across model sizes from 124M to 70B parameters and five attention architectures. Surprisingly, sparse Focus attention outperforms full attention at 124M scale (30.3 vs. 31.4 perplexity) and matches full attention when trained from scratch at 7B scale (13.82 vs. 13.89). Focus is also fast: top-k group membership gives a 2x speedup with better quality than the original pretrained model. Using our FlashAttention decomposition, Focus achieves an 8.6x speedup at 1M tokens without custom kernels.

Figures & tables

Method	Params	PPL $\downarrow$	HellaSwag	ARC-E	PIQA	LAMBADA
Pretrained (full attn)	0	42.8	31.1	39.5	62.5	32.6
Longformer ( Beltagy et al., 2020 )	0	38.9	30.0	37.5	58.9	6.6
Performer ( Choromanski et al., 2021 )	0	112.0	26.9	30.8	55.0	0.3
Routing Trans. ( Roy et al., 2021 )	0	37.4	29.6	38.3	58.4	6.4
Full attention FT	124M	36.4	30.0	37.8	59.9	7.8
Focus (ours)	100K	36.2	31.1	39.5	62.5	32.6

Table 1: Retrofit comparison on GPT-2 124M / PG-19. Focus is the only method that improves PPL and preserves all benchmarks.

Figure 1: Quality–speed Pareto frontier of efficient attention retrofits on GPT-2 124M / PG-19 (seq_len=64K). Y-axis is inverted: higher position means lower (better) PPL. Only Focus occupies the upper-right quadrant: better quality and faster than full attention.

Model	PG-19 PPL	HellaSwag	ARC-E	PIQA	LAMBADA
GPT-2 124M	42.8 $\to$ 34.2	31.1	39.5	62.5	32.6
GPT-2 774M	25.7 $\to$ 21.7	45.3	46.6	69.2	47.7
Mistral 7B	10.8 $\to$ 11.6	81.2	79.4	82.6	75.3
Qwen 2.5 7B	19.3 $\to$ 20.3	78.3	76.2	80.0	70.7
OLMo-2 7B	16.4 $\to$ 16.9	80.5	82.9	81.0	73.2
LLaMA-2 13B	11.7 $\to$ 11.7	79.6	76.5	80.4	76.6

Table 2: Focus composed onto seven models. Only centroids trained on PG-19; all pretrained weights frozen. PPL column shows pretrained

\to

Focus. Benchmark columns show Focus scores, which are identical to pretrained (zero degradation).

Method	Params	PPL $\downarrow$	HellaSwag	ARC-E	LAMBADA	PIQA
LoRA ( $r{=}4$ )	147K	31.6	$-0.4$	$-1.6$	$-2.1$	$-0.5$
LoRA ( $r{=}16$ )	590K	31.2	$-0.3$	$-1.6$	$-1.5$	$-0.4$
Centroids ( $d_{g}{=}16$ )	148K	34.2	$\pm 0.0$	$\pm 0.0$	$\pm 0.0$	$\pm 0.0$

Table 3: Focus vs LoRA on GPT-2 124M / PG-19.

Method	Params	PPL	TQA MC1	HellaSwag	ARC-E	LAMBADA
Instruct baseline	0	17.9	39.7	74.4	77.2	69.1
+ Centroids ( $K{=}2$ )	2.1M	18.0	40.0	74.4	77.2	69.1
+ LoRA ( $r{=}4$ , lr $=10^{-5}$ )	1.7M	16.1	40.1	72.6	76.4	69.1
+ LoRA ( $r{=}4$ , lr $=5{\times}10^{-5}$ )	1.7M	17.9	33.3	63.9	64.5	56.0
+ LoRA ( $r{=}4$ , lr $=10^{-4}$ )	1.7M	20.8	28.5	31.2	31.6	16.4

Table 4: Alignment preservation on Mistral-7B-Instruct (2000 training steps).

Model / Method	PPL	Params trained
GPT-2 124M
Full attention FT	31.4	124M
Focus FT	30.3	124M
GPT-2 Large 774M
Full attention FT	20.4	774M
Focus FT	20.7	774M

Table 5: PG-19 PPL across three scales (all parameters fine-tuned).

Context	1K	4K	16K	32K	65K	262K	1M
$K{=}4$ speedup	0.2 $\times$	0.5 $\times$	1.5 $\times$	2.2 $\times$	3.0 $\times$	4.0 $\times$	4.1 $\times$
$K{=}8$ speedup	0.2 $\times$	0.6 $\times$	1.8 $\times$	3.1 $\times$	4.7 $\times$	7.6 $\times$	8.6 $\times$

Table 6: Wall-clock speedup of Focus over full attention (both using FlashAttention) on H100-80GB.

Model	top- $k$	PPL	Speedup	Pairs retained	$\Delta$ vs Pretrained
GPT-2 124M	1 (argmax)	82.9	4.0 $\times$	26%	$+40.1$
	2	41.3	2.0 $\times$	60%	$\mathbf{-1.5}$
	3	47.2	1.3 $\times$	100%	$+4.4$
	4 (full)	47.2	1.0 $\times$	100%	$+4.4$
Mistral 7B	1 (argmax)	16.2	4.0 $\times$	25%	$+5.4$
	2	11.6	2.0 $\times$	73%	$\mathbf{+0.7}$

Table 7: Speed–quality tradeoff by varying top-

k

group membership at inference (

K{=}4

, PG-19). GPT-2 pretrained: 42.8 PPL; Mistral pretrained: 10.8 PPL.

Method	A	B	C	Outcome
Entropy + balance loss	Partial	✗	✗	Collapses by step 600
Stop-gradient on inputs	✗	✓	✗	Slow, not converging
EMA centroids + proj	✓	✗	✗	Proj erases structure
Recluster every 100 steps	Periodic	Periodic	✗	Balanced but unstable
Balance weight $\times$ 5	Partial	✗	✗	6 of 8 groups die
Sinkhorn (ours)	✓	✓	✓	Stable, semantic

Table 8: Three escape pathways and mitigations attempted.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Parameter	Value	Centroid PPL	Fine-tuned PPL	Dominance (centroid)	Dominance (full FT)
Groups $K$	4	36.8	30.1	40%	38%
	8	38.4	30.3	24%	23%
	16	40.4	30.4	20%	31%
	32	42.4	30.5	21%	30%
Window $w$	64	38.3	30.2	17%	17%
	128	38.4	30.2	26%	23%

Appendix

Table 9: Ablation study (GPT-2 124M, PG-19). Each row varies one hyperparameter. Fine-tuned PPL is stable (29.9–30.5) across all 16 configurations.

Method	PPL $\downarrow$	HellaSwag	ARC-E	PIQA	LAMBADA
Pretrained	42.8	31.1	39.5	62.5	32.6
SparQ [ Ribar et al., 2024 ]	52.8	31.3	39.4	62.4	34.3
SparQ (mean realloc.)	48.3	31.2	39.3	62.3	33.1
MagicPIG [ Chen et al., 2024 ]	52.8	31.3	39.4	62.5	34.0
Focus (ours)	36.2	31.1	39.5	62.5	32.6

Appendix

Table 10: Token-selection methods vs Focus on GPT-2 124M / PG-19 (

k{=}32

). Token-selection methods preserve downstream benchmarks but degrade PPL by 5–10 points. Focus improves PPL with zero benchmark degradation.

Explore similar work

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers
Jun 22, 2026Huzama Ahmad, Se-Young YunLong-Context Language Model Inference LLM Inference Acceleration
KAIST
Long Context Pre-Training with Lighthouse Attention
May 7, 2026Bowen Peng, Subho Ghosh, Jeffrey QuesnelleLanguage Model Pretraining Self-Attention
Nous Research
Depth-Attention: Cross-Layer Value Mixing for Language Models
Jun 3, 2026Boyi Zeng, Yiqin Hao, Zitong Wang +7Efficient Transformer Inference Decoder-Only Language Models
LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · Sun Yat-sen University · Shanghai Innovation Institute +2

Jun 22, 2026J/K move · Enter open · S save

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers

Huzama Ahmad, Se-Young Yun

KAIST

Long-Context Language Model Inference LLM Inference Acceleration Sparse Attention Block-Sparse Attention

Open paper AskPDF arXiv

Long contexts have become standard in pretrained LLMs, yet they remain expensive to run: prefill compute grows quadratically with sequence length, and every decode step re-reads a key-value cache that grows linearly with it. Sparse attention cuts these costs by attending only to a relevant subset of past tokens, but selecting that subset is itself expensive. We present SpotAttention, a lightweight selector that attaches to a frozen pretrained transformer and learns by KL distillation to estimate its attention distribution. The selector picks the top-K keys each query attends to, and because its estimate is a calibrated distribution, a dual top-p rule reads the per-query, per-layer budget directly from it. Across Qwen3 (dense, 4B-32B) and Qwen3.5 (hybrid linear/full attention, 4B-9B), SpotAttention matches dense accuracy at contexts up to 128K tokens, eight times the training length. Decode at L=128K runs 3.9x faster than FlashAttention and 1.8x faster than Twilight, the strongest training-free baseline. Quantizing the selector's K-cache to INT4 or FP4 microscale shrinks it 3.5x at no accuracy cost.

Abstract

Explore similar work

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers

Long Context Pre-Training with Lighthouse Attention

Depth-Attention: Cross-Layer Value Mixing for Language Models