Sliding Window Attention（SWA）の理論と実装 — 固定窓による局所注意とMistralでの採用

Transformerベースの大規模言語モデル（LLM）を使っていると、ある壁にぶつかります。入力トークン数を増やすほどGPUメモリが急激に膨らみ、ある長さを超えると推論すらできなくなるのです。この原因はSelf-Attentionの $O(n^2)$ という計算量にあります。しかし、本当に全てのトークンが遠く離れた全てのトークンに注目する必要があるのでしょうか。

日本語の文章を読むとき、あるフレーズの意味を理解するために10ページ前の単語を均等に参照することはまずありません。多くの場合、近くの数十〜数百トークンの文脈が最も重要であり、遠方のトークンとの関係は間接的に推論されます。この「注意の局所性」を正面から活用するのがSliding Window Attention（SWA、スライディングウィンドウアテンション）です。

2023年にMistral AIが公開したMistral 7Bは、SWAを中核技術として採用し、7Bのパラメータで13Bクラスのモデルを上回る性能を実証しました。SWAにより推論時のKVキャッシュサイズが系列長によらず一定に抑えられるため、エッジデバイスや長文脈でのデプロイが格段に容易になります。

仕組みは図のとおり単純です。各トークン（赤）は、自分を含む直前 $w$ 個のトークン（青）だけに注目し、窓の外（グレー）はまったく見ません。この「視界の制限」だけで計算量とメモリが劇的に変わる——本記事ではその理屈と実装を丁寧に追いかけます。

Sliding Window Attentionを理解すると、以下のような応用・知見が得られます。

• LLM推論のメモリ最適化: 窓サイズで決まる固定メモリ量でKVキャッシュを運用でき、長系列推論のメモリ見積もりが正確に行える
• Mistral/Mixtralの深い理解: オープンソースLLMの主流アーキテクチャであるMistralの設計思想を把握でき、vLLMやllama.cppでの運用最適化に直結する
• 効率的Attentionの体系的理解: Longformerのローカルアテンション、Flash AttentionのブロックスパースなどのSparse Attention手法との関係が整理され、用途に応じた手法選択が可能になる
• 受容野設計の理論: CNNの受容野と同様の考え方がTransformerにも適用できることを知り、モデルの表現力を理論的に分析できるようになる

本記事の内容

• Full Attentionの $O(n^2)$ 問題と注意の局所性仮説
• Sliding Window Attentionの数学的定義
• 窓サイズ $w$ と有効受容野の関係
• 多層スタックによる受容野の拡大メカニズム
• MistralにおけるSWAとGrouped-Query Attention（GQA）の組み合わせ
• Rolling Buffer KV Cacheの仕組み
• Longformer方式との比較
• PyTorchでのスクラッチ実装と計算時間・メモリ比較実験

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

Full Attentionの問題点

なぜ $O(n^2)$ が致命的なのか

Sliding Window Attentionの動機を理解するために、まずFull Attentionの何が問題なのかを改めて確認しましょう。

Scaled Dot-Product Attentionの計算式は次のとおりです。

$$ \text{Attention}(\bm{Q}, \bm{K}, \bm{V}) = \text{softmax}\left(\frac{\bm{Q}\bm{K}^\top}{\sqrt{d_k}}\right)\bm{V} $$

ここで $\bm{Q}, \bm{K} \in \mathbb{R}^{n \times d_k}$、$\bm{V} \in \mathbb{R}^{n \times d_v}$ です。$\bm{Q}\bm{K}^\top$ の結果は $n \times n$ の行列になるため、系列長 $n$ の二乗に比例する計算量 $O(n^2 d_k)$ とメモリ $O(n^2)$ が必要です。

たとえば系列長512なら注意行列のメモリは約1 MBですが、系列長32,768（Mistral 7Bの最大系列長）では約4 GBに達します。これは1ヘッド・1レイヤー分の話であり、32ヘッド×32レイヤーの全体では到底GPUメモリに収まりません。

LLM推論時のKVキャッシュ問題

訓練時の $O(n^2)$ 問題に加えて、推論時にはKVキャッシュの線形増大という別の問題が生じます。自己回帰生成では、新しいトークンを1つ生成するたびに、過去の全トークンのKey・Valueベクトルが必要です。Full Attentionでは生成が進むほどキャッシュが大きくなり、128,000トークンのコンテキストを持つモデルでは数十GBのKVキャッシュが必要になります。

バッチ推論を行う場合、この問題はさらに深刻です。バッチサイズ $B$ のリクエストを同時に処理するには、KVキャッシュも $B$ 倍必要です。GPUメモリの多くがKVキャッシュに奪われ、バッチサイズを上げられない — これがLLMサービングのボトルネックとしてよく知られています。

注意行列の実質的なスパース性

Full Attentionの注意行列を学習済みモデルで可視化すると、興味深い事実が浮かび上がります。ほとんどのAttentionヘッドで、注意重みが集中しているのは対角線付近（近傍トークン）と文頭の特殊トークン（BOS、句読点など）です。遠く離れたトークン間の注意重みは極めて小さく、softmaxの出力としてはほぼゼロです。

これは自然言語の統計的性質と合致しています。英語でも日本語でも、文法的な係り受けのほとんどは数単語〜数十単語の範囲に収まります。長距離の依存関係（照応解析、文書レベルの一貫性など）は存在しますが、それらは少数のトークンペアに限定されます。

つまり、$n \times n$ の注意行列のうち、本当に意味のある値を持つのはごく一部です。ならば、最初から「近くのトークンにだけ注目する」というマスクを掛けてしまえば、計算量を大幅に削減できるはずです。これがSliding Window Attentionの基本発想です。

ここまでで、Full Attentionが長系列で致命的なメモリ・計算コストを持つこと、そして注意行列が実質的に局所的な構造を持つことがわかりました。次に、この局所性を直接的に活用するSliding Window Attentionの数学的な定義を見ていきましょう。

Sliding Window Attentionの数学的定義

直感的な理解 — 「窓」を滑らせる

Sliding Window Attentionのアイデアは、1次元の畳み込み（CNN）のカーネルと同じ発想です。1D CNNでは、各位置の出力はその前後の固定幅 $k$ の入力だけから計算されます。同様に、SWAでは各トークンのQuery $\bm{q}_i$ が注目するKeyの範囲を、自分の位置 $i$ を中心とした固定幅の「窓」（ウィンドウ）に制限します。

窓の外にあるトークンは、まるで存在しないかのように扱われます。教室で先生の話を聞いているとき、自分の席の前後数列の友人の表情は見えますが、100メートル離れた別の教室の生徒は視界に入りません。SWAはこの「視界の範囲を固定する」という単純なアイデアを数式にしたものです。

形式的な定義

系列長 $n$、窓サイズ $w$ のSliding Window Attentionを正式に定義します。位置 $i$ のトークンが注目できるKey位置の集合 $\mathcal{W}(i)$ は次のように定義されます。

$$ \begin{equation} \mathcal{W}(i) = \{j \mid \max(0, i – w + 1) \leq j \leq i\} \end{equation} $$

ここで、因果的（causal）なマスクを考慮しているため、$j > i$ のトークン（未来のトークン）は含みません。つまり、各トークンは自分自身を含む直前 $w$ 個のトークンにのみ注目します。

この集合を使って、注意マスク $\bm{M} \in \{0, 1\}^{n \times n}$ を定義します。

$$ \begin{equation} M_{ij} = \begin{cases} 1 & \text{if } j \in \mathcal{W}(i) \\ 0 & \text{otherwise} \end{cases} \end{equation} $$

マスクされた注意スコアは、$M_{ij} = 0$ の位置を $-\infty$ に設定することで実現します。

$$ \begin{equation} \tilde{s}_{ij} = \begin{cases} \dfrac{\bm{q}_i^\top \bm{k}_j}{\sqrt{d_k}} & \text{if } M_{ij} = 1 \\[8pt] -\infty & \text{if } M_{ij} = 0 \end{cases} \end{equation} $$

softmaxを適用すると、$-\infty$ の位置は重みゼロになります。

$$ \begin{equation} \alpha_{ij} = \frac{\exp(\tilde{s}_{ij})}{\sum_{k=1}^{n} \exp(\tilde{s}_{ik})} \end{equation} $$

最終的な出力は、窓内のValueベクトルの重み付き和になります。

$$ \begin{equation} \bm{o}_i = \sum_{j \in \mathcal{W}(i)} \alpha_{ij} \bm{v}_j \end{equation} $$

計算量の分析

Full Attentionでは各トークンが $n$ 個のKeyに注目するため計算量は $O(n^2 d_k)$ でした。SWAでは各トークンが最大 $w$ 個のKeyにしか注目しないため、計算量は次のようになります。

$$ O(n \cdot w \cdot d_k) $$

$w$ がモデル設計時に固定される定数（Mistral 7Bでは $w = 4096$）であるため、系列長 $n$ に対して線形の計算量です。同様に、注意行列のメモリも $O(nw)$ — Full Attentionの $O(n^2)$ から劇的に削減されます。

メモリ使用量の具体的な比較

窓サイズ $w = 4096$（Mistral 7Bの設定）のとき、注意行列のメモリ使用量（FP32、1ヘッド分）を比較してみましょう。

系列長が窓サイズ $w$ を超えると、SWAのメモリ使用量は $n$ に線形にしか増えないため、削減効果が急激に大きくなることがわかります。系列長 $n = w$ のときはFull Attentionと同じ（全トークンが窓内に収まる）ですが、系列長が窓サイズの32倍になると97%近いメモリ削減を達成しています。

この表を対数スケールの棒グラフにすると、削減の構造がよく見えます。系列長4Kでは両者は同じ高さですが、系列長が伸びるほどFull Attention（赤）だけが二次関数的に伸び、SWA（青）との差が開いていきます。131Kトークンでは約97%の削減——「窓の外を捨てる」だけでこれだけの差になります。

ここまでで、SWAの数学的な定義と計算量の分析を行いました。しかし、「窓の外のトークンの情報は本当に失われないのか？」という疑問が自然に浮かびます。次に、多層のTransformerをスタックしたときに窓の外の情報がどのように伝播するかを見ていきましょう。

窓サイズと有効受容野

1層での受容野

CNNにおいて、複数の畳み込み層を重ねるとカーネルサイズ以上の広い範囲の情報を間接的に利用できるようになります。これと全く同じメカニズムがSliding Window Attentionにも成り立ちます。

まず、1層のSWA（窓サイズ $w$）における受容野を考えましょう。位置 $i$ のトークンは $\mathcal{W}(i) = \{i – w + 1, \ldots, i\}$ のトークンにのみ直接注目できます。したがって、1層での受容野は $w$ トークンです。

多層スタックによる受容野の拡大

2層のSWAを考えます。1層目で位置 $i$ のトークンは位置 $\{i-w+1, \ldots, i\}$ の情報を集約します。2層目では、この出力が再びSWAに入力されます。2層目の位置 $i$ は、1層目の出力 $\{i-w+1, \ldots, i\}$ に注目しますが、1層目の位置 $i-w+1$ は元の入力の位置 $\{i-2w+2, \ldots, i-w+1\}$ の情報を含んでいます。

したがって、2層を経由すると位置 $i$ は元の入力の位置 $\{i-2w+2, \ldots, i\}$ の情報にアクセスできます。これは約 $2w$ トークンの受容野に相当します。

この伝播の様子を図にしました。2層目の赤いトークンは1層目の窓内3トークン（薄赤）しか直接見ていませんが、その各トークンが入力の窓内を見ているため、間接的に約 $2w$ の範囲（青帯）の情報が届きます。層を1つ重ねるごとに、届く範囲が約 $w$ ずつ左に伸びていくわけです。

一般化すると、$L$ 層のSWAをスタックしたときの理論上の受容野は次のようになります。

$$ \begin{equation} R(L) = L \times (w – 1) + 1 \approx L \times w \end{equation} $$

$w – 1$ なのは、各層で自分自身の位置が重複するためです。$w$ が十分大きい場合は $R(L) \approx Lw$ と近似できます。

Mistral 7Bでの具体的な受容野

Mistral 7Bのパラメータ $w = 4096$、$L = 32$（レイヤー数）を代入してみましょう。

$$ R(32) = 32 \times (4096 – 1) + 1 = 131,009 \approx 131,\!072 $$

理論上の受容野は約131,000トークンに達します。つまり、窓サイズが4,096でも、32層をスタックすることで12万トークン以上の入力を間接的にカバーできるのです。これはMistral 7Bの最大コンテキスト長（32,768トークン）を大幅に上回ります。

理論的受容野と有効受容野の違い

ただし、ここで重要な注意点があります。理論上の受容野 $Lw$ は「情報が到達し得る最大範囲」を示すものであり、実際にモデルが情報を有効に利用できる範囲（有効受容野、effective receptive field）はこれより小さくなります。

CNNの有効受容野の研究（Luo et al., 2016）で知られているように、理論上の受容野の端に近い位置の影響は指数的に減衰します。SWAでも同様に、$L$ 層を経由して伝播する情報は各層のsoftmaxで重み付けされ、希釈されていきます。

直感的に言えば、理論上の受容野は「光が届く最大距離」、有効受容野は「実際に照らされて見える範囲」に相当します。32層スタックで理論上は13万トークンに届きますが、10層程度を超えた間接参照の情報はかなり薄まっていると考えるべきです。

この「有効受容野の減衰」はSWAの本質的な限界であり、完全なFull Attentionに対するトレードオフです。しかし、実用上は多くのタスクで局所的な注意が支配的であるため、このトレードオフは十分に合理的です。

ここまでで、SWAの理論的な受容野が $Lw$ に達すること、ただし有効受容野はそれより小さいことがわかりました。では、Mistral 7BではこのSWAをどのような設計哲学で採用しているのでしょうか。次に、Mistralのアーキテクチャ全体の中でのSWAの位置づけと、GQAとの組み合わせについて見ていきます。

MistralにおけるSWAの採用

Mistral 7Bのアーキテクチャ概要

Mistral 7BはMeta社のLLaMA 2をベースに、いくつかの重要な変更を加えたモデルです。SWAに関連する主要な設計パラメータを整理します。

最も注目すべきは、SWAとGQA（Grouped-Query Attention）を同時に採用している点です。それぞれの手法が異なる側面からメモリ効率を改善し、組み合わせることで相乗効果を生んでいます。

SWAとGQAの相乗効果

SWAは「各トークンが注目する範囲を $w$ に制限する」ことで、注意行列のサイズを $O(n^2)$ から $O(nw)$ に削減します。これは系列長方向のメモリ削減です。

GQAは「複数のQueryヘッドがKV（Key-Value）ヘッドを共有する」ことで、KVキャッシュのメモリを削減します。Mistral 7Bでは32個のQueryヘッドに対して8個のKVヘッドを共有するため、KVキャッシュのメモリは $\frac{8}{32} = \frac{1}{4}$ に削減されます。これはヘッド数方向のメモリ削減です。

両者を組み合わせたときの1レイヤーあたりのKVキャッシュサイズを計算してみましょう。Full MHA（Multi-Head Attention）+ Full Attentionの場合と比較します。

Full MHA + Full Attention:

$$ \text{KV Cache} = 2 \times n \times h \times d_k \times \text{sizeof(dtype)} $$

ここで $h$ はKVヘッド数、$d_k$ はヘッド次元です。$n = 32768$、$h = 32$（全ヘッド分）、$d_k = 128$、FP16（2バイト）として計算すると、

$$ 2 \times 32768 \times 32 \times 128 \times 2 = 536,\!870,\!912 \text{ bytes} \approx 512 \text{ MB} $$

GQA + SWA（Mistral 7B）:

SWAにより保持するトークン数は最大 $w = 4096$ に制限され、GQAによりKVヘッド数は8に削減されます。

$$ 2 \times 4096 \times 8 \times 128 \times 2 = 16,\!777,\!216 \text{ bytes} = 16 \text{ MB} $$

1レイヤーあたりのKVキャッシュが512 MBから16 MBへ、32分の1に削減されます。32レイヤー全体でも $16 \times 32 = 512$ MB — これなら推論用のGPUメモリに十分収まります。

2つの手法の関係を面積で表すと上図のようになります。SWAは横（系列長）方向を $1/8$ に、GQAは縦（ヘッド数）方向を $1/4$ に削る、互いに直交する削減です。だから効果は単純に掛け算になり、$\frac{1}{8} \times \frac{1}{4} = \frac{1}{32}$ が実現します。

なぜMistralはSWAを選んだのか

効率的なAttentionの手法は多数ありますが、MistralがSWAを選んだ理由はその実装の単純さと推論時のメモリ上限保証にあります。

1. 実装が極めて単純: SWAはAttentionマスクを変更するだけで実装できます。Flash Attentionのようなカスタムカーネルの修正も最小限で済みます。Longformerのようにグローバルトークンの特別扱いも不要です
2. メモリ上限が保証される: KVキャッシュのサイズが窓サイズ $w$ で上限が決まるため、系列長がいくら長くなってもメモリが際限なく増えることがありません。これはプロダクション環境でのメモリ計画において極めて重要です
3. Flash Attentionとの相性が良い: Flash Attention 2はブロック単位で注意行列を計算するため、SWAのような帯行列構造と自然に組み合わせられます。ブロックが窓の外であれば計算自体をスキップできるため、実速度の改善に直結します

これらの利点が、Mistralの「小さくて速いが性能は高い」という設計哲学と完全に一致したのです。

ここまでで、MistralがSWAとGQAを組み合わせて驚異的なメモリ効率を実現していることがわかりました。しかし、SWAでKVキャッシュのサイズを制限するだけでは、推論時にキャッシュをどう管理するかという実装上の問題が残ります。次に、Mistralが採用しているRolling Buffer KV Cacheの仕組みを見ていきましょう。

Rolling Buffer KV Cache

なぜ特殊なキャッシュ管理が必要か

通常のKVキャッシュでは、生成が進むにつれてKey・Valueベクトルがリストの末尾に追加されていきます。SWAを使うとき、窓サイズ $w$ 個より前のKVはもう参照されないため、削除してよいはずです。しかし、単純にリストの先頭から削除すると、メモリの再配置（コンパクション）が頻繁に発生し、GPUの性能を著しく損なう場合があります。

リングバッファによる解決

Mistralが採用するRolling Buffer KV Cacheは、固定サイズ $w$ のリングバッファ（循環バッファ）でKVキャッシュを管理します。新しいトークンのKVは、位置のモジュロ演算で上書き位置を決定します。

位置 $t$ のトークンが格納されるバッファのインデックスは次のとおりです。

$$ \begin{equation} \text{idx}(t) = t \bmod w \end{equation} $$

たとえば窓サイズ $w = 4$ の場合、各ステップでのバッファの状態を追跡してみましょう。

ステップ4以降は、最も古いKVが新しいKVで上書きされていきます。バッファのサイズは常に $w$ で一定であり、メモリの動的な確保や解放が一切発生しません。

上の表を図にしたものです。赤いスロットが各ステップで書き込まれた位置で、$t \bmod 4$ の巡回により左から右へ順に進み、一周すると先頭に戻って最古のKVを上書きします。どのステップでもスロット数は4のまま——「メモリが一定」という性質が一目でわかります。

Rolling Buffer KV Cacheのメリット

このリングバッファ方式には3つの大きなメリットがあります。

1. メモリサイズが定数: KVキャッシュのメモリは $w \times d$ で固定され、生成トークン数にかかわらず一定です。100トークン生成しても10万トークン生成しても、同じメモリで済みます。

2. メモリコピーが不要: 古いエントリを削除してメモリを詰める操作が不要です。新しいKVは計算した瞬間にバッファの所定位置に直接書き込まれます。

3. 連続メモリアクセス: GPUの性能はメモリアクセスパターンに大きく依存します。固定サイズの連続バッファにアクセスするため、キャッシュミスが少なくGPUの演算効率が高い状態を維持できます。

位置エンコーディングとの整合性

Rolling Bufferでは、バッファ内のKVの物理的な並び順が時間順と一致しません（上の表でステップ5の時点では $[\bm{k}_4, \bm{k}_5, \bm{k}_2, \bm{k}_3]$ — 位置2と3のKVはまだ古いものが残っている）。

これが問題にならない理由は、MistralがRotary Position Embedding（RoPE）を使用しているためです。RoPEは位置情報をQuery・Keyに直接エンコードするため、注意スコア $\bm{q}_i^\top \bm{k}_j$ を計算する際に物理的な格納位置ではなく元の位置 $i, j$ が使われます。つまり、バッファ内のKVがどの物理スロットに格納されていても、正しい相対位置の情報が注意スコアに反映されるのです。

ここまでで、Rolling Buffer KV Cacheの仕組みと、それがSWAの推論効率をさらに高める仕組みを理解しました。次に、SWAと同じく局所注意を活用するLongformerとの違いを明確にし、それぞれの手法の使い分けを考えましょう。

Longformer方式との比較

共通点 — 局所注意という基盤

Sliding Window AttentionとLongformer（Beltagy et al., 2020）は、どちらも「局所的な注意パターン」を基盤とする点で共通しています。Longformerのローカルアテンション（Local Attention）は、各トークンが前後 $w/2$ 個のトークンに注目する仕組みであり、SWAと本質的に同じ操作です。

相違点1 — グローバルトークンの有無

Longformerは局所注意に加えて、グローバルアテンションを持ちます。[CLS]トークンなどの特定のトークンが全トークンに注目し、逆に全トークンもそのグローバルトークンに注目します。これにより、文書全体の情報がグローバルトークンを経由して集約されます。

一方、SWA（Mistralの実装）にはグローバルトークンがありません。文書全体の情報は、多層スタックによる受容野の拡大に頼ります。グローバルトークンを使わないことで、実装がシンプルになり、Flash Attentionとの互換性も高くなります。

相違点2 — 双方向 vs 因果的

Longformerは主にエンコーダモデル（BERT系）の効率化のために設計されており、双方向の局所注意を使います。各トークンは前後の $w/2$ トークンに注目します。

SWAはデコーダモデル（GPT系、LLM）のために設計されており、因果的（causal）な局所注意を使います。各トークンは自分自身と直前の $w – 1$ トークンにのみ注目し、未来のトークンは見えません。

相違点3 — 設計上の対象タスク

どちらを選ぶべきか

分類や質問応答など、入力全体をエンコードするタスクでは、Longformerのグローバルトークンによる情報集約が有利です。一方、テキスト生成（LLM推論）では、因果マスク + Rolling Buffer KV CacheによるSWAが自然な選択です。

実際の使い分けはモデルアーキテクチャに依存します。BERTベースのモデルを長系列対応にするならLongformer、LLMを効率的に動かすならSWA（Mistral方式）が第一選択です。

ここまでで、SWAの理論的な側面を一通り整理しました。次に、Pythonで実際にSWAを実装して、Full Attentionとの計算量・メモリの違いを実験で確認しましょう。

Pythonでの実装

SWAマスクの生成

まず、Sliding Window Attentionの因果的マスクを生成する関数を実装します。このマスクがSWAの核心部分であり、Full Attentionとの唯一の違いです。

import numpy as np
import matplotlib, matplotlib.pyplot as plt
import torch
import torch.nn.functional as F
import time

# 図の日本語ラベル用フォント設定
for cand in ["Hiragino Sans", "Yu Gothic", "Noto Sans CJK JP", "IPAexGothic", "Meiryo"]:
    if any(cand == f.name for f in matplotlib.font_manager.fontManager.ttflist):
        plt.rcParams["font.family"] = cand
        break
plt.rcParams["axes.unicode_minus"] = False

def create_swa_mask(seq_len: int, window_size: int) -> torch.Tensor:
    """因果的 Sliding Window Attention マスクを生成する

    Args:
        seq_len: 系列長 n
        window_size: 窓サイズ w

    Returns:
        mask: (seq_len, seq_len) のブールマスク。True = 注目可能
    """
    # 因果マスク（下三角行列）
    causal_mask = torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool))

    # 窓マスク（対角線から w-1 行以内）
    row_idx = torch.arange(seq_len).unsqueeze(1)  # (n, 1)
    col_idx = torch.arange(seq_len).unsqueeze(0)  # (1, n)
    window_mask = (row_idx - col_idx) < window_size  # 距離が w 未満

    # 因果マスクと窓マスクの論理積
    swa_mask = causal_mask & window_mask
    return swa_mask

この関数は2つのマスクの論理積（AND）を取っています。causal_mask は「未来のトークンを見ない」という制約、window_mask は「距離が $w$ 以上離れたトークンを見ない」という制約です。両方を満たす位置だけが True になります。

マスクパターンの可視化

3つのAttentionパターン（Full Causal、SWA、Longformer風）を並べて可視化してみましょう。

n = 32  # 可視化用の小さな系列長
w = 8   # 窓サイズ

# Full Causal Attention マスク
full_mask = torch.tril(torch.ones(n, n, dtype=torch.bool))

# Sliding Window Attention マスク
swa_mask = create_swa_mask(n, w)

# Longformer風（双方向ローカル + グローバルトークン2個）
row_idx = torch.arange(n).unsqueeze(1)
col_idx = torch.arange(n).unsqueeze(0)
local_mask = (row_idx - col_idx).abs() < (w // 2)
global_positions = [0, 1]  # 先頭2トークンをグローバルに
longformer_mask = local_mask.clone()
for gp in global_positions:
    longformer_mask[gp, :] = True   # グローバルトークン → 全トークンに注目
    longformer_mask[:, gp] = True   # 全トークン → グローバルトークンに注目

fig, axes = plt.subplots(1, 3, figsize=(18, 6))

masks = [full_mask, swa_mask, longformer_mask]
titles = [
    f'Full Causal Attention\n(非ゼロ {full_mask.sum().item()} 個)',
    f'Sliding Window Attention (w={w})\n(非ゼロ {swa_mask.sum().item()} 個)',
    f'Longformer風 (w={w})\n(非ゼロ {longformer_mask.sum().item()} 個)'
]

for ax, mask, title in zip(axes, masks, titles):
    ax.imshow(mask.numpy(), cmap='Blues', aspect='equal', vmin=0, vmax=1)
    ax.set_title(title, fontsize=13)
    ax.set_xlabel('Key の位置')
    ax.set_ylabel('Query の位置')

plt.tight_layout()
plt.savefig('attention_mask_comparison.png', dpi=150, bbox_inches='tight')
plt.show()

左のFull Causal Attentionは下三角行列で、各トークンが自分以前の全トークンに注目しています。非ゼロ要素数は $n(n+1)/2 = 528$ 個です。中央のSWAは対角線付近の帯状構造になっており、非ゼロ要素数が大幅に減少しています。窓サイズ $w = 8$ で系列長 $n = 32$ の場合、最大でも $nw = 256$ 個程度の注意計算しか行いません。右のLongformer風マスクは、局所注意（双方向）に加えて先頭2トークンの行・列が全て接続されており、グローバルトークンのハブとしての役割が視覚的にわかります。

Scaled Dot-Product Attentionの実装

SWAマスクを使ったAttention計算を実装します。

def scaled_dot_product_attention(
    Q: torch.Tensor,
    K: torch.Tensor,
    V: torch.Tensor,
    mask: torch.Tensor = None
) -> torch.Tensor:
    """Scaled Dot-Product Attention（マスク対応）

    Args:
        Q: Query (batch, heads, seq_len, d_k)
        K: Key   (batch, heads, seq_len, d_k)
        V: Value (batch, heads, seq_len, d_v)
        mask: ブールマスク (seq_len, seq_len)。Trueの位置に注目

    Returns:
        output: (batch, heads, seq_len, d_v)
    """
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)

    if mask is not None:
        # Falseの位置を -inf に設定
        scores = scores.masked_fill(~mask.unsqueeze(0).unsqueeze(0), float('-inf'))

    attn_weights = F.softmax(scores, dim=-1)
    # NaN防止（全て -inf の行がある場合）
    attn_weights = attn_weights.nan_to_num(0.0)

    output = torch.matmul(attn_weights, V)
    return output

この実装のポイントは masked_fill でマスクの False 位置を $-\infty$ に設定する部分です。softmaxに通すと $e^{-\infty} = 0$ となるため、マスクされた位置の注意重みが自動的にゼロになります。系列の先頭付近で窓内のトークン数が少ない場合も、softmaxの正規化が正しく機能します。

Rolling Buffer KV Cacheの実装

推論時のRolling Buffer KV Cacheも実装してみましょう。

class RollingBufferKVCache:
    """Rolling Buffer KV Cache（Mistral方式）"""

    def __init__(self, window_size: int, num_heads: int, head_dim: int,
                 dtype=torch.float32):
        self.window_size = window_size
        self.num_heads = num_heads
        self.head_dim = head_dim

        # 固定サイズのリングバッファ
        self.k_cache = torch.zeros(1, num_heads, window_size, head_dim, dtype=dtype)
        self.v_cache = torch.zeros(1, num_heads, window_size, head_dim, dtype=dtype)
        self.current_len = 0  # これまでに格納したトークン数

    def update(self, new_k: torch.Tensor, new_v: torch.Tensor):
        """新しいトークンのKVを追加（上書き）

        Args:
            new_k: (1, num_heads, 1, head_dim)
            new_v: (1, num_heads, 1, head_dim)
        """
        # モジュロ演算で書き込み位置を決定
        idx = self.current_len % self.window_size
        self.k_cache[:, :, idx:idx+1, :] = new_k
        self.v_cache[:, :, idx:idx+1, :] = new_v
        self.current_len += 1

    def get_kv(self):
        """現在のKVキャッシュを取得

        Returns:
            k: (1, num_heads, valid_len, head_dim)
            v: (1, num_heads, valid_len, head_dim)
        """
        valid_len = min(self.current_len, self.window_size)
        if self.current_len <= self.window_size:
            return self.k_cache[:, :, :valid_len, :], self.v_cache[:, :, :valid_len, :]
        else:
            # リングバッファの内容を時間順に並べ替え
            start_idx = self.current_len % self.window_size
            indices = [(start_idx + i) % self.window_size
                       for i in range(self.window_size)]
            k = self.k_cache[:, :, indices, :]
            v = self.v_cache[:, :, indices, :]
            return k, v

    def memory_bytes(self):
        """KVキャッシュが使用するメモリ量（バイト）"""
        element_size = self.k_cache.element_size()
        return 2 * self.window_size * self.num_heads * self.head_dim * element_size

update メソッドでは、self.current_len % self.window_size によって書き込みインデックスを循環させています。バッファが一周した後も、古いKVは新しいKVで自動的に上書きされ、メモリサイズは一定です。get_kv メソッドでは、バッファの内容を時間順に並べ替えて返します。

Rolling Bufferの動作確認

実際にRolling Bufferがどのように動作するか、小さな例で確認しましょう。

# Rolling Buffer の動作確認
w = 4
num_heads = 1
head_dim = 2
cache = RollingBufferKVCache(w, num_heads, head_dim)

print(f"窓サイズ: {w}")
print(f"固定メモリ: {cache.memory_bytes()} bytes\n")

# 8トークンを順に追加
for t in range(8):
    # ダミーのKVベクトル（位置番号を値として使用）
    new_k = torch.full((1, num_heads, 1, head_dim), float(t))
    new_v = torch.full((1, num_heads, 1, head_dim), float(t))
    cache.update(new_k, new_v)

    k, v = cache.get_kv()
    positions = k[0, 0, :, 0].tolist()
    idx = t % w
    print(f"Step {t}: buf_idx={idx}, "
          f"cache=[{', '.join(f'{p:.0f}' for p in positions)}], "
          f"valid_len={k.size(2)}")

出力を見ると、ステップ0〜3ではバッファが順に埋まっていき、ステップ4以降は最も古いKVが上書きされていきます。ステップ7の時点でバッファには位置4〜7のKVだけが残っており、これはまさに窓サイズ $w = 4$ の範囲です。メモリサイズはステップ4以降ずっと同じ値であり、生成トークン数に依存しないことが確認できます。

ここまでで、SWAとRolling Buffer KV Cacheのコア実装が完成しました。次に、Full AttentionとSWAの計算時間・メモリ使用量を実験的に比較し、理論的な分析を数値で裏付けましょう。

計算時間・メモリ比較実験

実験の目的と設定

SWAがFull Attentionに対してどの程度の計算時間・メモリ削減を実現するかを、系列長を変えながら測定します。実装は上で定義した scaled_dot_product_attention を使い、CPU上で実験します（GPUで測定する場合もCUDAイベントで同期すれば同様の傾向が得られます）。

def benchmark_attention(seq_lengths, window_size, d_k=64, num_heads=8,
                        num_trials=3):
    """Full Attention と SWA の計算時間・メモリを比較する"""
    results = {
        'seq_len': [],
        'full_time': [],
        'swa_time': [],
        'full_mem_mb': [],
        'swa_mem_mb': []
    }

    for n in seq_lengths:
        # ランダムなQ, K, Vを生成
        Q = torch.randn(1, num_heads, n, d_k)
        K = torch.randn(1, num_heads, n, d_k)
        V = torch.randn(1, num_heads, n, d_k)

        # Full Causal マスク
        full_mask = torch.tril(torch.ones(n, n, dtype=torch.bool))
        # SWA マスク
        swa_mask = create_swa_mask(n, window_size)

        # メモリ使用量（注意行列のサイズ）
        full_mem = n * n * 4 / (1024 ** 2)  # FP32, MB
        swa_nnz = swa_mask.sum().item()
        swa_mem = swa_nnz * 4 / (1024 ** 2)  # 非ゼロ要素分のみ（理想値）

        # 計算時間の測定
        # Full Attention
        times_full = []
        for _ in range(num_trials):
            start = time.perf_counter()
            _ = scaled_dot_product_attention(Q, K, V, full_mask)
            times_full.append(time.perf_counter() - start)

        # SWA
        times_swa = []
        for _ in range(num_trials):
            start = time.perf_counter()
            _ = scaled_dot_product_attention(Q, K, V, swa_mask)
            times_swa.append(time.perf_counter() - start)

        results['seq_len'].append(n)
        results['full_time'].append(np.median(times_full) * 1000)  # ms
        results['swa_time'].append(np.median(times_swa) * 1000)    # ms
        results['full_mem_mb'].append(full_mem)
        results['swa_mem_mb'].append(swa_mem)

        print(f"n={n:>6d}: Full={results['full_time'][-1]:.1f}ms, "
              f"SWA={results['swa_time'][-1]:.1f}ms, "
              f"speedup={results['full_time'][-1]/results['swa_time'][-1]:.2f}x")

    return results

実験の実行

系列長を256から4096まで変化させ、窓サイズ $w = 256$ で測定します。

seq_lengths = [256, 512, 1024, 2048, 4096]
window_size = 256

print("=" * 60)
print(f"Benchmark: Full Attention vs SWA (w={window_size})")
print("=" * 60)
results = benchmark_attention(seq_lengths, window_size)

注意すべき点として、ここでの実装は密行列演算（torch.matmul）にマスクを適用しているため、実際の計算量はFull AttentionもSWAも同じ $O(n^2)$ のmatmulを実行しています。本来のSWAの速度向上は、Flash Attentionのブロックスパース実装やカスタムCUDAカーネルで窓外の計算を完全にスキップすることで得られます。この実験では主にメモリ使用量の理論値と、マスクパターンの違いを確認することが目的です。

結果の可視化

fig, axes = plt.subplots(1, 2, figsize=(14, 6))

# 計算時間の比較
axes[0].plot(results['seq_len'], results['full_time'], 'o-',
             label='Full Attention', linewidth=2, markersize=8, color='#e74c3c')
axes[0].plot(results['seq_len'], results['swa_time'], 's-',
             label=f'SWA (w={window_size})', linewidth=2, markersize=8, color='#3498db')
axes[0].set_xlabel('系列長', fontsize=12)
axes[0].set_ylabel('計算時間 (ms)', fontsize=12)
axes[0].set_title('計算時間の比較: Full vs SWA', fontsize=14)
axes[0].legend(fontsize=12)
axes[0].grid(True, alpha=0.3)
axes[0].set_xscale('log', base=2)
axes[0].set_yscale('log', base=2)

# メモリ使用量の比較（理論値）
axes[1].plot(results['seq_len'], results['full_mem_mb'], 'o-',
             label='Full Attention ($O(n^2)$)', linewidth=2, markersize=8, color='#e74c3c')
axes[1].plot(results['seq_len'], results['swa_mem_mb'], 's-',
             label=f'SWA (w={window_size}, $O(nw)$)', linewidth=2, markersize=8,
             color='#3498db')
axes[1].set_xlabel('系列長', fontsize=12)
axes[1].set_ylabel('注意行列メモリ (MB)', fontsize=12)
axes[1].set_title('注意行列メモリの比較: Full vs SWA', fontsize=14)
axes[1].legend(fontsize=12)
axes[1].grid(True, alpha=0.3)
axes[1].set_xscale('log', base=2)
axes[1].set_yscale('log', base=2)

plt.tight_layout()
plt.savefig('swa_benchmark.png', dpi=150, bbox_inches='tight')
plt.show()

左のグラフ（計算時間）では、系列長が増えるにつれてFull AttentionとSWAの差が広がる傾向が確認できます。ただし前述のとおり、密行列演算ベースの実装ではマスクの有無にかかわらずmatmulのコストが支配的であるため、差は限定的です。右のグラフ（メモリ使用量）では、Full Attentionが二次曲線的に増加するのに対し、SWAは線形に近い増加を示しています。系列長4096で窓サイズ256の場合、注意行列のメモリは約16分の1に削減されることが確認できます。

KVキャッシュサイズの比較

推論時のKVキャッシュサイズも、Full AttentionとSWA（Rolling Buffer）で比較してみましょう。

def compute_kv_cache_size(seq_len, num_layers, num_kv_heads, head_dim,
                          window_size=None, dtype_bytes=2):
    """KVキャッシュのメモリ使用量を計算する（FP16想定）

    Args:
        seq_len: 系列長
        num_layers: レイヤー数
        num_kv_heads: KVヘッド数
        head_dim: ヘッド次元
        window_size: SWAの窓サイズ（Noneの場合はFull Attention）
        dtype_bytes: データ型のバイト数（FP16=2）

    Returns:
        メモリ使用量（MB）
    """
    effective_len = window_size if window_size else seq_len
    # Key + Value の2つ
    size_bytes = 2 * num_layers * num_kv_heads * effective_len * head_dim * dtype_bytes
    return size_bytes / (1024 ** 2)

# Mistral 7B のパラメータ
num_layers = 32
num_kv_heads_full = 32   # Full MHA
num_kv_heads_gqa = 8     # GQA
head_dim = 128
window_size = 4096

seq_lengths_kv = [4096, 8192, 16384, 32768, 65536, 131072]

print("=" * 80)
print("KV Cache Memory Comparison (FP16)")
print("=" * 80)
print(f"{'Seq Len':>10} | {'Full MHA':>12} | {'GQA only':>12} | "
      f"{'SWA+GQA':>12} | {'Reduction':>10}")
print("-" * 80)

kv_full_list, kv_gqa_list, kv_swa_gqa_list = [], [], []

for n in seq_lengths_kv:
    full = compute_kv_cache_size(n, num_layers, num_kv_heads_full, head_dim)
    gqa = compute_kv_cache_size(n, num_layers, num_kv_heads_gqa, head_dim)
    swa_gqa = compute_kv_cache_size(n, num_layers, num_kv_heads_gqa, head_dim,
                                     window_size)
    reduction = (1 - swa_gqa / full) * 100

    kv_full_list.append(full)
    kv_gqa_list.append(gqa)
    kv_swa_gqa_list.append(swa_gqa)

    print(f"{n:>10,} | {full:>10.1f} MB | {gqa:>10.1f} MB | "
          f"{swa_gqa:>10.1f} MB | {reduction:>8.1f}%")

この比較表から、SWA + GQAの組み合わせがいかに強力かが数値で読み取れます。系列長32,768（Mistral 7Bの最大系列長）では、Full MHAのKVキャッシュが16,384 MBに対し、SWA + GQAではわずか512 MB — 32分の1です。系列長が伸びるほど差は広がり、131,072トークンでは65,536 MBと512 MBで128倍の差になります。SWA + GQAではKVキャッシュサイズが系列長に依存しない定数値（512 MB）に固定されるため、長文脈対応のメモリ計画が非常にシンプルになります。

KVキャッシュのメモリ推移の可視化

fig, ax = plt.subplots(figsize=(10, 6))

ax.plot(seq_lengths_kv, kv_full_list, 'o-', label='Full MHA',
        linewidth=2, markersize=8, color='#e74c3c')
ax.plot(seq_lengths_kv, kv_gqa_list, 's-', label='GQAのみ (KVヘッド8)',
        linewidth=2, markersize=8, color='#f39c12')
ax.plot(seq_lengths_kv, kv_swa_gqa_list, '^-', label='SWA + GQA (Mistral 7B)',
        linewidth=2, markersize=8, color='#3498db')

# 一般的なGPUメモリの参考線
ax.axhline(y=24 * 1024, color='gray', linestyle='--', alpha=0.5, label='24 GB GPU')
ax.axhline(y=80 * 1024, color='gray', linestyle=':', alpha=0.5, label='80 GB GPU')

ax.set_xlabel('系列長', fontsize=12)
ax.set_ylabel('KVキャッシュ (MB)', fontsize=12)
ax.set_title('KVキャッシュの比較: Full MHA vs GQA vs SWA+GQA', fontsize=14)
ax.legend(fontsize=11)
ax.grid(True, alpha=0.3)
ax.set_xscale('log', base=2)
ax.set_yscale('log', base=2)

# x軸のラベルを読みやすく
ax.set_xticks(seq_lengths_kv)
ax.set_xticklabels([f'{n//1024}K' for n in seq_lengths_kv])

plt.tight_layout()
plt.savefig('kv_cache_comparison.png', dpi=150, bbox_inches='tight')
plt.show()

このグラフからは、3つの明確な傾向が読み取れます。まず、Full MHA（赤線）は系列長に対して線形に増加し、32Kトークンで既にGPU1枚の24 GBに近づいています。GQA単体（黄線）はKVヘッド数を4分の1にすることで同じく線形増加ですが傾きが緩やかになります。SWA + GQA（青線）は窓サイズ以上では完全に水平になり、系列長がいくら伸びてもKVキャッシュサイズが一定です。この「水平線」こそがSliding Window Attentionの最大の実用的価値であり、長文脈LLMのメモリ計画を根本的に変える特性です。

受容野の拡大を可視化する

最後に、多層スタックによる受容野の拡大を可視化しましょう。各レイヤーで窓サイズ $w$ のSWAを適用したとき、情報がどのように伝播するかを追跡します。

def compute_receptive_field(num_layers: int, window_size: int,
                             seq_len: int, target_pos: int) -> np.ndarray:
    """多層SWAでの受容野を計算する

    各レイヤーの出力位置 target_pos が参照できる
    元の入力位置の集合を追跡する。

    Returns:
        reachable: (num_layers+1, seq_len) のバイナリ配列
                   reachable[l][j] = 1 なら l層目で位置jの情報に到達可能
    """
    reachable = np.zeros((num_layers + 1, seq_len), dtype=np.int32)
    reachable[0, target_pos] = 1  # 初期状態: 自分自身のみ

    for layer in range(1, num_layers + 1):
        for pos in range(seq_len):
            if reachable[layer - 1, pos]:
                # この位置が前のレイヤーで到達可能なら、
                # その位置の窓内の全位置も到達可能
                start = max(0, pos - window_size + 1)
                reachable[layer, start:pos + 1] = 1

    return reachable

# パラメータ設定
n_vis = 64
w_vis = 8
L_vis = 8
target = n_vis - 1  # 最後のトークン

reachable = compute_receptive_field(L_vis, w_vis, n_vis, target)

fig, ax = plt.subplots(figsize=(14, 6))
ax.imshow(reachable, cmap='Blues', aspect='auto', interpolation='nearest')
ax.set_xlabel('入力の位置', fontsize=12)
ax.set_ylabel('層', fontsize=12)
ax.set_title(f'受容野の拡大 (w={w_vis}, 注目トークン=位置{target})',
             fontsize=14)
ax.set_yticks(range(L_vis + 1))
ax.set_yticklabels([f'{l}層目' for l in range(L_vis + 1)])

# 各レイヤーの受容野サイズを注記
for l in range(L_vis + 1):
    rf_size = reachable[l].sum()
    ax.text(n_vis + 1, l, f'受容野={rf_size}', va='center', fontsize=10, color='#2c3e50')

plt.tight_layout()
plt.savefig('receptive_field_expansion.png', dpi=150, bbox_inches='tight')
plt.show()

このヒートマップでは、横軸が入力位置、縦軸がレイヤーを表しています。Layer 0（初期状態）では最後のトークン位置だけが青く、Layer 1では窓サイズ $w = 8$ 分だけ左に広がります。レイヤーが深くなるにつれて青い領域が左へ拡大していき、Layer 8では受容野が57位置まで広がっていることが確認できます。これは先に示した理論式 $R(L) = L(w-1) + 1 = 8 \times 7 + 1 = 57$ と正確に一致しています。重要な点は、受容野が線形に（指数的にではなく）拡大することです。CNNのdilated convolutionが受容野を指数的に拡大するのとは対照的であり、SWAの受容野拡大は穏やかです。この性質がSWAの表現力の限界でもあり、非常に長距離の依存関係を捉えるには十分なレイヤー数が必要です。

発展的な話題

Pre-fill と Chunked Prefill

Mistralの実装では、推論をpre-fill（プロンプト処理）フェーズとgeneration（トークン生成）フェーズに分けます。Pre-fillフェーズでは入力プロンプト全体を並列に処理し、SWAマスク付きのAttentionを計算します。この段階ではKVキャッシュをRolling Bufferに順次書き込みます。

長いプロンプト（窓サイズ $w$ を超えるもの）の場合、Chunked Prefillが使われます。プロンプトを窓サイズ $w$ ごとのチャンクに分割し、各チャンク内でSWAを計算します。チャンク間の依存関係はKVキャッシュを通じて引き継がれるため、メモリ使用量は常に $O(w)$ に収まります。

Mistral以降のSWA採用モデル

Mistral 7BでSWAの有効性が実証されて以降、多くのモデルがSWAまたはその変形を採用しています。

Mixtral 8x7BはMistral 7Bのアーキテクチャにスパースなmixture-of-experts（MoE）を組み合わせたモデルで、SWAもそのまま継承しています。Gemma（Google）も局所的なSliding Window Attentionのレイヤーを含むハイブリッド設計を採用しており、一部のレイヤーでFull Attentionを、残りのレイヤーでSWAを使う構成になっています。

このハイブリッドアプローチは、SWAの効率性とFull Attentionの長距離依存捕捉能力を両立する設計として注目されています。全レイヤーをSWAにするか、一部をFull Attentionにするかはモデルの用途とリソースに依存する設計選択であり、今後も発展が続く領域です。

スパース性を活用した高速化 — Flash AttentionとBlockSparse

SWAの理論的な計算量削減を実際の速度向上に変換するには、スパースなマスクパターンに対応したカスタムカーネルが必要です。Flash Attention 2のブロック分割アルゴリズムは、SWAの帯行列構造と自然に適合します。注意行列をブロック（タイル）に分割したとき、窓の外に完全に位置するブロックは計算自体をスキップできるため、ウォールクロック時間での高速化が実現します。

xFormersライブラリやFlash Attention 2の最新実装では、SWA用のカーネルが提供されており、window_size パラメータを指定するだけで帯行列構造のAttentionが効率的に計算されます。

まとめ

本記事では、Sliding Window Attention（SWA）の理論と実装を詳しく解説しました。

• 基本原理: SWAは各トークンの注意範囲を固定幅の窓 $w$ に制限することで、Attentionの計算量を $O(n^2)$ から $O(nw)$ に削減します。窓外のトークンには $-\infty$ マスクを適用してsoftmaxでゼロにする、極めてシンプルな手法です
• 受容野の拡大: 1層では $w$ トークンの受容野ですが、$L$ 層をスタックすると理論上 $R(L) \approx Lw$ の受容野を持ちます。Mistral 7B（$w = 4096$、$L = 32$）では約13万トークンの受容野を確保しています
• Rolling Buffer KV Cache: 固定サイズ $w$ のリングバッファでKVキャッシュを管理することで、メモリの動的確保を排除し、生成トークン数によらないメモリ上限を保証します
• GQAとの相乗効果: SWA（系列長方向の削減）とGQA（ヘッド数方向の削減）を組み合わせることで、Mistral 7BのKVキャッシュはFull MHAの32分の1に削減されます
• Longformerとの違い: 同じ局所注意でも、SWAはデコーダ向け（因果マスク、グローバルトークンなし）、LongformerはBERTエンコーダ向け（双方向、グローバルトークンあり）と設計思想が異なります

SWAは「Attentionの局所性」という自然言語の統計的性質を正面から活用した手法であり、その単純さゆえに実装が容易で、Flash AttentionやvLLMなどの推論フレームワークとの親和性が高いという実用上の強みがあります。

次のステップとして、以下の記事も参考にしてください。