Sparse Attention

スパースアテンションは、シーケンス内のすべてのキーではなく、選択したサブセットだけに対して各クエリがスコアを計算するアテンション機構の一群です。N個のトークンに対する完全なセルフアテンションでは、ヘッドごとにN²個のスコアを計算するため、文書の長さが2倍になるとアテンションの計算コストは4倍になります。一方、各トークンがk個の他のトークンに注目するパターンでは、コストはN·kです。パターンは、ローカルウィンドウのように事前に固定することも、クエリごとに最も高いスコアのキーを選ぶように、内容に基づいて決めることもできます。OpenAIの2019年のSparse Transformerは、数万ステップのシーケンスをモデル化するため、O(N√N)のコストがかかるストライドパターンを採用しました。現在の大規模言語モデル (LLMs)は、10万トークンを超えるコンテキストウィンドウを処理するためにスパースレイヤーを使用しています。

一般的なスパースパターン#

固定パターンでは、スコアを計算するペアを事前に決めます。スライディングウィンドウアテンションでは、各トークンが直前のw個のトークンだけを参照します。Mistral 7Bは4,096トークンのウィンドウを使用し、OpenAIのオープンウェイトgpt-ossモデルは、密なレイヤーと128トークンのウィンドウレイヤーを交互に配置します。これは、GPT-3が採用した、密なレイヤーと局所的な帯状レイヤーの交互配置と同じです。各ウィンドウレイヤーは情報をさらにwトークン先まで伝えるため、深く積み重ねても遠く離れたトークン同士がつながります。ストライドパターンでは、s個ごとに過去のトークンを追加することで、単一レイヤー内で情報を遠くまで伝えます。グローバルトークンはすべての位置に注目し、すべての位置から注目されます。Longformerは分類トークンなどのタスクトークンをグローバルにし、BigBirdはランダムな接続も加えます。また、アテンションシンクでは、ストリーミング生成を安定させるため、最初の数個のトークンをすべてのクエリから参照できるようにします。

内容ベースのパターンでは、データに基づいてペアを選びます。DeepSeekのNative Sparse Attention (NSA、2025年)は、キーのブロックを圧縮した要約、最も関連性の高いブロックの選択、スライディングウィンドウを組み合わせ、すべてをエンドツーエンドでトレーニングします。DeepSeek-V3.2のDeepSeek Sparse Attentionは、小規模な「lightning indexer」を実行して、クエリごとに過去のすべてのトークンをスコアリングし、上位2,048個だけに対して完全なアテンションを計算します。そのため、コンテキストがどれだけ長くなっても、各クエリの主要なアテンションコストは一定です。本番環境のモデルでは、これらのルールを組み合わせることがよくあります。

パターンを計算量の削減につなげる#

密なスコア行列にマスクを適用しても、すべてのスコアが計算されます。そのため、高速化には、行列の空のタイルをスキップするカーネルが必要です。Flash attentionカーネルはスライディングウィンドウを直接サポートし、PyTorchのFlexAttentionは任意のマスク関数をブロックスパースなGPUカーネルにコンパイルします。そのcreate_block_maskは、1,024トークンのウィンドウと4つのシンクトークンを使う場合に、カーネルがどのブロックをスキップするかを示します。

from torch.nn.attention.flex_attention import create_block_mask

seq_len, window, sinks = 8192, 1024, 4


def local_plus_sinks(b, h, q_idx, kv_idx):
    """Causal attention to the previous 1,024 tokens plus the first 4 tokens."""
    causal = q_idx >= kv_idx
    return causal & ((q_idx - kv_idx <= window) | (kv_idx < sinks))


# FlexAttentionは、カーネルが空のブロックをスキップできるよう、マスクを128 x 128のブロックに分割します
block_mask = create_block_mask(local_plus_sinks, B=None, H=None, Q_LEN=seq_len, KV_LEN=seq_len, device="cpu")
print(block_mask)

出力には、8,192 × 8,192のスコア行列のうち、許可されたペアを含まない128 × 128ブロックの割合であるsparsity=85.47%が表示されます。また、ウィンドウを表す対角線の帯と、シンクトークンを表す左端の列としてパターンを描画します。flex_attentionをtorch.compile上で実行すると、カーネルはそれらのブロックを完全にスキップします。

スパースアテンションと関連用語の比較#

• Flash attention: 高速なオンチップメモリーで処理をタイル分割し、すべてのペアに対する正確なアテンションを計算する、I/Oを考慮したカーネルです。そのタイルはブロックスパースカーネルがスキップする単位でもあります。
• 線形アテンション: すべてのペアを保持しながらsoftmaxを特徴マップに置き換えるため、過去のすべてのキーと値を固定サイズの状態に圧縮できます。
• 変形可能アテンション: 固定パターンに従って注目するのではなく、各クエリが特徴マップ上の学習済みオフセットにある少数のサンプリング点に注目するビジョン向けの機構です。
• 混合エキスパート (MoE): パラメーターをスパースにする手法です。各トークンは少数のエキスパートネットワークだけを活性化しますが、アテンションパターンは密なままです。

用途#

長いコンテキストを扱う言語モデルが主な用途です。スパースレイヤーは計算量を抑え、ウィンドウレイヤーはKVキャッシュも抑えるため、モデルは1回のリクエストで長いレポートを要約したり、大規模なコードベースを検索したりできます。BigBirdは、一般的な文よりはるかに長い数千塩基対の入力からプロモーター領域を予測するため、スパースパターンをDNA配列に適用しました。コンピュータービジョンでは、画像解像度に応じてトークン数が増えるため、モデルはアテンションを局所領域に制限します。Swin Transformerは、ずらしたウィンドウ内でアテンションを行います。バッファロー大学と中国科学院大学の研究者が開発し、Ultralyticsパッケージで実行できる、アテンション中心の検出器YOLO12はエリアアテンションを使用します。ストライド16のバックボーン段階で、640ピクセル画像の40 × 40の特徴マップを400トークンずつの横長の4領域に分割し、各領域内だけでアテンションを行います。これにより、1,600トークンすべてに対する完全なアテンションで必要となる256万ペアの4分の1だけをスコアリングします。