Why Attend to Everything? Focus is the Key

TL;DR

Focus introduces a learnable gating mechanism with centroids to selectively attend token pairs, enabling efficient, high-quality attention in pretrained models without degradation and with significant speedups.

Contribution

It presents Focus, a novel method that adds minimal learnable parameters to pretrained models to improve attention efficiency while maintaining or improving performance.

Findings

Focus maintains zero degradation on downstream benchmarks.

Sparse Focus attention outperforms full attention at 124M scale.

Focus achieves up to 8.6x speedup with FlashAttention decomposition.

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.