8. 計算効率化と高速化技術

二乗の壁（$O(N^2)$問題）とメモリボトルネック

Transformer は驚異的な言語理解能力を持つ一方で、致命的な計算負荷の課題を抱えています。 系列の長さ（トークン数）を $N$ と置いたとき、すべての単語同士の組み合わせを計算するため、Self-Attention の計算量とメモリ消費量は $N$ の 2 乗（$O(N^2)$）に比例して増大します。

トークン数 N = 1,000      -> アテンション行列: 100 万要素
トークン数 N = 10,000     -> アテンション行列: 1 億要素
トークン数 N = 100,000    -> アテンション行列: 100 億要素（メモリ枯渇）

書籍 1 冊分の文章や長い対話履歴、高解像度の画像を処理しようとすると、GPU のメモリ（VRAM）は一瞬でパンクしてしまいます。 この「二乗の壁」を乗り越え、実用的な速度とコストで動かすために、さまざまな効率化技術が生み出されてきました。

推論の生命線: KVキャッシュ

文章を 1 単語ずつ生成する対話型 AI では、生成のたびに文頭からの全単語をニューラルネットワークに入力し直すと、過去の単語の計算が何重にも重複して無駄になります。

ここで活躍するのが「KV キャッシュ（Key-Value Cache）」です。 自己回帰生成では、過去に登場した単語の Key と Value のベクトルは以降の計算でも変化しません。 そのため、一度計算した過去の Key と Value を GPU メモリ上に保存（キャッシュ）しておきます。

graph LR
    subgraph 従来（キャッシュなし）
        All["過去の全単語 + 新単語"] --> Calc["全単語の Q, K, V を毎回再計算"]
    end
    subgraph KVキャッシュ利用
        NewToken["新単語 (1トークン)"] --> CalcNew["新単語の Q, K, V のみ計算"]
        Cache["過去の K, V キャッシュ"] -.->|連結| Attn["Attention 計算"]
        CalcNew --> Attn
    end

新しい単語を 1 つ生成するときは、その 1 単語の Query だけを計算し、キャッシュされた過去の Key・Value と突き合わせるだけで済みます。 これにより、1 トークン生成あたりの計算量が劇的に削減されます。

メモリ消費を削る: MQAとGQA

KV キャッシュは推論を高速化する一方で、利用者が増えたり文章が長くなったりすると、キャッシュ自体が膨大な GPU メモリを消費するという新たな壁を生み出しました。 このメモリ消費を劇的に削減するために考案されたのが、Multi-Query Attention（MQA）と Grouped-Query Attention（GQA）です。

【MHA (従来のマルチヘッド)】
Head 1: Q1, K1, V1
Head 2: Q2, K2, V2
Head 3: Q3, K3, V3
Head 4: Q4, K4, V4
(全ヘッドが独立した K と V を保持するため、メモリ消費大)

【GQA (グループド・クエリアテンション)】
Group 1: Q1, Q2  ──> 共通の K1, V1 を参照
Group 2: Q3, Q4  ──> 共通の K2, V2 を参照
(グループ内で K と V を共有し、メモリ消費を大幅削減)

【MQA (マルチクエリアテンション)】
Q1, Q2, Q3, Q4   ──> たった 1 組の K, V を全員で共有

▸ Grouped-Query Attention（GQA）の勝利

MQA はメモリを最小化できますが、表現力が落ちてモデルの知能が低下しやすいという問題がありました。 そこで登場したのが、両者の中間に位置する「GQA」です。 ヘッドをいくつかのグループに分け、グループ内で Key と Value を共有します。 Meta の LLaMA 2（70B）や Llama 3、Mistral など、近年の最先端 LLM の大半が GQA を採用しており、モデルの賢さを維持したまま、推論時のメモリ消費量を数分の一に圧縮することに成功しています。

ハードウェアの真実とFlashAttention

2022 年、スタンフォード大学の Tri Dao らは「FlashAttention」を発表し、深層学習業界に革命を起こしました。 従来の高速化研究の多くは「数式をいかに近似して計算量を減らすか」に注目していましたが、FlashAttention は「GPU のハードウェア構造」に徹底して着目しました。

▸ メモリ帯域幅のボトルネック（Memory-Bound）

現代の GPU は計算速度（FLOPS）が極めて速い反面、大容量メモリ（高帯域幅メモリ: HBM）とプロセッサチップ上の超高速メモリ（SRAM）との間でデータを読み書きする速度（転送帯域）がボトルネックになっています。 従来の Attention 実装では、$N \times N$ の巨大なアテンション行列をいちいち HBM に書き出し、次の処理で再び読み込んでいました。 計算そのものよりも、メモリへの読み書き待ちで GPU が遊んでいたのです。

▸ タイリング（Tiling）と Online Softmax

FlashAttention は、データを小さなブロック（タイル）に分割し、超高速なオンチップ SRAM の中だけで Attention の計算を完結させます。 巨大な $N \times N$ 行列を HBM に一切書き出さないため、メモリの往復転送が劇的に減少します。 数式の計算量自体は $O(N^2)$ のまま一切の近似（精度の劣化）を行わないにもかかわらず、実行速度を 2〜4 倍に加速し、メモリ使用量を大幅に削減しました。 現在、FlashAttention（および改良版の FlashAttention-2 / 3）は、あらゆる Transformer の学習と推論におけるデファクトスタンダードとなっています。

次章では、これまでの技術進化を踏まえ、Transformer が直面している根本的な課題と、ポスト Transformer を目指す次世代アーキテクチャの展望を総括します。