HyperAIHyperAI

Command Palette

Search for a command to run...

スライディングウィンドウは線形アテンションを凌駕する

Alexia Jolicoeur-Martineau Pashmina Cameron Rhea Sanjay Sukthanker Emy Gervais

概要

二次アテンションの性質上、大規模言語モデル(LLM)は多くのメモリとエネルギーを消費する。新しいトークンごとにコストは前のトークンよりも増加し、追加のトークンごとにキーとバリューを無期限にメモリに保存する必要があり、持続不可能である。この二次スケーリング問題を解決するためにいくつかの代替案が提案されており、その一つがLLMを線形アテンションに改造することである。このアイデアは、低コストで最先端の性能を維持しながら二次スケーリング問題を解決する可能性を約束するため、多くの注目を集めている。しかし、この研究の系統は、より単純なベースラインと適切に比較されていない。本研究では、シンクを備えたスライディングウィンドウアテンション(SWA)が、ポストトレーニングされた線形アテンションモデルと同等以上の性能を発揮することを示す。この観察は、複数のLLMと様々な下流タスクにわたって行われる。長文脈推論タスク(針山干し草探しとBABILong)では、SWAは線形アテンションよりも大幅に高い性能(2倍から10倍)を達成する。SWAはポストトレーニングを必要とせず、非常に高速で、低メモリで済むため、非常に安価で信頼性の高い解決策となる。推論メモリコストを削減するために、ポストトレーニングされた線形モデルではなくSWAへの切り替えを強く推奨する。線形アテンションモデルはある程度の可能性を示しているかもしれないが、SWAに匹敵するためには、ゼロからのトレーニングまたは大規模なポストトレーニングが必要になる可能性が高い。

One-sentence Summary

Researchers from Microsoft Applied Sciences Group and an independent affiliation show that Sliding Window Attention (SWA) with sinks matches or outperforms post-trained Linear Attention across multiple LLMs and downstream tasks, achieving 2102\text{--}10210 higher performance on long-context reasoning while requiring no post-training and lower memory, making SWA a cheaper and more reliable alternative.

Key Contributions

  • Demonstrates that Sliding Window Attention (SWA) with attention sinks matches or outperforms post-trained linear attention models across multiple LLMs and downstream tasks, without requiring any post-training.
  • On long-context reasoning benchmarks (Needle-in-a-Haystack and BABILong), SWA achieves 2 to 10 times higher performance than linear attention, recovering 20% to 25% of baseline performance at context length 256, compared to 2.2% and 5% for the linear attention method LoLCATs.
  • Shows that SWA with sinks is a cheaper and more reliable alternative to linear attention post-training, offering higher decoding speed, lower memory cost, and 99% baseline performance recovery on short-context reasoning tasks.

Introduction

Large Language Models face a significant bottleneck: their attention mechanisms scale quadratically with context length, causing ever growing memory and compute costs due to the KV cache. Linear attention methods offer a linear alternative, but they suffer from lower expressivity, the difficult problem of deciding what to retain or forget, and expensive training requirements. Prior work like LoLCATs attempted to convert pretrained models to linear attention with minimal fine-tuning, yet comparisons against simpler baselines were incomplete.

The authors address this gap by directly comparing post-trained linear attention models to Sliding Window Attention (SWA) with attention sinks, a training-free approach. They demonstrate that SWA with sinks, which attends to the previous k tokens plus the first 4 tokens, matches or outperforms most linearized models on short-context reasoning tasks, recovering 99% of baseline performance. On long-context tasks, SWA achieves dramatically higher accuracy, recovering 20% and 25% of baseline performance on S-NIAH-3 and BABILong respectively, compared to LoLCATs' 2.2% and 5%. The authors show that pretrained models can use SWA at inference time without any post-training or specialized kernels, providing a simpler and more effective solution for fixed memory cost inference.

Method

Method

The authors build upon two complementary attention mechanisms to design an efficient hybrid architecture: Sliding Window Attention (SWA) and Linear Attention. Each addresses distinct limitations of standard softmax self-attention, and their combination enables sub-quadratic inference cost while preserving strong performance.

Sliding Window Attention with Sinks

Instead of attending to all previous tokens, Sliding Window Attention restricts each query to attend only to the previous www tokens. This constraint mirrors the local receptive field of convolutional networks: after lll layers, the effective receptive field grows to lwl \cdot wlw, allowing the model to aggregate information across distant positions through depth rather than direct pairwise attention. Formally, SWA computes:

xt=i=max(1,tw+1)texp(qtki/d)vii=max(1,tw+1)texp(qtki/d),t[1,,L].\mathbf{x}_t = \frac{\sum_{i = \max(1, t - w + 1)}^{t} \exp\left(\mathbf{q}_t \mathbf{k}_i^{\top} / \sqrt{d}\right) \mathbf{v}_i}{\sum_{i = \max(1, t - w + 1)}^{t} \exp\left(\mathbf{q}_t \mathbf{k}_i^{\top} / \sqrt{d}\right)}, \qquad t \in [1, \dots, L].xt=i=max(1,tw+1)texp(qtki/d)i=max(1,tw+1)texp(qtki/d)vi,t[1,,L].

Empirically, SWA improves long-term memorization and length extrapolation by forcing the model to learn dependencies beyond its local receptive field, rather than relying on a global attention pattern that may encourage shortcut learning.

However, a critical failure mode arises: large language models assign disproportionately high attention to the first few tokens, even when those tokens are semantically irrelevant. These so-called attention sinks serve as repositories for excess attention mass. If the sliding window moves past these sink tokens, performance degrades catastrophically. The authors adopt a simple and effective fix: in addition to the w4w - 4w4 tokens in the sliding window, the model always attends to the first s=4s = 4s=4 tokens. This guarantees that sink tokens remain visible at every position, preventing the collapse observed when they fall outside the window. Importantly, this work focuses exclusively on training-free SWA with fixed sinks, avoiding any additional post-training or learnable sink parameters.

Linear Attention

Linear Attention replaces the softmax kernel with a feature map ϕ\phiϕ such that exp(qtki)ϕ(qt)ϕ(ki)\exp(\mathbf{q}_t \mathbf{k}_i^{\top}) \approx \phi(\mathbf{q}_t) \phi(\mathbf{k}_i)^{\top}exp(qtki)ϕ(qt)ϕ(ki). This factorization enables the attention computation to be rewritten as a recurrent update:

xt=ϕ(qt)i=1tϕ(ki)viϕ(qt)i=1tϕ(ki)=ϕ(qt)stϕ(qt)zt,\mathbf{x}_t = \frac{\phi(\mathbf{q}_t) \sum_{i=1}^{t} \phi(\mathbf{k}_i)^{\top} \mathbf{v}_i}{\phi(\mathbf{q}_t) \sum_{i=1}^{t} \phi(\mathbf{k}_i)^{\top}} = \frac{\phi(\mathbf{q}_t) \mathbf{s}_t}{\phi(\mathbf{q}_t) \mathbf{z}_t},xt=ϕ(qt)i=1tϕ(ki)ϕ(qt)i=1tϕ(ki)vi=ϕ(qt)ztϕ(qt)st,

where the state variables are updated incrementally at each step:

st=st1+ϕ(kt)vt,zt=zt1+ϕ(kt).\mathbf{s}_t = \mathbf{s}_{t-1} + \phi(\mathbf{k}_t)^{\top} \mathbf{v}_t, \quad \mathbf{z}_t = \mathbf{z}_{t-1} + \phi(\mathbf{k}_t)^{\top}.st=st1+ϕ(kt)vt,zt=zt1+ϕ(kt).

This formulation yields O(1)\mathcal{O}(1)O(1) inference cost with respect to sequence length, since only the fixed-size state pair (st,zt)(\mathbf{s}_t, \mathbf{z}_t)(st,zt) needs to be stored and updated over time. This removes the growing memory footprint and latency associated with quadratic attention over long contexts.

The practical challenge lies in designing a kernel ϕ\phiϕ that balances three necessary properties: expressiveness, spikiness, and monotonicity. The authors adopt the Hedgehog kernel, which applies a learnable linear projection followed by a dual-sided exponential transformation:

ϕ(x)(exp(f(x)),exp(f(x))),\phi(x) \leftarrow \left(\exp(f(x)), \exp(-f(x))\right),ϕ(x)(exp(f(x)),exp(f(x))),

where fff is a linear projection from dimension DDD to D/2D/2D/2. This construction produces a kernel that is sufficiently expressive to capture complex attention patterns while maintaining the monotonic and spiky characteristics needed for stable recurrent updates.

Post-Training for Hybrid Attention

Training a linear attention Transformer from scratch is prohibitively expensive, and most existing software and hardware stacks are optimized for softmax attention. Instead, the authors convert pretrained quadratic-attention LLMs into linear attention models through post-training. Using Low-Rank Adaptation (LoRA), they linearize models with as little as 40M tokens of additional training, recovering a substantial portion of the baseline's performance. The key to making this work is combining an expressive kernel such as Hedgehog with a small Sliding Window Attention component. The SWA branch preserves local, high-fidelity information that the recurrent linear attention might otherwise blur, while the linear branch provides efficient global context aggregation. This hybrid design forms the foundation of the authors' approach, enabling efficient inference without sacrificing the quality of the original pretrained model.

Experiment

The experiments compare linear attention variants, sliding-window attention (SWA), and full attention across general knowledge, long-context reasoning, and efficiency benchmarks. SWA consistently outperforms linear methods on downstream tasks, recovering the most baseline performance with no training tokens, while linear approaches like LoLCATs require fine-tuning and still lag, especially at longer contexts. In long-context tasks, SWA maintains superior accuracy over linear methods, and in speed and memory tests, SWA is fastest with lower or comparable memory costs at smaller window sizes. Overall, SWA emerges as the most effective and efficient alternative to full attention.

Sliding Window Attention (SWA) consistently outperforms linear attention methods on general knowledge and reasoning benchmarks, recovering the most baseline performance on MMLU and nearly all of the average benchmark performance. While some linear methods like QRWKV6 match or slightly exceed SWA in specific cases, SWA achieves the best trade-off between performance and training efficiency, requiring no additional tokens. SWA achieves the highest average downstream performance in 9 out of 11 cases, with only marginal exceptions from LoLCATs and QRWKV6. SWA recovers 93.2% of MMLU baseline performance, the highest among all methods, and nearly all (99.0%) of the average baseline performance. SWA requires zero post-training tokens, whereas the next best efficient method, LoLCATs, uses 40M tokens to recover 83.2% of MMLU and 97.5% of average performance. QRWKV6 matches the baseline on MMLU for Qwen2.5-32B-Instruct, while SWA shows a slight drop, and DiJiang outperforms SWA on MMLU for Llama2.0-7B.

Sliding Window Attention (SWA) consistently outperforms linear attention variants on general knowledge and reasoning benchmarks, achieving the best average performance in most cases while requiring no fine-tuning tokens. SWA also recovers a high percentage of the baseline model's performance, particularly on average metrics, and is the most training-efficient option among the compared methods. SWA achieves the highest average downstream performance in 9 out of 11 cases, with only marginal exceptions for LoLCATs on Phi-1.5 and QRWKV6 on Qwen2.5. SWA recovers 99.0% of the average baseline performance and 93.2% of MMLU baseline performance without any fine-tuning tokens. LoLCATs, the closest competitor in training efficiency, requires 40M tokens to recover 83.2% of MMLU and 97.5% of the average baseline performance. On MMLU, SWA is the best performer across most base models, except for Llama2-7B where DiJiang slightly outperforms it.

Across all tested window sizes and context lengths, SWA consistently matches or outperforms LoLCATs and Liger-GLA on the Single Needle-in-a-Haystack tasks. At the longest context length, SWA retains a meaningful portion of full attention accuracy, while the other methods drop to near zero. SWA achieves equal or higher accuracy than LoLCATs and Liger-GLA at every window size and context length. At 4K context, SWA recovers 17.2-23% of full attention accuracy, whereas LoLCATs and Liger-GLA reach at most 5.8% and 0.8%. Larger window sizes generally improve accuracy for all models, but SWA maintains the largest advantage.

On the BABILong benchmark, LoLCATs(+SWA) slightly outperforms SWA at short context lengths (0K and 1K), but SWA shows a clear advantage at longer contexts (2K and 4K). Relative to full attention, both methods recover a substantial portion of accuracy at 0K, but at 4K SWA retains much more performance than LoLCATs. At 0K and 1K context, LoLCATs(+SWA) scores slightly higher than SWA (e.g., 56% vs 55% at 0K). At 2K and 4K context, SWA outperforms LoLCATs by a large margin (e.g., 15% vs 3% at 4K). At 0K, both methods recover about 74-76% of full attention accuracy, but at 4K SWA recovers 25% while LoLCATs recovers only 5%.

SWA consistently outperforms linear attention methods on general knowledge and reasoning benchmarks, achieving the best average performance in most cases while requiring no fine-tuning tokens, and recovers a high percentage of baseline performance, particularly on average metrics. On long-context tasks, SWA maintains a clear advantage over alternatives like LoLCATs and Liger-GLA, especially at longer context lengths where other methods drop to near zero accuracy, though LoLCATs slightly edges out SWA at very short contexts on BABILong. Overall, SWA offers the best trade-off between performance and training efficiency, with only marginal exceptions from specific methods on certain benchmarks.


AIでAIを構築

アイデアからローンチまで — 無料のAIコーディング支援、すぐに使える環境、最高のGPU価格でAI開発を加速。

AI コーディング補助
すぐに使える GPU
最適な料金体系

HyperAI Newsletters

最新情報を購読する
北京時間 毎週月曜日の午前9時 に、その週の最新情報をメールでお届けします
メール配信サービスは MailChimp によって提供されています