スライディングウィンドウアテンション(スライディングウィンドウアテンション)
スライディングウィンドウアテンション(Sliding Window Attention, SWA)とは、各トークンが全トークンではなく周囲の固定幅ウィンドウ(通常1,024〜4,096トークン)内のトークンのみに注意を向けるローカルアテンション手法である。計算量をO(n^2)からO(n * w)に削減しつつ、複数層を積み重ねることで実効的な受容野を拡大する。Mistral 7Bがウィンドウ幅4,096で採用し、長文処理の効率化手法として広く知られるようになった。
スライディングウィンドウアテンション(SWA)とは
スライディングウィンドウアテンション(Sliding Window Attention, SWA)は、各トークンのアテンション範囲を固定幅のウィンドウに制限するローカルアテンション手法である。全トークン間のアテンション(フルアテンション)がO(n^2)の計算量を要するのに対し、SWAはウィンドウ幅wに対してO(n * w)で計算でき、長いシーケンスでも線形に近いスケーリングを実現する。
概要と動作原理
SWAの動作は直感的である:
- 各トークンは自身の前後wトークン(因果マスクの場合は自身以前のwトークン)のみに注意を向ける
- アテンション行列はn x nの密行列ではなく、帯行列(band matrix)となる
- ウィンドウ外のトークンとの直接的なアテンションはゼロ
- 複数層を重ねることで、L層のモデルの実効受容野(effective receptive field)はL x wトークンに拡大
Mistral 7Bの例: ウィンドウ幅w=4,096、32層の構成で、実効受容野は最大32 x 4,096 = 131,072トークンに達する。
フルアテンションとの比較
| 特性 | フルアテンション | SWA (w=4096) | SWA + Sink Token |
|---|---|---|---|
| 計算量 | O(n^2 * d) | O(n * w * d) | O(n * (w+s) * d) |
| 32Kトークン時メモリ | ~4GB | ~512MB | ~530MB |
| 128Kトークン時メモリ | ~64GB | ~2GB | ~2.1GB |
| 長距離依存性 | 直接参照可能 | L層で間接的に拡大 | Sink経由で補完 |
| 適用例 | GPT-4o, Llama 3.1 | Mistral 7B, Longformer | StreamingLLM |
主要なSWA実装モデル
| モデル | パラメータ | ウィンドウ幅 | 層数 | 実効受容野 | 特記事項 |
|---|---|---|---|---|---|
| Longformer | 149M | 512 | 12 | 6,144 | グローバルアテンション併用 |
| BigBird | 125M-4.3B | 512 | 12-48 | 最大24K | ランダムアテンション併用 |
| Mistral 7B | 7.3B | 4,096 | 32 | 131,072 | GQA + SWA の組合せ |
| Mixtral 8x7B | 46.7B MoE | 4,096 | 32 | 131,072 | MoE + SWA |
| Gemma 2 | 9B/27B | 4,096/8,192 | 26/46 | 可変 | 交互にグローバル/ローカル |
| Phi-3 Mini | 3.8B | 2,048 | 32 | 65,536 | LongRoPE併用で128K対応 |
SWAの拡張手法
純粋なSWAの限界(ウィンドウ外の情報へのアクセス不可)を解消するため、複数の拡張手法が提案されている:
- グローバルアテンション(Longformer): 特定のトークン([CLS]、質問トークン等)はフルアテンションで全トークンを参照可能にする
- ランダムアテンション(BigBird): ウィンドウ外のランダムな位置にもアテンション接続を追加(情報の長距離伝播を促進)
- Attention Sink(StreamingLLM): 先頭の数トークン(sink token)を常にアテンション範囲に含めることで、無限長シーケンスへの対応を実現
- 交互構成(Gemma 2): 偶数層でグローバルアテンション、奇数層でSWAを交互に配置し、効率と品質を両立
- Dilated SWA: ウィンドウ内の位置を間引き(dilation)して、同じ計算量でより広い範囲をカバー
FlashAttentionとの組み合わせ
FlashAttention-2はSWAのマスクパターンに対応しており、SWA + FlashAttentionの組み合わせで二重の高速化が得られる:
- ウィンドウ外のブロックをスキップすることでFlashAttentionのタイル計算量を削減
- Mistral 7BではFlashAttention-2 + SWAにより、128Kトークンの処理が実用的な速度で可能
- vLLMのSWA対応: PagedAttention + FlashAttention + SWAの三重最適化
よくある質問(FAQ)
Q1: ウィンドウ幅はどう決めればよいですか? A: 一般的にはモデルの層数Lとの積(L x w)が想定される最大シーケンス長以上になるように設定する。Mistral 7Bの例ではw=4,096、L=32で実効受容野131,072トークン。タスクに依存するが、w=2,048〜8,192の範囲が実用的。
Q2: SWAはフルアテンションより品質が劣りますか? A: 単一層では情報範囲が制限されるが、多層積み重ねにより実効受容野が拡大するため、実際のベンチマーク(MMLU、HumanEval等)ではフルアテンションモデルと同等かそれ以上の性能を達成している。Mistral 7Bは同サイズのLlama 2 13Bを上回った。
Q3: SWAとRoPE(Rotary Position Embedding)の関係は? A: RoPEは位置情報のエンコーディング手法であり、SWAとは直交する技術。SWAモデルでもRoPEは標準的に使用される。ただしRoPEの外挿性能(学習時より長いシーケンスへの汎化)がSWAの実効受容野の実効性に影響するため、LongRoPE等の拡張手法との組み合わせが重要。
まとめ
- SWAは固定幅ウィンドウでアテンション範囲を制限し、計算量をO(n * w)に削減する
- 複数層の積み重ねで実効受容野がL x wトークンに拡大(Mistral 7B: 131Kトークン)
- グローバルアテンション、Attention Sink、交互構成などの拡張手法で長距離依存性を補完
- FlashAttention-2との組み合わせで二重の高速化が可能