Transformerのセルフアテンション計算を固定幅のローカルウィンドウに制限し、計算量をO(n²)からO(n×w)に削減する手法。Mistral 7Bで採用されたことで広く知られ、Longformer、BigBirdなどの先行研究を実用レベルに昇華させた。ロングコンテキスト対応の基盤技術の一つ。
Sliding Window Attention(SWA)は、Transformerのセルフアテンション機構において、各トークンが参照する範囲を固定幅のウィンドウ(例: 4,096トークン)に制限する手法である。標準的なFull Attentionでは全トークンペアに対してアテンションスコアを計算するためO(n²)の計算量が必要だが、SWAではO(n×w)(wはウィンドウ幅)に削減される。
Mistral 7Bが2023年9月にSWAを採用し、7Bパラメータのモデルで32Kトークンのコンテキストを効率的に処理できることを実証した。このアプローチは「各トークンが直近のw個のトークンのみを参照する」というシンプルな制約だが、Transformerの多層構造により、上位レイヤーでは間接的に広い範囲の情報が統合される。
| アテンション方式 | 計算量 | メモリ使用量 | 参照範囲 | 代表モデル |
|---|---|---|---|---|
| Full Attention | O(n²) | O(n²) | 全トークン | GPT-4, Claude 3 |
| Sliding Window | O(n×w) | O(n×w) | ウィンドウ幅w | Mistral 7B |
| Longformer | O(n×(w+g)) | O(n×(w+g)) | ローカル+グローバル | Longformer |
| BigBird | O(n×(w+g+r)) | O(n×(w+g+r)) | ローカル+グローバル+ランダム | BigBird |
| Ring Attention | O(n²/P) | O(n²/P) | 全トークン(分散) | 研究段階 |
| Sparse Attention | O(n√n) | O(n√n) | ストライドパターン | GPT-3(初期) |
wはウィンドウ幅、gはグローバルトークン数、rはランダムアテンション数、Pはデバイス数を示す。
Mistral 7Bでは以下の設計でSWAを実装している:
この設計により、Mistral 7Bは13Bパラメータ級のLLaMA 2と同等以上の性能を、約半分のメモリ使用量で達成した。
Longformer(2020年、Allen AI)はSWAの先駆的研究である。Longformerは以下の3種類のアテンションパターンを組み合わせる:
Longformerは最大4,096トークンの処理に焦点を当てていたが、同じ原理をMistral 7Bは32Kトークンに拡張し、さらにGQAやRolling Buffer Cacheと組み合わせることで実用性を高めた。
SWAの本質的な制約は、ウィンドウ外のトークンへの直接参照ができない点である。多層による間接的な情報伝播で補うが、Full Attentionに比べて長距離依存の捕捉精度は劣る場合がある。
対策として以下のハイブリッドアプローチが研究されている:
SWAのRolling Buffer Cacheでは、ウィンドウ幅を超えた過去のKVキャッシュが破棄される。これは推論時のメモリ効率を大幅に改善するが、長い対話セッションでは初期の文脈が完全に失われる。
対策:
SWAはFlashAttention-2/3と組み合わせることで、さらなる高速化を実現する。FlashAttentionはアテンション計算をGPUのSRAM上で効率的にタイリングする技術で、メモリ帯域のボトルネックを解消する。
| 組み合わせ | TFLOPS | メモリ使用量 | 備考 |
|---|---|---|---|
| Full + 標準実装 | ベースライン | O(n²) | 最も遅い |
| Full + FlashAttention-2 | 2-4倍高速 | O(n) | 2024年の標準 |
| SWA + FlashAttention-2 | 4-8倍高速 | O(w) | Mistral方式 |
| SWA + FlashAttention-3 | 6-12倍高速 | O(w) | Hopper GPU最適化 |
vLLM、SGLang、TensorRT-LLMなどの推論フレームワークは、SWA+FlashAttentionの組み合わせをネイティブサポートしている。
Q1: Sliding Window Attentionを使うとモデルの精度は下がりますか? A: 短〜中距離の依存関係が中心のタスク(対話、コード生成、要約など)では、Full Attentionとほぼ同等の精度を維持する。ただし、コンテキスト全体を均一に参照する必要があるタスク(長文の事実検索、Needle in a Haystackテスト)では精度が低下する場合がある。Mistral 7Bのベンチマーク結果ではLLaMA 2 13Bを上回る総合スコアを記録している。
Q2: ウィンドウ幅はどう決めるべきですか? A: ウィンドウ幅は計算コストと精度のトレードオフ。Mistral 7Bでは4,096トークンが採用されているが、用途によって最適値は異なる。コード生成では関数スコープに合わせて2,048〜4,096、長文要約では8,192〜16,384が推奨される。学習後の変更は困難なため、事前に用途を想定した設計が重要。
Q3: SWAとSparse Attentionの違いは何ですか? A: SWAは連続した固定幅ウィンドウ内のトークンのみに注目する「ローカル」パターン。Sparse Attentionは一定間隔でスキップしながら広い範囲を参照する「ストライド」パターンを含む、より一般的な枠組み。SWAはSparse Attentionの一種と見なすことができるが、実装のシンプルさとハードウェア効率の高さから、2024年以降はSWA(またはその変種)が主流。