一文要約: シーケンス長が数千から数百万に跳ね上がると O(N²) の完全 Attention は壁にぶつかる — Sparse Attention はどのトークンに注目するかを選択し、Linear Attention は計算順序を変え、Infini Attention は大きくならない圧縮メモリを追加します。

24.1 完全 Attention が壁にぶつかる

24.1.1 数字は正直に語る

標準的な Self-Attention はすべてのトークンペアのスコアを計算しなければなりません。長さ N のシーケンスに対して N×N マトリクスが必要です:

シーケンスが 10 倍になると計算量は 100 倍になります。O(N²) が「二次」と呼ばれる理由です — 入力の成長より速く加速します。

24.1.2 メモリが真のボトルネック

計算だけなら時間をかければ済みますが、メモリは伸ばせません。FP16 での完全な Attention マトリクスのコスト:

128k トークンコンテキストを処理する 32 層モデルは、すべてのレイヤーの Attention マトリクスを保持するだけで 1 TB 超 が必要です。プロダクションにそんな GPU はありません。

第21章の Flash Attention は完全なマトリクスを HBM に実体化しないことで助けになりますが、根本的な O(N²) の計算を排除できるわけではありません — すべてのペアは依然として処理されます。

24.1.3 需要は現実にある

実際のアプリケーションは標準的な学習よりはるかに長いコンテキストを必要とします:

• 長いコードレビュー: 差分を含む大きなプルリクエストは数万行に及ぶことがある
• ドキュメント理解: 技術マニュアル、法的書類、書籍分量のコンテンツ
• エージェントのタスク履歴: 数百のツール呼び出しを実行したエージェントは長いトレースを蓄積する
• 動画: 1 時間のエンコードされた動画フレームは数十万トークンになり得る

これが Sparse Attention、Linear Attention、Infini Attention が解決するために設計された問題です。

24.2 Sparse Attention: すべてのトークンがすべてのトークンを見る必要はない

24.2.1 直感

技術文書の 500 ページ目を読むとき、1 ページ目からすべての単語に同時にアクセスする必要があるでしょうか？ いいえ。必要なのは:

1. 周囲のコンテキスト（直前の数段落）
2. 特定のグローバルアンカー（セクション見出し、重要な定義）
3. ときどき、特定の以前の内容

Sparse Attention はこの構造を Attention マスクに組み込みます。密な N×N の接続マトリクスの代わりに、慎重に選ばれたスパースなパターンを使います。

24.2.2 三つの基本的なスパースパターン

Random Attention: 各トークンがランダムにサンプリングされた他のトークンの集合に Attention を当てます。無秩序に見えますが、良い理論的特性を持ちます。単一のトークンがシーケンス全体を見なくても、情報はグラフ上を O(1) ホップで伝播できます。

スライディングウィンドウ（ローカル）Attention: 各トークンが両側の w トークンに Attention を当てます。これは自然言語の強い局所性を反映します — 構文的・意味的な依存関係のほとんどは短距離です。計算量: O(N × w)。

グローバル Attention: 少数の特別なトークン（[CLS]、タスク固有のトークン、各文の最初のトークンなど）がすべての他のトークンに Attention を当て、またすべての他のトークンからも Attention を受けます。これらのグローバルな「ハブ」がシーケンス全体の情報を収集・再配布します。すべてのトークンをすべてと接続する必要がありません。

24.3 Longformer と BigBird

24.3.1 スパース Attention マトリクスをステップごとに構築する

ステップ 1 — ランダムトークン: 長距離の情報フローを確保するために疎なランダム接続を追加します（黄色で表示）。

ステップ 2 — ウィンドウトークン: 局所的な接続の対角バンドを追加します。各トークンが隣接するトークンに Attention を当てます。

ステップ 3 — グローバルトークン: いくつかの位置をグローバルハブとして指定します。それらの行と列が全部埋まります。すべてが最大一つのグローバルハブを経由して他のすべてに到達できます。

最終的なマトリクスはスパースです。黄色のセルは計算されます。灰色のセルは完全にスキップされます。

24.3.2 BigBird: ランダム + ウィンドウ + グローバル

BigBird は三つのパターンをすべて組み合わせます:

BigBird Attention = ランダム + ウィンドウ + グローバル

この組み合わせは理論的に完全なチューリング計算と同等であることが証明されています — つまり BigBird は原理的に、完全 Attention が表現できるものをより少ない操作で表現できます。

ウィンドウサイズ w、g 個のグローバルトークン、トークンあたり r 個のランダム接続での計算量:

O(N × (w + g + r))

w、g、r が定数（N とともに成長しない）の場合、これは O(N) — シーケンス長に対して線形です。

24.3.3 Longformer と BigBird の比較

Longformer はシンプルでタスク固有です。QA モデルでは質問トークンをグローバルとして指定し、文書全体を Attention できるようにします。BigBird のランダムコンポーネントが提供するのは、任意の 2 トークンが有界なホップ数以内で通信できるという理論的保証です。

24.4 スライディングウィンドウ Attention の実際

24.4.1 局所性の仮定

言語には強い局所構造があります:

• 主語-動詞の一致はほとんどの文で少数の単語にまたがる
• 代名詞の解消は通常、最も近い候補名詞を指す
• 同じ段落内の文は同じトピックを扱う

トークンの大多数にとって、最も関連性の高いコンテキストは近くにあります。サイズ w のスライディングウィンドウはこれを O(N²) ではなく O(N × w) で捉えます。

24.4.2 実装

def sliding_window_attention(Q, K, V, window_size):
    n = Q.shape[0]
    output = []
    for i in range(n):
        start = max(0, i - window_size // 2)
        end   = min(n, i + window_size // 2 + 1)
        local_K = K[start:end]
        local_V = V[start:end]
        scores  = Q[i] @ local_K.T / math.sqrt(d_k)
        weights = softmax(scores)
        output.append(weights @ local_V)
    return torch.stack(output)

24.4.3 レイヤーを重ねると受容野が広がる

w = 5 の単一スライディングウィンドウレイヤーは 5 トークンを見ます。2 層重ねると 9 トークン見えます。一般に L 層後の実効受容野は:

1 + L × (w - 1)

w = 512 の 32 層モデルは理論上 16,000 トークン超の受容野に達します — ほとんどの実用タスクで十分な量です。

24.4.4 Mistral 7B のローリングバッファキャッシュ

Mistral 7B は w = 4096 のスライディングウィンドウ Attention を使用します。便利なトリック: KV キャッシュは直近の w トークンだけを保存すればよく、会話の長さにかかわらずキャッシュメモリが一定になります。

標準 KV キャッシュ: すべての履歴を格納 → 時間とともに線形に成長
ローリングバッファ: 直近の w トークンのみ格納 → 一定サイズ

これは永遠に埋まり続けるメモリと、会話が成長するにつれて前に進むメモリの違いです。

24.5 Linear Attention: 計算順序を変える

24.5.1 O(N²) がどこから来るのか

標準 Attention:

Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d_k)) @ V

計算順序:

1. Q @ K^T → N×N マトリクス — これが O(N²)
2. softmax(...) → N×N マトリクス
3. 結果 @ V → N×d マトリクス

中間の N×N マトリクスが二次計算と二次メモリの両方の源です。

24.5.2 核心的な洞察: 結合則

行列の積は結合則が成り立ちます:

(Q @ K^T) @ V  =  Q @ (K^T @ V)

K^T @ V を先に計算すると:

• K^T @ V の形状は d × d — N×N ではない
• Q @ (K^T @ V) の形状は N × d

中間マトリクスが N×N ではなく d × d になります。N >> d（長いシーケンスではよくある）の場合、これは大きな勝利です。

計算量の比較:

• 標準: N×N マトリクスで O(N² × d)
• Linear（順序変更）: d×d マトリクスで O(N × d²)

N = 100,000、d = 64 の場合: 標準は約 6,400 億演算、Linear は約 4 億 1,000 万演算。3 桁の差です。

24.5.3 Softmax の問題

障害があります: softmax が結合則を壊します。

softmax(Q @ K^T) @ V  ≠  Q @ softmax(K^T) @ V

Linear Attention はこれを softmax をカーネル関数 φ で置き換える ことで解決します:

標準 Attention:  softmax(Q @ K^T) @ V
Linear Attention: φ(Q) @ φ(K)^T @ V  =  φ(Q) @ (φ(K)^T @ V)

ここで φ は通常 elu(x) + 1 のようなもの — 値を正に保つ（有効な Attention 分布に必要）が順序の入れ替えを可能にする関数です。

24.5.4 Linear Attention のコード

def linear_attention(Q, K, V, feature_map):
    """
    Q, K, V: [batch, seq_len, d_model]
    feature_map: 例 lambda x: F.elu(x) + 1
    """
    Q_prime = feature_map(Q)   # [batch, seq_len, d_model]
    K_prime = feature_map(K)   # [batch, seq_len, d_model]

    # K^T @ V を先に計算: [batch, d_model, d_model]
    KV = torch.einsum('bnd,bnm->bdm', K_prime, V)

    # Q @ (K^T @ V): [batch, seq_len, d_model]
    output = torch.einsum('bnd,bdm->bnm', Q_prime, KV)

    # 正規化
    Z = torch.einsum('bnd,bd->bn', Q_prime, K_prime.sum(dim=1))
    return output / Z.unsqueeze(-1)

24.5.5 Linear Attention がデフォルトにならない理由

Linear Attention は理論的にエレガントです。しかし実際には:

• 表現力が落ちる — softmax をカーネルで置き換えることで、完全 Attention が生み出せる鋭く集中した Attention 分布が失われます。タスクによってはその鋭さが必要です。
• 数値的不安定性 — 特定のカーネル選択が学習中に勾配の問題を引き起こす可能性があります。
• ベンチマークでの品質差 — ほとんどの標準 NLP タスクで、Linear Attention は完全 Attention より劣ります。その差は目立つことが多いです。

これが、理論的な魅力にもかかわらず Linear Attention がプロダクションモデルで完全 Attention に取って代わっていない理由です。その貢献は概念的なものです。計算を並べ替えて N×N マトリクスを排除するというアイデアが、Infini Attention に直接インスピレーションを与えました。

24.6 Infini Attention: 固定メモリで無制限のコンテキスト

24.6.1 KV キャッシュが解決できない問題

スライディングウィンドウ Attention を使っても、成長する会話はやがてウィンドウを超えます。モデルは最初の部分を忘れます。

長いコードベースを読み通して質問に答えているエージェントを想像してください。次々とターンが積み重なりコンテキストが増えていきます。従来の KV キャッシュは線形に成長します — 十分なターンの後、メモリが尽きるか、最古のコンテキストがウィンドウから外れるかします。

私たちが欲しいのは: 完全な履歴を捉えながら固定された有限のストレージを使うメモリです。

24.6.2 核心的なアイデア

Google の Infini Attention（2024）は各 Attention レイヤーに固定サイズ d × d の圧縮メモリマトリクスを与えます。新しいセグメントが届くたびに情報がこのマトリクスに圧縮されます。マトリクスは大きくなりません — これまでに見たすべてのものの「ローリングサマリー」を保持します。

各 Attention レイヤーは二つのコンポーネントを持ちます:

1. ローカル Attention — 現在のセグメントに対する標準（または Flash）Attention。直近の正確なコンテキストを捉えます。

2. 圧縮メモリ — すべてのセグメント境界でクエリと更新が行われる固定の d × d マトリクス。過去のすべてのセグメントから要約した情報を運びます。

最終的な出力は両方をブレンドします:

出力 = β × メモリ出力 + (1 - β) × ローカル出力

ここで β はレイヤーごとの学習可能なゲートです。

24.6.3 圧縮メモリの仕組み

Infini Attention は Linear Attention のアイデアを借ります: K^T @ V はキーバリューペアの d × d サマリーです。各セグメント後にこれを捨てる代わりに、Infini Attention はそれを保持して蓄積します:

class InfiniAttentionLayer:
    def __init__(self, d_model):
        self.memory     = torch.zeros(d_model, d_model)
        self.normalizer = torch.zeros(d_model)

    def forward(self, Q, K, V):
        # 圧縮メモリからの取得（Linear Attention スタイル）
        memory_out = Q @ self.memory / (Q @ self.normalizer + eps)

        # 現在のセグメントに対するローカル Attention（標準 Attention）
        local_out = standard_attention(Q, K, V)

        # 学習可能なゲートでブレンド
        beta = torch.sigmoid(self.gate)
        output = beta * memory_out + (1 - beta) * local_out

        # 圧縮メモリを更新（インクリメンタル）
        self.memory     = self.memory     + K.T @ V
        self.normalizer = self.normalizer + K.sum(dim=0)

        return output

メモリの更新は加算的です — 新しい K/V 情報が折り込まれ、古い情報は新しい更新によって徐々に「希釈」されます。これが圧縮です: マトリクスは常に d × d のままですが、もはや特定の過去のトークンを正確に表現していません。

24.6.4 なぜ「無制限」のコンテキストが実現できるのか

メモリマトリクスのサイズは d × d です。シーケンス長や会話のターン数に関わらず成長しません。100 万トークンを処理しても 1,000 万トークンを処理しても — 同じ固定マトリクスがサマリーを保持します。

これは人間のワーキングメモリに類似しています。去年読んだ本のすべての文を思い出せないかもしれませんが、その内容について推論できる圧縮されたサマリーを保持しています。詳細は失われていますが、本質は保たれています。

24.6.5 パフォーマンス比較

*総履歴長に対して O(1)。各セグメントは内部的に O(セグメント²) または O(セグメント) の計算が必要。

24.6.6 Gemini 1.5 との関係

Google は Infini Attention 論文の直後に 100 万トークンコンテキストウィンドウをサポートする Gemini 1.5 をリリースしました。アーキテクチャの詳細は公開されていませんが、タイミングと技術的整合性から、多くの研究者が Infini スタイルのメモリが百万トークンコンテキストをプロダクションで実現する一因だと考えています。

Infini Attention が率直に認めているトレードオフ: 圧縮は情報損失を意味します。文書の冒頭の詳細の正確な想起を必要とするタスク — 長い会話のターン 1 からの特定の数字など — では圧縮メモリは十分でない可能性があります。長い履歴にわたる一般的なテーマやパターンを推論するタスクでは、多くの場合うまく機能します。

24.7 技術を組み合わせる: 実践的なツールキット

プロダクションでは、これらの技術は排他的ではありません。一般的な組み合わせ:

10k〜50k トークンの場合:

Flash Attention + GQA + スライディングウィンドウ

成熟しており、サポートが充実していて、品質リスクが最小限。

50k〜200k トークンの場合:

Flash Attention + Sparse Attention (BigBird/Longformer スタイル) + RoPE または ALiBi

スパースパターンのより慎重なチューニングが必要。スパースパターンはハードウェアフレンドリーである必要があります。

200k+ トークンまたは無制限に成長する会話の場合:

Infini Attention + Flash Attention

真に無制限の成長を扱えるのは圧縮メモリだけです。

ツールキットのその他のレバー:

• 小ウィンドウ事前学習 + RoPE スケーリングによる大ウィンドウファインチューニング（第25章）
• KV キャッシュ量子化（INT8/FP8）でキャッシュフットプリントを削減
• PagedAttention（vLLM）— OS が仮想メモリを管理するように KV ページを管理
• スパース MLP ルーティング（ライト/ヘビーブランチ）— 単純なトークンをより安価なパスに

これらは積み重ねられます。100 万トークンコンテキストの実サービングシステムでは、Flash Attention（効率的なカーネル）、GQA（小さいキャッシュ）、NTK スケーリング付き RoPE（位置外挿）、ページ KV 管理（リクエスト間のキャッシュ再利用）を組み合わせることがあります。

24.8 具体例

24.8.1 完全 Attention とスライディングウィンドウ: 計算カウント

シーケンス長 N = 8、次元 d = 4。

完全 Attention:

• Q @ K^T: (8×4) × (4×8) → 64 出力、それぞれ 4 回の乗算 = 256 演算
• 結果 @ V: (8×8) × (8×4) → 32 出力、それぞれ 8 回の乗算 = 256 演算
• 合計: 約 512 回の乗算

スライディングウィンドウ（w = 3）:

• 各トークンが 3 つの隣接トークンに Attention を当てる
• トークンごとのスコア: 3 × d = 12 回の乗算; 8 トークン = 96 演算
• トークンごとのバリュー合計: 3 × d = 12 回の乗算; 8 トークン = 96 演算
• 合計: 約 192 回の乗算 — 完全 Attention より 62% 少ない

24.8.2 スケールにおける Linear Attention の優位性

N = 1000、d = 64 の場合:

このスケールで Linear Attention は約 87.5% 少ない演算。優位性は N が大きくなるにつれて拡大します。

24.9 よくある誤解

「Sparse Attention は重要な情報を失う」 — BigBird のような適切に設計されたスパースパターンは、グラフ理論的な保証を提供します: グローバルまたはランダム接続を経由して、任意の 2 トークンが有界なホップ数以内で通信できます。情報は伝播できます。ただし複数のレイヤーを経由することになります。

「Linear Attention は標準 Attention と等価」 — そうではありません。softmax をカーネル関数で置き換えると、実効的な Attention 分布が変わります。遠くから特定のトークンをコピーするような鋭い Attention を必要とするタスクでは、Linear Attention はしばしば劣ります。これは異なる帰納バイアスであり、無料の近似ではありません。

「Infini Attention はすべての履歴を完全に保持する」 — 圧縮は常に情報を失います。Infini Attention は読書中に取るメモのような圧縮されたサマリーとして考えるのが最善です。初期トークンの正確な想起は保証されません。

「これらの技術は互いに排他的」 — 積み重ねるように設計されています。プロダクションシステムでは Flash Attention（効率的なカーネル）+ GQA（小さいキャッシュ）+ スライディングウィンドウ（削減された計算）+ ページメモリ（キャッシュ管理）を日常的に組み合わせます。それぞれが長いコンテキスト問題の異なる部分に対処します。

24.10 章のまとめ

チャプターチェックリスト

この章を終えた後、以下ができるようになっているはずです:

• O(N²) が実行不可能になる理由とコストの源を説明できる。
• ランダム、ウィンドウ、グローバルのスパース Attention パターンを説明できる。
• BigBird の組み合わせが情報フロー保証を持つ O(N) 計算量を提供する理由を説明できる。
• Linear Attention が計算を並べ替える方法と何を諦めるかを説明できる。
• Infini Attention の圧縮メモリが固定サイズのままである理由を説明できる。
• 与えられたコンテキスト長目標に対して適切な技術の組み合わせを選べる。

参考文献

1. Longformer: The Long-Document Transformer (Beltagy et al., 2020)
2. Big Bird: Transformers for Longer Sequences (Zaheer et al., 2020)
3. Efficient Attention: Attention with Linear Complexities (Shen et al., 2018) — arXiv 1812.01243
4. Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention (Google, 2024)
5. Mistral 7B Technical Report — ローリングバッファ KV キャッシュとスライディングウィンドウ

次の章へ

長いシーケンスのメモリと計算コストに対処してきました。しかし、これまで脇に置いていた関連する問題があります: 位置エンコーディングです。

モデルを 4k トークンから 128k トークンに拡張するとき、位置エンベディングは学習中に見たことのない位置に遭遇します。第25章では位置エンコーディングがどのように進化したかを説明します — 正弦波エンコーディングから RoPE、ALiBi、YaRN まで — この外挿の課題を特に扱うために。次章でお会いしましょう。