行列積のコスト — AIの計算はほぼここに帰着する

なぜGEMMが全てなのか。O(n³)の内訳、メモリ帯域と演算強度、GPUが速い理由、タイル化の直感。最後に、AIモデルの学習・推論に必要なFLOPsを自分の手で見積もれるところまで持っていきます。

比喩: 加工機とトラック

工場に、猛烈に速い加工機を1台入れたとします。1秒で1000個の部品を削れる。ところが材料を運ぶトラックは1秒に10個しか届けられません。この工場の生産量は、加工機の性能ではなくトラックの本数で決まります。加工機は99%の時間、手を止めて待っている。

現代の計算機はこの工場そのものです。演算器は猛烈に速く、メモリからデータを運ぶ帯域はそれに追いつきません。だから「何回計算するか」だけを数えても実行時間は当たらず、何バイト運ぶかを同時に見る必要があります。

そしてAIの計算時間は、その大半が1種類の演算に集中しています。行列積です。

なぜ行列積が全てなのか

Transformerの中身を分解すると、出てくるのはほとんど行列積です。Q・K・Vへの射影、注意スコアの計算、注意の重み付き和、出力射影、そしてFFNの2枚の線形層。畳み込みも、im2colやそれに類する変形で行列積に落として実行されるのが普通です。

だからライブラリの世界では、この演算に GEMM（General Matrix Multiply, 一般行列積）という専用の名前が付き、何十年も最適化され続けてきました。「AIを速くする」という工学課題の実体は、かなりの部分が「GEMMを速くする」ことです。

$O(n^3)$ の内訳

$m \times k$ の行列 $A$ と $k \times n$ の行列 $B$ を掛けて $m \times n$ の $C$ を作ります。出力の1マスの定義はこれです。

「$C$ の $i$ 行 $j$ 列は、$A$ の $i$ 行目と $B$ の $j$ 列目を端から掛けて足したもの」——つまり内積1回です。$C$ にはマスが $mn$ 個あり、1マスにつき掛け算 $k$ 回と足し算がほぼ $k$ 回。掛け算と足し算を1回ずつ数えると、必要な演算回数は

$$\text{FLOPs} = 2\,m\,k\,n$$

つまりこの式は「出力のマスの数 $mn$ × 1マスを埋めるのに要る積和の回数 $k$ × 積和1回で2演算」を並べて掛けただけのものです。難しいことは何も起きていません。

正方行列（$m = k = n$）なら $2n^3$ です。$n$ が2倍になれば計算量は8倍。これが $O(n^3)$ の中身で、AIモデルの「計算量」と呼ばれるものの正体でもあります。

数を入れると規模感が出ます。$n = 4096$ の正方行列積は約 $1.4 \times 10^{11}$ FLOPs——たった1回の掛け算に1000億回超の演算が要る。Transformerの層はこれを何十枚も重ね、それをトークンごと、層ごとに繰り返します。

行列積の原子が内積であることは、手で動かして掴んでおく価値があります。

ここまでが演算回数の話です。しかし実行時間は $2n^3$ だけでは決まりません。冒頭のトラックの話——運ぶ側の制約が、ここから効いてきます。

演算強度: 1バイト運ぶごとに何回計算できるか

冒頭の工場の話を式にします。演算強度（arithmetic intensity）を次で定義します。

つまり、分子が「どれだけ働いたか」、分母が「そのために何バイト運んだか」。$I$ はデータの燃費だと思ってください。「運んだデータ1バイトあたり、何回の演算に使い回せたか」という比率です。これが小さい演算は、演算器がいくら速くても帯域で頭打ちになります（メモリ律速）。大きければ演算器の性能をそのまま使えます（演算律速）。実際に出る性能はおおよそ次で押さえられます。

$$\text{性能} \le \min\big(\text{演算器のピーク性能},\ \ I \times \text{メモリ帯域}\big)$$

つまり、実際に出る速さは「演算器が出せる上限」と「運べるデータで許される上限」の低いほうで頭打ちになる、ということです。前者は機械の限界、後者は「帯域で運べる量 × 1バイトあたりの仕事量」。この2本の直線で性能の上限を描いたものがルーフラインモデルです。

行列積で計算してみます。$n \times n$ の正方行列なら、運ぶ要素は $A, B, C$ の3枚で $3n^2$ 個。fp32（1要素4バイト）なら $12n^2$ バイトです。FLOPsは $2n^3$ なので

$$I = \frac{2n^3}{12n^2} = \frac{n}{6}$$

要するに、行列積の演算強度は「一辺の長さを6で割った数」でしかありません。$n$ に比例して増えます。行列が大きいほど、1回運んだデータを何度も使い回せる——これが行列積の際立った性質です。$n = 4096$ なら $I \approx 680$。演算律速の側にしっかり入ります。

対照的に、行列とベクトルの積（GEMV）は $2n^2$ FLOPsに対して運ぶのが約 $4n^2$ バイト、つまり $I \approx 0.5$ で固定です。行列の各要素を1回しか使わないから使い回しようがない。LLMのトークン生成が1件ずつだと遅く、バッチをまとめると劇的に効率が上がるのは、GEMVがGEMMに変わって $I$ が上がるからです。

なぜGPUが速いのか

理由は2つに整理できます。演算器をとにかく多く並べたこと、そして行列積の形で命令を切ったことです。

CPUは1つの命令をなるべく速く終わらせる設計で、分岐予測や投機実行に多くの面積を割きます。GPUは逆に、同じ演算を大量のデータに一斉に当てる前提で、制御回路を削って演算器を敷き詰めました。さらに近年のアクセラレータは、小さな行列同士の積和を1命令で実行する専用ユニットを持ちます。$I$ が高い演算だけを速くする、という割り切りです。

ここで効いてくるのが $I$ です。演算器を1000倍並べても、$I$ が小さい演算は帯域で止まるので速くなりません。GPUが速いのは行列積が高い演算強度を持つからであって、あらゆる計算が速くなるわけではない。この非対称性は、実装の判断に直結します。

タイル化: 大きな行列を、キャッシュに載る塊に切る

$I = n/6$ は「$A$ も $B$ も $C$ も全部が高速メモリに載っている」理想の値でした。現実には $n = 4096$ の行列はキャッシュに入りません。素朴な三重ループは、内側で $B$ の列を舐めるたびに主記憶へ取りに行き、実効的な $I$ は 0.25 程度まで落ちます。

対策がタイル化（ブロッキング）です。行列を $T \times T$ の小さなブロックに切り、「$A$ のブロック・$B$ のブロック・$C$ のブロック」の3枚が高速メモリに同時に載るようにする。載せてしまえば、そのブロック内では $T^3$ 回の演算を $3T^2$ 個のデータで賄えます。全体で見ると

$$I \approx \frac{T}{6}$$

つまり、実効的な燃費を決めているのは行列全体の一辺 $n$ ではなく、高速メモリに載せきれた塊の一辺 $T$ のほうだ、ということです。$T$ が大きいほど良く、上限はキャッシュ（GPUなら共有メモリ）の容量です。BLASライブラリが何十年も磨いてきたのは、突き詰めればこの切り方と、レジスタまで含めた多段のタイル設計です。アルゴリズムは同じ $O(n^3)$ のまま、データの動かし方だけで数倍から一桁の差がつく——計算量の話で触れた「オーダーの外側」が最も露骨に効く領域です。

なお $O(n^3)$ 自体を下げる方法もあります。Strassenのアルゴリズムは分割統治で $O(n^{2.807})$ を達成し、非常に大きな行列では実際に使われることがあります（数値的な安定性は素朴な方法より劣ります）。理論上の指数はさらに2.4を切るところまで下げられていますが、定数が巨大すぎて実用にならない「銀河系アルゴリズム」です。

FLOPsを自分で見積もる

ここが到達点です。まず道具を用意します。

def gemm_cost(m, k, n, bytes_per_elem=4):
    flops = 2 * m * k * n                          # 積和1回 = 2 FLOP
    moved = (m * k + k * n + m * n) * bytes_per_elem
    return flops, flops / moved                    # (FLOPs, 演算強度)

print(gemm_cost(4096, 4096, 4096))   # → 演算強度 682.7 [FLOP/byte]
print(gemm_cost(4096, 4096, 1))      # → 演算強度   0.5 [FLOP/byte]

同じ行列 $A$ を使っても、相手が行列かベクトルかで演算強度が1000倍以上違う。バッチ化の効き目がそのまま数字に出ます。

モデル全体は、パラメータ数 $N$ から概算できます。パラメータ1個は「掛けて足す」で2 FLOPを生むので、推論の順伝播はトークン1個あたり約 $2N$ FLOPs。学習では逆伝播が入力側と重み側の2つの勾配を計算するため、順伝播の約2倍かかり、合計でトークン1個あたり約 $6N$ FLOPsです（誤差逆伝播）。したがって $D$ トークンで学習する総計算量は

つまり「パラメータ1個とトークン1個の組み合わせごとに、だいたい6回の演算」。学習の総コストは、モデルの大きさ $N$ と読ませた文章量 $D$ の掛け算でほぼ決まる、ということです。これは注意機構の $s^2$ 項や正規化を無視した近似で、$d_{\text{model}} > s/12$、つまり系列長 $s$ が $d_{\text{model}}$ の約12倍を下回る範囲で使えます。長文脈では注意の項が無視できなくなる、と覚えておいてください。

数字を入れます。$N = 7 \times 10^9$、$D = 10^{12}$ トークンなら $C \approx 4.2 \times 10^{22}$ FLOPs。もし持続的に $10^{15}$ FLOP/s（1 PFLOP/s）を出せる環境なら $4.2 \times 10^{7}$ 秒、約1.3年かかります。学習に数百〜数千台の加速器を並べる理由が、この一行で腑に落ちるはずです。

実務で効く4つ

1. バッチサイズを疑う 推論が遅いとき、まず $I$ を見ます。バッチ1の生成はGEMVでメモリ律速なので、演算器を増やしても速くなりません。まとめて流すのが唯一の効く手です。

2. 「実効FLOPs / ピークFLOPs」を見る 実測時間から逆算した実効性能をピーク性能で割ると、どれだけ演算器を遊ばせているかが出ます。この比が低いなら、原因は計算量ではなくデータの動かし方です。

3. 精度を下げるのは帯域を買う行為 fp32をfp16やint8にすると、演算が速くなるだけでなく運ぶバイト数が半分・4分の1になるので $I$ が上がります。量子化がメモリ律速の場面で特に効くのはこのためです。

4. 自分でGEMMを書かない タイル化・プリフェッチ・レジスタ割り当てを含む最適化は、既存のBLAS実装に何年ぶんもの蓄積があります。書くべきなのは、GEMMを呼ぶ形のほうです。

まとめ

• AIの計算はほぼGEMMに帰着する。$m \times k$ と $k \times n$ の積で $2mkn$ FLOPs、正方なら $2n^3$
• 実行時間は演算回数だけでは決まらない。演算強度 $I$ = FLOPs / 運ぶバイト数 が帯域律速か演算律速かを分ける
• 行列積は $I \approx n/6$ で大きいほど有利、行列ベクトル積は $I \approx 0.5$ で固定。バッチ化が効くのはこの差
• タイル化は $I \approx T/6$ をキャッシュ容量まで引き上げる工夫。$O(n^3)$ のまま実行時間が桁で変わる
• 概算式は推論 $2N$ / トークン、学習 $6ND$。これで必要な計算資源が自分で見積もれる

次は、この計算を多数の演算器へどう分けるか——並列化の限界と分散学習の通信を扱います。