AI・機械学習
上級

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

Transformerの自己注意機構で、各トークンが全トークンではなく固定幅のウィンドウ内のトークンのみに注意を向ける手法。計算量をO(n²)からO(n×w)に削減し(wはウィンドウ幅)、長系列の効率的な処理を可能にする。Mistralモデルで実用化された。

0 回閲覧
0 いいね

Sliding Window Attention(SWA)とは

Sliding Window Attention(SWA、スライディングウィンドウアテンション)とは、Transformerの自己注意機構において、各トークンが注意を向ける範囲を固定幅のウィンドウに制限する手法である。標準的なSelf-Attentionでは全トークン対全トークンのO(n²)計算が必要だが、SWAではウィンドウ幅wに制限することでO(n×w)に計算量を削減する。

SWAの基本原理

SWAの核心は「言語の局所性」という観測に基づく。自然言語では、あるトークンの意味は近傍のトークンに最も強く依存し、遠方のトークンへの依存度は急速に減衰する。この性質を利用し:

  • 各トークンは前後w/2トークン(合計wトークン)のみに注意を向ける
  • ウィンドウ外のトークンへの注意重みは0に固定(計算自体をスキップ)
  • 層を重ねることで受容野(receptive field)が層数×ウィンドウ幅に拡大

例: ウィンドウ幅4,096、32層のモデルでは理論上の受容野は4,096×32=131,072トークン。これがMistral 7B(SWA w=4,096、32層)が実効的に長い系列を処理できる理由である。

SWA採用モデルの比較

モデルウィンドウ幅層数理論受容野公称コンテキスト長備考
Mistral 7B4,09632131,07232,768SWA実用化の先駆者
Mistral Large 2可変非公開非公開128,000SWA+Full Attentionのハイブリッド
Mixtral 8x7B4,09632131,07232,768MoE+SWAの組み合わせ
Gemma 24,096/8,192交互可変8,192偶数層SWA/奇数層Full
Phi-3-mini2,0483265,536128,000SWA+Longrope併用

SWAの実装パターン

パターン1: 純粋SWA(Mistral 7B方式)

全層でSWAを適用。最もメモリ効率が高いが、長距離依存の捕捉は受容野に依存:

  • KVキャッシュはウィンドウ幅分のみ保持(4,096トークン分)
  • 古いKV値はFIFO方式で破棄(ローリングバッファ)
  • 推論時のメモリ消費がシーケンス長に依存しない

パターン2: ハイブリッドSWA(Gemma 2方式)

SWA層とFull Attention層を交互に配置:

  • 偶数層: SWA(局所パターン捕捉)
  • 奇数層: Full Attention(大域パターン捕捉)
  • KVキャッシュはFull Attention層のみシーケンス全体を保持
  • メモリ消費は純粋SWAより大きいが、長距離依存の精度が向上

パターン3: チャンク化SWA

シーケンスを固定長チャンクに分割し、チャンク内はFull Attention、チャンク間はSWA:

  • チャンクサイズ2,048〜4,096が一般的
  • チャンク境界でのトークン情報損失を防ぐためオーバーラップを設ける
  • Flash Attention 2のブロック処理と相性が良い

Flash Attentionとの統合

SWAはFlash Attention 2/3と組み合わせることで最大の性能を発揮する:

  • Flash Attentionのタイルベース計算でウィンドウ外のタイルをスキップ
  • HBM(High Bandwidth Memory)へのアクセスをウィンドウ幅分に制限
  • Mistral 7Bの場合、Flash Attention + SWAでFull Attentionの約3.5倍の推論スループット
  • vLLMやTensorRT-LLMがSWA対応のFlash Attentionカーネルを提供

SWAの制約と対策

  • 長距離依存の限界: 受容野を超える依存関係は原理的に捕捉不可。対策としてハイブリッド方式や、初期トークン(BOS等)を常時保持するSink Token手法がある
  • 位置エンコーディングとの整合: RoPEとの組み合わせでは、ウィンドウ内の相対位置のみが使用されるため、ウィンドウ幅を超える位置情報は伝播しない
  • チャンク境界問題: ローリングバッファ方式では、バッファ先頭と末尾の連続性が失われる。対策としてチャンクオーバーラップやSink Tokenを使用
  • 並列化効率: SWAのスパース注意パターンはGPUのワープレベル並列性と相性が悪い場合がある。Flash Attentionのブロックサイズをウィンドウ幅の約数に設定することで緩和

よくある質問(FAQ)

Q1: SWAのウィンドウ幅はどう決めるべきか? A: タスクの局所依存性に依存する。コード生成・翻訳など局所パターンが主要なタスクでは2,048〜4,096で十分。長文書要約や多文書比較では8,192以上が推奨される。Mistral 7Bのw=4,096は多くのタスクで良好なバランスを示す。

Q2: SWAを使うとPerplexityは悪化するのか? A: 短〜中程度のテキスト(数千トークン)では影響は軽微(perplexity増加1%未満)。ただし長距離の事実参照や因果推論を含むタスクでは、Full Attentionに比べて精度低下が顕著になる場合がある。

Q3: SWAとSparse AttentionやLinear Attentionの違いは? A: SWAは「固定幅の局所窓」という最もシンプルなスパースパターン。Sparse Attention(BigBird等)は局所窓+ランダム+グローバルトークンを組み合わせた複合パターン。Linear Attentionはカーネル近似によりO(n)を実現するが精度面でSelf-Attentionに劣る。SWAは実装のシンプルさと性能のバランスが優れている。

まとめ

  • SWAは固定幅ウィンドウで注意範囲を制限し、計算量をO(n×w)に削減する手法
  • Mistral 7Bで実用化され、Gemma 2やPhi-3でハイブリッド方式に発展
  • 層の積み重ねにより受容野が拡大し、実効的な長距離依存の捕捉が可能
  • Flash Attention 2/3との統合で最大3.5倍のスループット向上
  • KVキャッシュのローリングバッファにより推論時メモリがシーケンス長に非依存