Technical Summary: Short Window Attention Enables Long-Term Memorization

Problem Statement

Modern Large Language Models (LLMs) based on softmax attention rely on a Key-Value (KV) cache that grows linearly with sequence length, incurring prohibitive compute and memory costs for long-context tasks. Conversely, linear Recurrent Neural Networks (RNNs) and State Space Models (SSMs) maintain constant memory and compute costs but often suffer from inferior recall abilities compared to Transformers.

Recent hybrid architectures attempt to combine local sliding window attention (SWA) with global attention or linear RNN layers to balance efficiency and performance. However, the impact of the sliding window length on the interplay between local and long-term memory remains under-studied. Specifically, it is unclear how the window size influences the model's ability to learn long-term dependencies when linear RNNs are present, and whether larger windows inherently improve performance or hinder the model's ability to extrapolate to longer contexts.

Methodology

The authors investigate hybrid architectures that interleave Sliding Window Attention (SWA) layers with linear RNN layers, specifically utilizing xLSTM (a variant of mLSTM with matrix memory) as the long-term memory component. The proposed architecture, termed SWAX, alternates between SWA layers and xLSTM layers in a 1:1 ratio.

The study focuses on two primary methodological contributions:

1. Window Size Analysis: The authors train SWAX models with fixed sliding window sizes ranging from 128 to 2048 tokens. They evaluate these models on short-context reasoning benchmarks (e.g., ARC, Hellaswag, MBPP) and long-context retrieval tasks (e.g., RULER Needle-In-A-Haystack).
2. Stochastic Window Training: To address the trade-off between short-context precision (favored by larger windows) and long-context extrapolation (favored by smaller windows), the authors propose a training strategy where the sliding window size is stochastically sampled for each batch. In their experiments, the window size alternates between a short size (128) and a long size (2048) with a specific probability (0.5 for 1.4B models, 0.75 for 7B models). Additionally, they employ an annealing schedule where the stochastic sampling is removed for the final 10% of training to stabilize short-context performance.

The experiments are conducted on models with 1.4B and 7B parameters, trained on 150 billion tokens with a fixed sequence length of 16k during pre-training.

Key Contributions

1. Counter-Intuitive Finding on Window Size: The paper demonstrates that for hybrid architectures combining SWA and linear RNNs, larger sliding windows can degrade long-context performance. Contrary to the belief that larger windows provide more context, the authors find that large windows (e.g., 2048) allow the model to rely too heavily on the precise local softmax attention, causing the linear RNN layers to under-train on long-term dependencies.
2. Mechanism of Short-Window Benefit: Shorter windows (e.g., 128) force the model to rely on the linear RNN layers for information retrieval that falls outside the local window. This "supervisory signal" compels the linear layers to specialize in modeling long-term dependencies, thereby improving the model's ability to extrapolate to sequence lengths far exceeding the training window.
3. Stochastic Training Strategy: The authors introduce a training procedure that stochastically varies the window size. This approach forces the model to leverage both the short-term precision of the window and the long-term memory of the RNN, effectively acting as a form of dropout on the attention mechanism.

Results

• Long-Context Performance: Models trained with a fixed large window (2048) exhibit a catastrophic drop in performance on long-context retrieval tasks (e.g., RULER NIAH) when tested on sequences longer than the window. In contrast, models trained with shorter windows (128) maintain high accuracy even at sequence lengths of 131k tokens. The stochastic training method achieves performance comparable to or better than the short-window-only models on these long-context tasks.
• Short-Context Performance: While short windows (128) alone sometimes underperform on short-context reasoning benchmarks compared to larger windows, the stochastic training method recovers this performance. At both 1.4B and 7B scales, stochastic training yields short-context performance comparable to, or in some cases better than, models trained with a fixed 2048 window.
• Generalization: The findings hold across different linear attention backbones (xLSTM and Gated DeltaNet) and different hybrid configurations (alternating SWA with full attention layers). The stochastic training strategy consistently outperforms fixed-window baselines on both short and long-context benchmarks.

Significance and Claims

The paper claims that the choice of sliding window size in hybrid architectures is not merely a trade-off between computational cost and performance, but a critical factor in how the model learns to distribute memory responsibilities.

The authors assert that:

• Short windows are essential for training long-term memory: By restricting the local attention field, the model is forced to utilize its linear RNN layers for long-range dependencies, leading to better length extrapolation.
• Stochastic training resolves the trade-off: By dynamically changing the window size during training, the model learns to utilize both the high-fidelity local attention for short contexts and the global linear memory for long contexts.
• Over-reliance on local attention is detrimental: If a model is trained with a large window, it fails to learn the necessary mechanisms to retrieve information from beyond that window, leading to poor generalization on longer sequences.

The work suggests that for hybrid architectures to achieve state-of-the-art long-context performance without unbounded compute costs, the training procedure must actively prevent the model from over-relying on local attention mechanisms, even if those mechanisms are more precise for short-range tasks.