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

TL;DR

RATTENTION targets the efficiency–accuracy Pareto frontier of local-global transformers by integrating Sliding Window Attention (SWA) with Residual Linear Attention (RLA). The resulting RAttention hybrid captures information from out-of-window tokens while preserving constant memory, enabling a small window size of $512$ to match or exceed full-attention performance on 3B and 12B models and across long-context tasks such as the RULER benchmark. The approach achieves favorable training and inference efficiency thanks to specialized kernels and a memory-aware chunked design, while remaining parameter-efficient by reusing existing projections. Open-source Pallas kernels and model codes are released to foster broader adoption and further research.

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.

Figures (7)

• Figure 1: RAttention combines Sliding Window Attention (SWA) for local context with Residual Linear Attention (RLA) to gather information from out-of-window tokens. Apart from 4 in-window tokens, RLA compresses the information of token 1 for query token 5; token 1,2 for query token 6.
• Figure 2: Interleaved state-saving pattern used in our training kernels. For every $m$ chunks, only the state of the last chunk (${\mathbf{S}}_{[m]}$) is stored in HBM. The intermediate chunk states are recomputed on-chip as needed, using the most recent stored state ${\mathbf{S}}_{[i-1]}$ from the previous group of $m$ chunks. This approach make it more flexible to balances memory I/O cost and matmul computation.
• Figure 3: Model specifications of our 3B and 12B models.
• Figure 4: Comparison of MMLU 5-shot performance scores across different window sizes at 3B scale with pretraining context length 4096. The horizontal purple dashed line represents the baseline using only global attention. The blue line shows Local-Global with sliding window attention (SWA), while the red line demonstrates the performance of Local-Global with RAttention. When window size $=0$, Local-Global with RAttention reduces to Local-Global with only linear attention.
• Figure 5: Main results at 12B scale with pretraining context length 8192. Performance of zero- and one-shot tasks are summarized in Average (0/1-shot).