Sliding Window Attention

スライディングウィンドウアテンション（SWA）は、各トークンが固定数の隣接トークンのみを参照できるようにします。因果的な言語モデルでは、系列全体に対する完全な自己注意の代わりに、直近のw個のトークンを参照します。これにより、注意計算量は系列長に対して二次ではなく線形に増加し、生成モデルがKVキャッシュに保持する必要があるキーと値も直近w個のトークン分だけになります。Longformerは2020年にスライディングウィンドウを注意機構の中核に採用し、Mistral AIのMistral 7Bは2023年、ウィンドウサイズ4,096トークンで、この方式を広く利用されるオープンな大規模言語モデル（LLM）に導入しました。

層の深さに伴うウィンドウの拡大#

SWAの1層が参照するのはw個のトークンだけですが、Transformerでは層を積み重ねます。第2層では、あるトークンが、第1層ですでにw個のトークンを要約した隠れ状態を読み取るため、畳み込みニューラルネットワーク（CNN）の層を積み重ねると受容野が広がるのと同じように、層が増えるごとに情報は約w − 1位置先まで伝わります。PyTorchでは、ウィンドウはブール値の注意マスクとして表され、そのマスクを自乗すると情報が届く範囲を追跡できます。

import torch

seq_len, window = 16, 4
pos = torch.arange(seq_len)
offset = pos[:, None] - pos[None, :]

# Causal sliding window: each token attends to itself and the window - 1 tokens before it
mask = (offset >= 0) & (offset < window)
print(f"Scores per layer: {mask.sum().item()} with the window, {(offset >= 0).sum().item()} with full causal attention")

# Stacking layers widens the reach: each layer adds window - 1 more tokens of history
reach = mask
for layer in range(1, 5):
    print(f"Layer {layer}: last token draws on {reach[-1].sum().item()} tokens")
    reach = (mask.int() @ reach.int()) > 0

16トークンに対するスコア計算は、ウィンドウによって1層あたり136個ではなく58個になります。また、最後のトークンが参照できる範囲は、層ごとに3トークンずつ増えて、4、7、10、13となります。同じマスクをPyTorchのscaled_dot_product_attentionにattn_maskとして渡せます。Mistral 7Bでは、4,096トークンのウィンドウを持つ32層によって、理論上約131Kトークンの注意範囲が得られます。ただし、遠く離れた位置の情報が届くには、多数の層を経由して保持される必要があります。

Mistral 7Bは、キーと値もw個のスロットからなるローリングバッファに保存します。位置iのデータをスロットi mod wに書き込み、最も古いエントリーを上書きします。32Kトークンの系列では、品質を損なわずにキャッシュメモリを8分の1に削減でき、キャッシュサイズが固定されるため、生成がどれだけ長く続いてもメモリ使用量は一定です。

ローカル層とグローバル層の交互配置#

純粋なウィンドウ方式では、遠く離れた位置にある情報を一度の処理でコピーできないため、多くのモデルはSWAレイヤーとグローバルレイヤーを組み合わせています。GoogleのGemma 2は4,096トークンのウィンドウで両者を交互に配置し、Gemma 3はグローバルレイヤー1つごとにローカルレイヤーを5つ使用し、ウィンドウはわずか1,024トークンです。これにより、最大128KトークンのコンテキストウィンドウでもKVキャッシュを小さく保てます。コンテキスト全体のキーと値を保存するのはグローバルレイヤーだけなので、6レイヤー中5レイヤーはそれぞれ1,024トークンのみを保持します。

ビジョントランスフォーマーも、高解像度画像では同じ組み合わせを使います。この場合、各パッチが1つのトークンになります。Ultralyticsパッケージで画像セグメンテーションを実行できるMetaのSAM 3は、画像エンコーダーの32ブロック中28ブロックで、パッチの重ならない24 × 24ウィンドウ内に注意を向けます。また、ブロック7、15、23、31では、1,008ピクセル画像の72 × 72パッチグリッド全体に対してグローバルに注意を向けます。

SWAと関連用語#

• Swin Transformer：画像を固定サイズの重ならないウィンドウに分割し、各ウィンドウ内のすべてのトークンがそのウィンドウを共有します。層ごとにグリッドをずらすことで、隣接するウィンドウ間で情報を交換します。
• スパースアテンション：固定パターン全般を指します。SWAはその最も単純な手法で、ほかの手法では拡張ウィンドウ、グローバルトークン、ランダムな接続などを追加します。
• アテンションシンク：シーケンスの先頭にある、余ったアテンションの重みを吸収するトークンです。フルアテンションで学習したモデルに推論時のウィンドウを適用する際、これらを追い出すと生成が破綻します。StreamingLLMは、ウィンドウの隣にあるキャッシュ内にこれらのトークンを保持します。
• 変形可能アテンション：各クエリは、オフセットを学習する少数のサンプリング点に注意を向けるため、近傍領域はコンテンツに応じて変わります。
• Flash attention：厳密なアテンションを実行するGPUカーネルです。実装でウィンドウサイズを指定できるため、SWAモデルはウィンドウ外のブロックを完全にスキップします。