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 ↓
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 → 34.2
31.1
39.5
62.5
32.6
GPT-2 774M
25.7 → 21.7
45.3
46.6
69.2
47.7
Mistral 7B
10.8 → 11.6
81.2
79.4
82.6
75.3
Qwen 2.5 7B
19.3 → 20.3
78.3
76.2
80.0
70.7
OLMo-2 7B
16.4 → 16.9
80.5
82.9
81.0
73.2
LLaMA-2 13B
11.7 → 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 → Focus. Benchmark columns show Focus scores, which are identical to pretrained (zero degradation).
Method
Params
PPL ↓
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 ( dg=16 )
148K
34.2
±0.0
±0.0
±0.0
±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×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 ×
0.5 ×
1.5 ×
2.2 ×
3.0 ×
4.0 ×
4.1 ×
K=8 speedup
0.2 ×
0.6 ×
1.8 ×
3.1 ×
4.7 ×
7.6 ×
8.6 ×
Table 6: Wall-clock speedup of Focus over full attention (both using FlashAttention) on H100-80GB.
Model
top- k
PPL
Speedup
Pairs retained
Δ vs Pretrained
GPT-2 124M
1 (argmax)
82.9
4.0 ×
26%
+40.1
2
41.3
2.0 ×
60%
−1.5
3
47.2
1.3 ×
100%
+4.4
4 (full)
47.2
1.0 ×
100%
+4.4
Mistral 7B
1 (argmax)
16.2
4.0 ×
25%
+5.4
2
11.6
2.0 ×
73%
+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 × 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 ↓
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.
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.
Training causal transformers at extreme sequence lengths is bottlenecked by the quadratic time and memory of scaled dot-product attention (SDPA). In this work, we propose Lighthouse Attention, a training-only symmetrical selection-based hierarchical attention algorithm that wraps around ordinary SDPA and can be easily removed towards the end of the training. Our hierarchical selection is also gradient-free, which exempts us from dealing with a complicated and potentially inefficient backward pass kernel. Our contribution is three-fold: (i) A subquadratic hierarchical pre- and post-processing step that does adaptive compression and decompression of the sequence. (ii) A symmetrical compression strategy that pools queries, keys and values at the same time, while preserving left-to-right causality, which greatly improves parallelism. (iii) A two stage training approach which we pre-train for the majority of the time with Lighthouse Attention and recover a full attention model at the end with a short training. We run preliminary small scale LLM pre-training experiments that show the effectiveness of our method compared to full attention training with all other settings matched, where we achieve a faster total training time and lower final loss after the recovery phase. Full code is available at: https://github.com/ighoshsubho/lighthouse-attention
Self-attention selects information freely across the sequence, but across depth, Transformers merely add each layer's output to the residual stream, so later layers cannot selectively reuse earlier-layer representations. Recent cross-layer methods improve this flow but operate on hidden states outside attention, adding state beyond the key-value cache at inference--a cost that becomes increasingly salient as modern LLMs compress the cache with grouped-query and multi-head latent attention. We introduce Depth-Attention, which performs this selection inside the attention module itself: before a layer attends over the sequence, its query attends over the keys of earlier layers at the same token position and mixes their values into the value that self-attention then reads. Because Depth-Attention reuses the standard attention queries, keys, and value-cache slots, storing depth-mixed values in place of the original values, it adds no parameters and introduces no persistent inference state beyond the standard key-value cache--the same cache size as a vanilla decoder and less than hidden-state-based cross-layer methods. On Qwen3-style decoders at 1.5B and 3B parameters, Depth-Attention attains the lowest perplexity and the highest average downstream accuracy, improving over the vanilla Transformer by up to 2.3 accuracy points and surpassing strong cross-layer baselines in perplexity and average accuracy, while adding under 0.01% extra arithmetic FLOPs and no additional persistent inference state. The gains hold from 360M to 3B parameters and extend to looped Transformers.
Boyi Zeng, Yiqin Hao, Zitong Wang +7
LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · Sun Yat-sen University · Shanghai Innovation Institute +2