Learning to Focus: Focal Attention for Selective and Scalable Transformers

TL;DR

Focal Attention enhances transformer models by sharpening attention distributions, leading to better focus on relevant tokens, improved scalability, and significant performance gains on long-context tasks with fewer parameters and less data.

Contribution

This paper introduces Focal Attention, a novel method that sharpens attention distributions in transformers, improving scalability and performance especially on long-context tasks.

Findings

Achieves up to 42% fewer parameters for the same accuracy.

Reduces training data requirements by up to 33%.

Improves long-context task performance by 17% to 82%.

Abstract

Attention is a core component of transformer architecture, whether encoder-only, decoder-only, or encoder-decoder model. However, the standard softmax attention often produces noisy probability distribution, which can impair effective feature selection at every layer of these models, particularly for long contexts. We propose Focal Attention, a simple yet effective modification that sharpens the attention distribution by controlling the softmax temperature, either as a fixed hyperparameter or as a learnable parameter during training. This sharpening enables the model to concentrate on the most relevant tokens while suppressing irrelevant ones. Empirically, Focal Attention scales more favorably than standard transformer with respect to model size, training data, and context length. Across diverse benchmarks, it achieves the same accuracy with up to 42% fewer parameters or 33% less training data. On long-context tasks, it delivers substantial relative improvements ranging from 17% to 82%, demonstrating its effectiveness in real world applications.