RATTENTION: Towards the Minimal Sliding Window Size in Local-Global Attention Models

TL;DR

This paper introduces RATTENTION, a local-global attention variant that uses a specialized linear attention mechanism to effectively reduce window size while maintaining performance, improving efficiency especially in short-context scenarios.

Contribution

We propose RATTENTION, a novel local-global attention method that captures out-of-window tokens, enabling smaller window sizes without performance loss, and demonstrate its effectiveness at large scales.

Findings

RATTENTION with a window size of 512 matches full-attention performance.

Achieves better efficiency-performance tradeoff at 3B and 12B scales.

Maintains training speed comparable to existing methods.

Abstract

Local-global attention models have recently emerged as compelling alternatives to standard Transformers, promising improvements in both training and inference efficiency. However, the crucial choice of window size presents a Pareto tradeoff: larger windows maintain performance akin to full attention but offer minimal efficiency gains in short-context scenarios, while smaller windows can lead to performance degradation. Current models, such as Gemma2 and Mistral, adopt conservative window sizes (e.g., 4096 out of an 8192 pretraining length) to preserve performance. This work investigates strategies to shift this Pareto frontier, enabling local-global models to achieve efficiency gains even in short-context regimes. Our core motivation is to address the intrinsic limitation of local attention -- its complete disregard for tokens outside the defined window. We explore RATTENTION, a variant of local attention integrated with a specialized linear attention mechanism designed to capture information from these out-of-window tokens. Pretraining experiments at the 3B and 12B scales demonstrate that RATTENTION achieves a superior Pareto tradeoff between performance and efficiency. As a sweet spot, RATTENTION with a window size of just 512 consistently matches the performance of full-attention models across diverse settings. Furthermore, the recurrent nature inherent in the linear attention component of RATTENTION contributes to enhanced long-context performance, as validated on the RULER benchmark. Crucially, these improvements do not compromise training efficiency; thanks to a specialized kernel implementation and the reduced window size, RATTENTION maintains training speeds comparable to existing state-of-the-art approaches. We open-sourced our Pallas kernels along with model codes to facilitate further research effort.