AI・機械学習
中級

スライディングウィンドウアテンション(スライディングウィンドウアテンション)

Transformerの自己注意機構において、各トークンが固定幅のウィンドウ内にある隣接トークンのみに注意を向けるローカルアテンションの実装手法。Mistral 7Bで広く知られ、FlashAttention-2以降でハードウェアレベルの最適化が進んだ。計算量はO(n×w)でwはウィンドウ幅。

0 回閲覧
0 いいね

スライディングウィンドウアテンション(SWA)とは

スライディングウィンドウアテンション(Sliding Window Attention, SWA)は、Transformerの自己注意機構におけるローカルアテンションの最も代表的な実装である。各トークンが左右w/2個ずつ(合計w個)の隣接トークンのみに注意を向けることで、計算量をFull AttentionのO(n²)からO(n×w)に削減する。2023年にMistral AI社がMistral 7Bで採用し、7Bパラメータクラスでのロングコンテキスト対応のスタンダードとなった。

基本的な仕組み

SWAでは、アテンション行列のうち対角線から±w/2の帯状領域のみを計算し、それ以外の要素はマスク(-∞)で埋める。

ウィンドウ幅wの設計

パラメータ説明典型的な値
ウィンドウ幅 w各トークンが注意する隣接トークン数256 - 8,192
レイヤ数 LTransformerの積み重ね数12 - 80
受容野w × L(全レイヤ通過後の情報到達範囲)4,096 - 655,360
実効コンテキスト受容野の範囲で実際に使える文脈長受容野の50-80%

ウィンドウ幅4,096、レイヤ数32のMistral 7Bの場合:

  • 受容野 = 4,096 × 32 = 131,072トークン
  • 公式コンテキスト長: 32,768トークン(受容野の約25%で安全マージン確保)

多層スタックによる長距離情報伝達

SWAの重要な性質は、単一レイヤではw個のトークンしか見えないが、L層のスタックによりL×w個のトークン情報が間接的に到達する点である。

レイヤ1: トークンiはi-w/2 〜 i+w/2の情報を持つ レイヤ2: トークンiはi-w 〜 i+wの情報を持つ(レイヤ1の隣接トークンの情報が伝播) レイヤL: トークンiはi-Lw/2 〜 i+Lw/2の情報を持つ

主要モデルでの採用状況

モデルSWA設定コンテキスト長発表備考
Mistral 7Bw=4,09632K2023/09SWA普及の立役者
Mistral Nemo 12Bw=4,096128K2024/07NVIDIA共同開発
Mixtral 8x7Bw=4,09632K2023/12MoE + SWA
Mixtral 8x22Bw=4,09665K2024/04大規模MoE + SWA
Phi-3-mini 3.8Bw=2,048128K2024/04Microsoft、小型モデル
Gemma 2 9B/27Bw=4,0968K2024/06交互SWA/Full
Qwen2.5 7Bw=4,096128K2024/09YaRN位置エンコーディング
DeepSeek-V3動的128K2025/01NSAの一部として

FlashAttentionとの統合

FlashAttention-2(Tri Dao, 2023)はSWAをネイティブにサポートしており、ウィンドウ外の計算を完全にスキップするハードウェアレベルの最適化が実現されている。

パフォーマンス比較(A100 80GB GPU)

方式シーケンス長8K16K32K64K128K
Full Attention (naive)12ms48ms192ms768msOOM
FlashAttention-2 Full4ms14ms52ms200ms780ms
FlashAttention-2 SWA(w=4096)3ms6ms12ms24ms48ms
速度比 (SWA/Full)1.3x2.3x4.3x8.3x16.3x

シーケンス長が増えるほどSWAの利点が顕著になる。128Kトークンでは約16倍の高速化が得られる。

KVキャッシュの最適化

SWAの大きなメリットとして、KVキャッシュのサイズをウィンドウ幅wに制限できる点がある。標準的なFull Attentionではシーケンス長に比例してKVキャッシュが増大するが、SWAでは古いトークンのKVをウィンドウから外れた時点で破棄できる。

KVキャッシュサイズ比較(FP16、32ヘッド、128次元/ヘッド)

シーケンス長Full AttentionSWA (w=4096)削減率
8K512MB256MB50%
32K2GB256MB87.5%
128K8GB256MB96.9%
1M64GB256MB99.6%

128Kトークンの場合、KVキャッシュが8GBから256MBに削減される。これはVRAMの節約に直結し、RTX 4090(24GB)でも大規模モデルのロングコンテキスト推論が可能になる。

SWAの限界と対策

SWAの最大の弱点は、ウィンドウ外のトークンへの直接的なアテンションが不可能な点である。多層スタックで間接的に情報は伝播するが、伝播の過程で情報が減衰する「情報ボトルネック」が発生する。

対策手法

  • Sink Token: 先頭トークンに常にアテンションを向ける特別枠を設ける(StreamingLLM)。推論時のパープレキシティ安定化に有効
  • ローカル+グローバル混合: 一部のレイヤでFull Attentionを使用(Gemma 2アプローチ)
  • SSMハイブリッド: SWAレイヤとMamba(SSM)レイヤを交互に配置(Jamba)
  • YaRN位置エンコーディング: 位置エンコーディングのスケーリングでウィンドウ外の情報もゆるやかに保持(Qwen2.5)
  • 動的ウィンドウ: タスクやレイヤに応じてウィンドウ幅を動的に変更(DeepSeek-V3 NSA)

よくある質問(FAQ)

Q1: ウィンドウ幅はどう選べばよいか? A: 一般的にはGPUのL2キャッシュサイズに合わせる。A100(40MB L2)ではw=4096が最適、H100(50MB L2)ではw=8192まで効率的。実用上はw=4096がデファクトスタンダードで、Mistral・Gemma・Qwenなど主要モデルが採用している。

Q2: SWAだけで128Kトークンのコンテキストは本当に活用できるか? A: 完全には活用できない。Mistral 7Bのw=4096×32層=受容野131Kでカバーはしているが、ウィンドウ外の情報は間接伝達のため減衰する。Needle-in-a-Haystack評価では、128K後半の検索精度がFull Attentionに比べ5-15%低下する報告がある。

Q3: 自作PCでSWAモデルを高速に動かすコツは? A: llama.cppの最新版はSWAのKVキャッシュ最適化に対応しており、--ctx-size 32768と--flash-attnフラグの併用が最も効果的。RTX 4070 Ti Super(16GB VRAM、約13万円)でMistral 7B Q4_K_Mが32Kコンテキストで約40 token/sの生成速度を実現する。

まとめ

  • SWAは各トークンが固定幅ウィンドウ内のみに注意を向けるシンプルかつ高速な手法
  • Mistral 7Bの採用で広く普及し、w=4096がデファクトスタンダード
  • FlashAttention-2との統合で128Kトークンでも16倍以上の高速化
  • KVキャッシュをウィンドウ幅に制限でき、VRAM使用量を大幅削減
  • ウィンドウ外の情報減衰が弱点だが、Sink Token・SSMハイブリッド等で対策可能