一文でまとめると：Sparse Attention は計算する接続を減らし、Linear Attention は類似度の定義を変え、Infini-attention は segment をまたぐ古い K/V を有限状態へ圧縮します。三者が解く問題は同じではなく、どれも無制限の履歴を無料で保存する仕組みではありません。

24.1 Full Attention には三つの請求書がある

24.1.1 接続数

長さ $N$ の Full Self-Attention は、すべての Query とすべての Key を比較します。双方向 Attention の接続数は $N^2$ 、causal Attention で許される下三角の接続数は次の通りです。

1+2+\cdots+N=\frac{N(N+1)}{2}

「論理的には下三角だけが有効」であっても、通常の dense matmul が自動的に残り半分を省くわけではありません。Mask された領域を本当に計算しないかどうかは kernel によります。

$N$	双方向 $N^2$	Causal の有効接続 $N(N+1)/2$
1,024	1,048,576	524,800
4,096	16,777,216	8,390,656
32,768	1,073,741,824	536,887,296
1,000,000	1,000,000,000,000	500,000,500,000

Sequence が 10 倍になると、接続数はおよそ 100 倍になります。これが一つ目のコスト、二次の算術です。

24.1.2 中間結果の保存

一つの Head の $N\times N$ score matrix を FP16 または BF16 で実体化すると、次の容量になります。

$N$	1 Head の dense score matrix
4,096	32 MiB
8,192	128 MiB
32,768	2 GiB
131,072	32 GiB

これは一つの Head、一つの matrixの数字であり、layer 全体ではありません。全 Head を実体化すれば Head 数を掛ける必要があり、training では backward 用の状態もあります。

第21章で見た FlashAttention は、この表を必須ではなくしました。完全な score matrix を HBM に書かずに、正確な Softmax Attention を計算できます。I/O と作業用 storage は sequence length に対してほぼ線形まで下がりますが、Query–Key の二次の組を処理する算術そのものは残ります。

Full Attention の接続数、score の実体化、KV Cache という三つの異なるコスト

24.1.3 KV Cache は別の軸

第23章の KV Cache は token 数に対して線形に増えます。Full Attention の training / Prefill における pair 数は二次に増えます。別々に考えましょう。

FlashAttention：Full Attention の数学を保ち、完全な score matrix を実体化しない。
GQA/MQA：KV Head 数と Cache bytes を減らす。
Sparse Attention：許可する Query–Key 接続を減らす。
Linear Attention：Softmax 類似度の計算形式を置き換える。
Infini-attention：segment 間の履歴を有限状態に圧縮し、現在 segment 内には局所的な正確な Attention を残す。

この区別ができれば、後の complexity を混同しなくなります。

24.2 Sparse Attention：O(N) と書く前に graph を描く

輪島塗の古い修復台帳を読んでいるとします。現在の記録は、まず前後の作業記録に依存します。大きな損傷が見つかった日の見出しは global anchor になります。数年離れた修復箇所への少数の参照が、遠い記録同士の距離を縮めます。Sparse Attention は、この「誰が誰を見られるか」を graph として明示します。

Sliding Window、Global、Random の Sparse Attention 接続パターン

よく使われる三種類の edge は次の通りです。

Window / Local：各 Query は近傍のおよそ $w$ 個の Key を見る。
Global：少数の $g$ 個の位置が sequence 全体と双方向に接続する。Longformer では task に応じて global token を指定する。
Random：各位置がさらに $r$ 個のランダムな位置へ接続する。BigBird は Window、Global と組み合わせる。

各行の edge 数が固定なら、主な edge 数は次の通りです。

O\!\left(N(w+g+r)\right)

$w,g,r$ が $N$ とともに増えない場合だけ、 $O(N)$ と省略できます。Window を sequence と一緒に拡大すれば complexity も変わります。

24.2.1 Longformer と BigBird が実際に示した範囲

Longformer は Sliding Window と task 固有の Global Attention を組み合わせます。原論文は主に長文 document encoding、masked language modeling、LED encoder–decoder を扱いました。BigBird は Random 接続を追加し、こちらも主に BERT 型の長い sequence encoding task を評価しました。

BigBird 論文は、global token、depth、precision などを含む所定の graph 構成と仮定の下で、Sparse Transformer が universal approximation と Turing completeness の性質を保てることを証明します。しかし、次のことまでは証明しません。

任意の Random pattern で、どの二 token も固定 $O(1)$ layer 以内に結ばれる。
有限 depth の学習済み BigBird が常に Full Attention と同じ品質になる。
すべての task で accuracy loss が小さい。

理論上の表現力、graph の経路長、学習後の品質は別々の主張です。

24.2.2 Mask は speedup ではない

通常の $N\times N$ score matrix を計算してから不許可セルに -inf を入れれば、Sparse Attention の正しい結果は得られます。しかし dense matrix はすでに計算されています。実際に仕事を減らすには、空 block を飛ばす Block-sparse または FlexAttention 型 kernel が必要です。Block size、sparsity、shape、hardware utilization が実速度を決めます。Sparsity が少なければ、schedule の overhead が理論上の利益を消すこともあります。

Dense score matrix に Sparse mask を加える場合と structured sparse kernel の違い

24.3 Sliding Window：考え方は簡単、境界は一つずれやすい

Decoder のcausal windowだけを考えます。ここでは window=3 を「自分自身と直前の二位置まで見られる」と定義します。

q0 -> k0
q1 -> k0 k1
q2 -> k0 k1 k2
q3 ->    k1 k2 k3
...

長さ 8 の有効接続数は $8\times3=24$ ではなく、 $1+2+6\times3=21$ です。左端で三本少なくなります。

import math
import torch


def causal_window_mask(length, window, device=None):
    if length <= 0 or window <= 0:
        raise ValueError("length and window must be positive")
    positions = torch.arange(length, device=device)
    distance = positions[:, None] - positions[None, :]
    return (distance >= 0) & (distance < window)


def masked_attention(q, k, v, mask):
    """Dense teaching reference; use a structured sparse kernel for real savings."""
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.size(-1))
    scores = scores.masked_fill(~mask, float("-inf"))
    return torch.softmax(scores, dim=-1) @ v

これは mask semantics を確認する dense reference であり、高速 Sparse 実装ではありません。PyTorch の boolean mask も API 間で意味が同一ではありません。SDPA/FlexAttention と MultiheadAttention では、True が「残す」のか「遮る」のかを現在の API で確認する必要があります。

24.3.1 Layer を重ねた receptive field

上のように「現在位置を含めて window token」と定義し、dilation も global position もない場合、 $L$ layer 後の最大 causal span は次の通りです。

1+L(w-1)

論文によっては $W$ を後方への距離として定義し、およそ $LW$ と書きます。見かけの差が定義だけの場合もあります。Window の diameter、radius、現在位置を含むかどうかを混ぜないことが大切です。

24.3.2 Mistral 7B という歴史的な例

Mistral 7B v0.1 の technical report は、32 layer と $W=4096$ の Sliding Window を示し、約 131K token の理論的な情報伝播 spanを導きます。同じ config table の context_len は 8192 です。131K は、v0.1 checkpoint が 131K のすべての位置で正確な token recall を検証済みという意味ではありません。

同報告は Rolling Buffer Cache も説明しています。位置 $i$ の K/V は $i\bmod W$ に書かれ、Window 外へ出た entry は上書きされます。Window が埋まれば Cache はそれ以上増えません。Window 外の正確な K/V は消えますが、古い内容が新しい hidden state を介して上層へ間接的に影響することはあります。

Causal window が layer ごとに receptive field を広げ、Rolling Buffer KV Cache を上書きする流れ

24.4 Linear Attention は Softmax の括弧を変えただけではない

Softmax がなければ、行列の結合則により：

(QK^T)V=Q(K^TV)

左の順序は $N\times N$ 、右の順序は $d_k\times d_v$ の中間 matrix を作ります。しかし標準 Attention は次の式です。

\operatorname{softmax}(QK^T/\sqrt{d_k})V

Softmax は Query ごとに行全体を正規化するため、括弧の中を通り抜けられません。Kernelized Linear Attention は代わりに非負 feature map $\phi$ を選び、別の similarity を定義します。

\operatorname{sim}(Q_i,K_j)=\phi(Q_i)^T\phi(K_j)

Causal Attention では二つの prefix state を保ちます。

S_i=\sum_{j\le i}\phi(K_j)V_j^T,\qquad Z_i=\sum_{j\le i}\phi(K_j)

出力は：

O_i=\frac{\phi(Q_i)^TS_i}{\phi(Q_i)^TZ_i}

Causal Linear Attention が N×N Softmax matrix の代わりに prefix state S と Z を使う

論文にある具体的な feature map の一つが $\phi(x)=\operatorname{ELU}(x)+1$ です。

import torch
import torch.nn.functional as F


def phi(x):
    return F.elu(x) + 1.0


def causal_kernel_attention(q, k, v):
    """One head: q/k [N,C], v [N,M]."""
    state = q.new_zeros(q.size(1), v.size(1))
    normalizer = q.new_zeros(q.size(1))
    outputs = []
    for q_t, k_t, v_t in zip(phi(q), phi(k), v):
        state = state + torch.outer(k_t, v_t)
        normalizer = normalizer + k_t
        outputs.append((q_t @ state) / (q_t @ normalizer).clamp_min(1e-12))
    return torch.stack(outputs)

Feature dimension を $C$ 、Value dimension を $M$ とすると、主項は $O(NCM)$ です。Streaming Decode が持つ state は $CM+C$ で、履歴長に対して固定です。Full Attention より速いかどうかは $N,C,M$ 、parallelism、kernel によります。最も重要なのは、この code が Kernel Attention を正確に実装しているのであって、標準 Softmax Attention を誤差なく並べ替えたものではないことです。

本章の実行可能な test は、streaming prefix state と明示的な lower-triangular Kernel Attention を比較します。Float32/float64 の出力と勾配は一致します。

24.5 Infini-attention：固定されるのは segment 間の state

Infini-attention は入力を固定長 segment に分けます。各 layer、各 Head は二つの仕事をします。

現在 segment 内で正確な causal scaled dot-product Attention を行い、局所的な細部を保つ。
前 segment の圧縮 state から検索し、現在 segment の処理後に state を更新する。

Infini-attention が segment 内の local Attention と segment 間の M・z state を組み合わせる

一つの Head の state は単なる d_model × d_model ではありません。

M: [d_key, d_value]
z: [d_key]

古い state から検索します。

A_{mem}=\frac{\phi(Q)M_{s-1}}{\phi(Q)z_{s-1}}

Segment の出力後、その K と V で更新します。

M_s=M_{s-1}+\phi(K)^TV,\qquad z_s=z_{s-1}+\sum_t\phi(K_t)

論文は Delta update も評価していますが、ここでは直接的な Linear 版を使います。Local output $A_{dot}$ と memory output は Head ごとの一つの学習可能 scalar で混ぜます。

A=\sigma(\beta)A_{mem}+(1-\sigma(\beta))A_{dot}

import torch
import torch.nn.functional as F


def phi(x):
    return F.elu(x) + 1.0


def retrieve_memory(q, memory, normalizer):
    q_features = phi(q)
    return (q_features @ memory) / (
        q_features @ normalizer
    ).unsqueeze(-1).clamp_min(1e-12)


def infini_segment(q, k, v, memory, normalizer, gate_logit):
    """One head and one segment; q/k [N,D_k], v [N,D_v]."""
    memory_output = retrieve_memory(q, memory, normalizer)
    local_output = F.scaled_dot_product_attention(
        q[None, None], k[None, None], v[None, None],
        is_causal=True, dropout_p=0.0,
    )[0, 0]
    gate = torch.sigmoid(gate_logit)
    output = gate * memory_output + (1.0 - gate) * local_output

    k_features = phi(k)
    next_memory = memory + k_features.T @ v
    next_normalizer = normalizer + k_features.sum(dim=0)
    return output, next_memory, next_normalizer

順序が重要です。現在 segment はまず $M_{s-1},z_{s-1}$ を読み、出力を作った後に $M_s,z_s$ を書きます。そのため memory branch は現在 segment の未来位置を漏らしません。現在 segment の causality は local Attention が担当します。

24.5.1 「Infinite」の意味

Layer ごと、Head ごとの跨 segment state は、処理済み segment 数に関係なく $d_kd_v+d_k$ 個の値です。Segment length を固定 $S$ 、総 length を $T$ とすると：

Local Attention の総算術はおよそ $O(TSd)$ であり、総計算が $O(1)$ になるわけではない。
圧縮 state は $T$ に対して $O(1)$ 。
一時的な local-Attention memory は $S$ と kernel に依存する。

「Infinite」は、segment の総数に事前上限を置かずに stream を処理でき、recurrent state が bounded であるという意味です。過去の全 token を無損失に random access できるという意味ではありません。一つの matrix に多数の K/V binding が加算されるため、干渉と情報損失は圧縮の一部です。

24.5.2 論文の結果は Gemini 1.5 の architecture 公開ではない

Infini-attention 論文は、対応する training 後に、1B model で 1M length の passkey retrieval、8B model で 500K length の book summarization を報告しました。これは論文内の特定 model、task、training 条件です。

Gemini 1.5 technical report は 2024年3月、Infini-attention の最初の arXiv 版は 2024年4月に公開されました。どちらの公開 report にも Gemini 1.5 が Infini-attention を使うとは書かれていません。どちらも長い context を扱うという理由だけで同じ architecture とみなすのは推測です。

24.6 各方式を正しい軸に戻す

方式	何を変えるか	総 length $T$ の主項	過去を正確に保存するか	重要な条件
Full + FlashAttention	Kernel / I/O	$O(T^2d)$ 算術	Window 内は正確	Full score matrix を HBM に書かない
Sliding Window	接続 graph	$O(Twd)$	Window 外は間接影響のみ	固定 $w$ と構造を利用する kernel
Longformer / BigBird	Local、Global、Random graph	$O(T(w+g+r)d)$	許可した edge のみ正確	原論文の claim には task と graph の境界がある
Kernel Linear Attention	Similarity と計算順序	$O(TCM)$	$S,Z$ に圧縮	Softmax の単純な等価式ではない
Infini-attention	Segment 間の recurrent state	固定 segment なら $T$ に線形	Segment 間は圧縮	適応 training が必要、local Attention は残る

数学的に組み合わせられる方式はありますが、すべての framework が効率的な fused kernel を持つとは限りません。FlashAttention、GQA、Block-sparse mask、Rolling Cache、compressed memory は別々の layer に作用します。追加するたびに correctness、quality、peak memory、実測 throughput を検証します。

24.7 章のまとめ

Full Attention には二次の pairwise work がある。実体化した $N^2$ score storage と線形 KV Cache は別のコスト。
Sparse Attention が仕事を減らすのは、空 block を飛ばす kernel と組み合わせたときだけ。Dense matrix + mask は Sparse な結果を作るだけ。
Longformer と BigBird が線形になるのは Window、Global、Random の数が固定の場合。理論上の表現力は任意 task の品質保証ではない。
Linear Attention は feature map で新しい Kernel similarity を定義する。標準 Softmax Attention の括弧を変えただけではない。
Infini-attention の $M,z$ state は総履歴に対して bounded だが、総計算は segment 数とともに増え、segment 間の履歴は lossy。
Gemini 1.5 が Infini-attention を使うという公開証拠はない。

参考文献

Longformer: The Long-Document Transformer（Beltagy et al., 2020）
Big Bird: Transformers for Longer Sequences（Zaheer et al., 2020）
Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention（Katharopoulos et al., 2020）
Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention（Munkhdalai et al., 2024）
Mistral 7B（Jiang et al., 2023）
Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context（Gemini Team, 2024）
PyTorch FlexAttention documentation

次の章へ

この章では、モデルが相対位置を知っていることを暗黙に前提としていました。Window には進む方向があり、causality には順序があり、segment にも position scheme が必要です。第25章ではその基礎へ戻り、位置 encoding が absolute position から relative distance、rotation、linear bias へどう発展したかを追います。