第7章 — 効率と Transformer の派生
『LLM Primer II: 数学で読み解く言語モデル』を章ごとに紹介していくシリーズ、第7回です。第II部を閉じる章 — 数学がハードウェアと出会う場面、と言ってもよいかもしれません。
この章の問い
第4章で導出した Attention の式には、短い系列では問題にならなかったけれど、長い系列では一気に重くのしかかってくる性質があります。計算量もメモリも、系列長の2乗で伸びるのです。文脈を2倍にすれば、Attention 層は4倍のものを要求してまいります。
200トークンのチャットなら、ほとんど気づきません。20万トークンの文脈窓となると、すぐに気づきます — おそらく1回の順伝播を終える前に GPU のメモリが尽きてしまいます。
第7章は、その壁の章です。コストについて数学が何を語っているのか。現代のシステムは、その壁をどう押し下げているのか。そして Attention の式そのものを賢く書き換えた「Transformer の派生」たちには、どんな顔ぶれがあるのか。控えめに見渡してまいります。
7.1 Attention の計算量
7.1 節では、コストを丁寧に計算します。長さ n、埋め込み次元 d の系列に対して、QKT の計算は O(n²·d)。softmax がさらに O(n²) を上乗せします。V との加重和でまた O(n²·d)。合計 O(n²·d) は、n が数百のうちは穏やかですが、数千で痛み始め、それより先は手に余ります。
実務では、メモリのほうがもっと辛いことが多いです。中間生成物である n × n の Attention 行列が、softmax が終わるまで GPU のメモリに居座らなければなりません。長文脈では、計算量がボトルネックになる前に、この行列がボトルネックになります。
7.2 GPU メモリとスループットの算数
7.2 節では、いったん Attention の式を脇に置き、ハードウェアの作りを丁寧にほどきます。GPU には、速くて小さなオンチップメモリ(SRAM)、ずっと大きいけれど少し遅い高帯域メモリ(HBM)、そしてさらに遅い CPU / ディスクという、三段のメモリ階層があります。性能は、これらの階層のあいだでデータを動かす回数で大方決まってまいります。
節のなかでは、計算量とメモリアクセスの比 — 算術強度(arithmetic intensity) — と、与えられたアルゴリズムと GPU で「計算律速」か「メモリ帯域律速」かを予言してくれるルーフラインモデルを紹介します。ここで効くのは次の結論です。素朴な Attention は、現代の文脈長において、計算律速になる前にメモリ律速になっています。
この数学的なお膳立てが、次の節を可能にしてくれます。
7.3 FlashAttention とメモリを意識した計算
FlashAttention は、近年の Transformer 史のなかでも、もっとも実務に効いた工夫のひとつです。7.3 節では、ここを丁寧に導出します。
核心は、小さくて美しい一手です。標準の Attention は、softmax を取る前に n × n の Attention 行列を HBM 上に丸ごと展開します。FlashAttention は、計算順を組み替えて、softmax をブロックごとに、SRAM に収まる小さなタイルで進めます。途中の統計量を運びながら更新していく工夫を入れて、結果は標準の softmax と数学的に同じにそろえる、というのが仕掛けです。
出力は、浮動小数点の丸めを除けば、ビット単位で同じ数字です。けれども HBM へのアクセス回数は大きく減ります。実機では、モデルを変えずに長文脈の Attention が 2 〜 4 倍ほど速くなる、という場面がよくあります。章のなかでは、これを支える再帰の式と、なぜきちんと働くのかを示します。
7.4 マルチクエリ Attention とゲート付きアーキテクチャ
7.4 節では、「計算の仕方を変える」から「Attention そのものを変える」へ、足を踏み出します。
マルチクエリ Attention(MQA)は、クエリだけマルチヘッドのまま残し、キーとバリューを全ヘッドで1組だけ共有します。これにより、推論の作業領域である KV キャッシュが、ヘッド数ぶんだけ小さくなり、生成時のレイテンシとメモリコストがそのまま下がります。グループクエリ Attention(GQA)は、標準のマルチヘッドと MQA の中間に位置し、ヘッドをグループ分けしてキー・バリューを共有します。
節の後半では、ゲート付きアーキテクチャ — Attention や FFN の残差への寄与を、学習されたゲートで動的に調整する仕組み — にも触れています。どの層をどれだけ使うかを、モデル自身に少しだけ決めさせる工夫、というふうに捉えていただけるかと思います。
7.5 低ランク・近似 Attention
7.5 節は、近似の世界を広めに見渡して章を閉じます。Attention が O(n²) で、それが問題なのだとしたら、それを O(n·log n) や O(n) でほどよく近似できる分解には、原理的にすべて出番があります。
章のなかでは、主だった系譜を順番に並べます。線形 Attention は、softmax Attention をカーネル計算として書き直し、組み替えて n に線形にします。Performer は、ランダム特徴量で softmax カーネルを近似します。スパース Attention は、各トークンが見にいける相手を、あらかじめ決めた小さなパターンに絞ります。スライディング窓 Attention は、近隣のみに目を限ります。それぞれが、正確さの一部とコストを交換しています。章のなかでは、その代金がいくらで、どの場面でその近似が向く・向かないのかを、まとめてあります。
この章を踏まえて
第7章を読み終えるころには、この分野が戦ってきたコストの輪郭と、定番の対処法の見取り図が、手元に残っているはずです。ここで第II部は閉じます。ここから先は、別種の数学 — モデルがどのように学ぶのかを語る最適化の理論と、何ができるようになるのかを予言してくれるスケーリング則 — に話が移ってまいります。
次回 — 第8章: モデルはどう学ぶのか。 第III部の最初の章です。過剰パラメータのモデルがなぜ汎化するのか(本来は「するはずがなかった」のです)、勾配法の暗黙のバイアス、学習が始まる前から能力を予言する経験則、そして残された未解決の数学的な問い — 控えめに紹介してまいります。
