行列積のコスト — 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バイトあたり、何回の演算に使い回せたか」という比率です。これが小さい演算は、演算器がいくら速くても帯域で頭打ちになります（メモリ律速）。大きければ演算器の性能をそのまま使えます（演算律速）。実際に出る性能はおおよそ次で押さえられます。