行列演算(tensor)

実装: llm/tensor/ / 実行: go test ./llm/tensor/

ニューラルネットの計算は、行列積と、その周りの数本の関数でできている。numpy に頼らず Go で書くと、多次元に見えても、メモリの上では1本の連続した配列でしかないことが手触りで分かる。しかも計算量が同じでも、回る順を変えるだけで実測 約1.5 倍の差が出る。非線形を挟まないと何層重ねても1層に潰れることも、テストで確かめる。

この章で作るもの

LLM は、BPEトークナイザで作ったトークン ID の列を受け取って、「次に来るトークンの候補ごとの点数」を返す。この入口から出口へ向かう一方向の計算を forward pass と呼び、返ってくる点数の列を logits と呼ぶ。この編ではその forward pass を下から順に自作していく。まず作るのは一番下の層、行列演算だ。numpy も PyTorch も使わない。

書くのは4本だけだ。行列積(MatMul)、softmax、LayerNorm、GELU。この4本で、この先の章で作る attention も mini-GPT も組める。

数が少ないのには理由がある。ニューラルネットの計算のほとんどは「入力ベクトルに重み行列を掛ける」の繰り返しで、それを層の数だけ並べたものだからだ。中心にあるのは行列積1つで、残りはその周りを整える係になる。

  (3, 4) の行列               メモリ上の並び(行優先)

    ┌────┬────┬────┬────┐     [ 00 01 02 03 10 11 12 13 20 21 22 23 ]
    │ 00 │ 01 │ 02 │ 03 │       └───行0───┘└───行1───┘└───行2───┘
    ├────┼────┼────┼────┤
    │ 10 │ 11 │ 12 │ 13 │     (r, c) は Data[r*Cols + c]
    ├────┼────┼────┼────┤
    │ 20 │ 21 │ 22 │ 23 │     行方向は隣どうし。列方向は Cols 個ぶん飛ぶ
    └────┴────┴────┴────┘


  行列積 C = A·B を、どの順で回るか

    ① i → j → p(教科書どおり)     C の1マスごとに内積を取る
         内側の p が B を列方向にたどる → 1回進むたび n 個ぶん飛ぶ

    ② i → p → j(こちらを採る)     A の1マスを取り出して B の1行に配る
         内側の j が B を行方向にたどる → 隣どうしを読む

    計算量は同じ。読む場所が連続するかどうかだけが違う

2次元に見えても、メモリの上では1本の配列になる。どの順で回るかで、読む場所が連続するかどうかが変わる

順に見ていく。

1. 中心は行列積1つ: attention も全結合層も、突き詰めると同じ形になる
2. 多次元に見えても、メモリは1本: どの順で回るかで速さが変わる。実測で約 1.5 倍
3. 非線形が無ければ何層重ねても1層: 行列積だけを重ねても、1つの行列に潰れてしまう

① 中心は行列積1つ

まず表現を決める。行優先の1次元配列に、行数と列数を添えるだけになる:

// Tensor は行優先(row-major)で並べた2次元行列。
// Data[r*Cols + c] が (r, c) 成分。
type Tensor struct {
	Rows, Cols int
	Data       []float32
}

// New はゼロ埋めの (rows, cols) 行列を作る。
func New(rows, cols int) *Tensor {
	return &Tensor{Rows: rows, Cols: cols, Data: make([]float32, rows*cols)}
}

// FromRows は2次元スライスから行列を作る。
func FromRows(rows [][]float32) *Tensor {
	r := len(rows)
	c := len(rows[0])
	t := New(r, c)
	for i := range rows {
		copy(t.Data[i*c:], rows[i])
	}
	return t
}

// At は (r, c) 成分を返す。
func (t *Tensor) At(r, c int) float32 { return t.Data[r*t.Cols+c] }

// Set は (r, c) 成分を書く。
func (t *Tensor) Set(r, c int, v float32) { t.Data[r*t.Cols+c] = v }

Data[r*Cols + c] で (r, c) にたどり着く。多次元に見える計算も、メモリの上では1本の連続した配列でしかない。この「連続したメモリ + 添字計算」は、ディスクとページのページと同じ考え方になる。

その上に行列積を書く:

// MatMul は行列積 a·b。a が (m, k)、b が (k, n) なら結果は (m, n)。
// Transformer の計算のほとんどはこれ。attention も全結合層も、突き詰めれば MatMul。
func MatMul(a, b *Tensor) *Tensor {
	if a.Cols != b.Rows {
		panic("tensor: MatMul shape mismatch")
	}
	m, k, n := a.Rows, a.Cols, b.Cols
	out := New(m, n)
	for i := 0; i < m; i++ {
		for p := 0; p < k; p++ {
			aip := a.At(i, p)
			for j := 0; j < n; j++ {
				out.Data[i*n+j] += aip * b.At(p, j)
			}
		}
	}
	return out
}

結果の1マスは「左の行列の行」と「右の行列の列」の内積になる。この単純な計算の積み重ねが、GPT の何十億回の演算の正体だ。

形が合わなければ掛けられない(A の列数 = B の行数)。実装では合わなければ止める。この形の取り違えは、実際の LLM 実装でいちばん多い間違いだ。

試す: 結果 C のマスにマウスを乗せると、掛け合わされる A の行(青)と B の列(緑)が光り、その内積が値になっているのが見える。

デモ行列積 A · B2×3 · 3×2 = 2×2

A

123456

×

B

123456

=

C

22284964

C[0][0] = 1×1 + 2×3 + 3×5 = 22

結果の1マスは「A の行」と「B の列」の内積。C のマスにマウスを乗せると、 掛け合わされる行(青)と列(緑)が光る。Transformer の計算はほぼこれの積み重ね。

② 多次元に見えても、メモリは1本

上の実装は i → p → j の順で回っている。教科書に載っているのは i → j → p、つまり「結果の1マスごとに内積を取る」形のほうだ。書いてみると、こうなる:

// MatMulByDot は結果の1マスごとに内積を取る、教科書どおりの並べ方。
//
// 結果は MatMul と同じになる。違うのは内側の回りかただけで、
// こちらは b を列方向にたどるので、1回進むたびに Cols 個ぶん離れた場所を読む。
// 連続した場所を読むほうが速いので、同じ計算量でも差が出る。
func MatMulByDot(a, b *Tensor) *Tensor {
	if a.Cols != b.Rows {
		panic("tensor: MatMulByDot shape mismatch")
	}
	m, k, n := a.Rows, a.Cols, b.Cols
	out := New(m, n)
	for i := 0; i < m; i++ {
		for j := 0; j < n; j++ {
			var s float32
			for p := 0; p < k; p++ {
				s += a.Data[i*k+p] * b.Data[p*n+j]
			}
			out.Data[i*n+j] = s
		}
	}
	return out
}

計算量はまったく同じになる。掛け算の回数も足し算の回数も変わらない。違うのは、内側の回りかたで読む場所が連続するかどうかだけだ。

内積を取る形は、B を列方向にたどる。1回進むたびに Cols 個ぶん離れた場所へ飛ぶので、読み込んだ塊のうち1つしか使わずに次へ行くことになる。i → p → j なら B を行方向にたどるので、読み込んだ塊を端から使い切れる。

256×256 の行列積で測るとこうなった(Apple M4 Pro、Go、float32):

回りかた	1回あたり
i → j → p(内積)	12.63 ms
i → p → j(行に配る)	8.16 ms

1.55 倍。同じ計算を、同じ回数やって、これだけ違う。計り直すと数%は動くが、大小が入れ替わることはない。テストで、2つの結果が一致すること(丸めの誤差を除いて)を固定した。

実物の行列積が BLAS や GPU で桁違いに速いのは、この延長線上にある。ブロックに切って、キャッシュに乗る大きさで回し、複数の掛け算を1命令で束ねる。アルゴリズムは同じで、メモリの触り方だけが違う。

③ 非線形が無ければ何層重ねても1層

残りの3本を書く。softmax は行ごとに確率へ:

// SoftmaxRows は各行を確率分布に変換する(行ごとに独立)。
// llm-sampling 編の softmax と同じ。max 引きで overflow を防ぐ。
// attention で「どのトークンにどれだけ注目するか」を確率にするのに使う。
func SoftmaxRows(x *Tensor) *Tensor {
	out := New(x.Rows, x.Cols)
	for r := 0; r < x.Rows; r++ {
		maxv := float32(math.Inf(-1))
		for c := 0; c < x.Cols; c++ {
			if x.At(r, c) > maxv {
				maxv = x.At(r, c)
			}
		}
		var sum float32
		for c := 0; c < x.Cols; c++ {
			e := float32(math.Exp(float64(x.At(r, c) - maxv)))
			out.Set(r, c, e)
			sum += e
		}
		for c := 0; c < x.Cols; c++ {
			out.Set(r, c, out.At(r, c)/sum)
		}
	}
	return out
}

softmax は、点数の並びを合計 1 の確率に変える関数だ。最大値を引いてから exp を取るのは値が大きいときの桁あふれを避けるためで、全体から同じ数を引いても結果の確率は変わらない。LLM Sampling で書いたものと同じで、数値安定化の手当てまで同じになる。

違うのは使いどころだ。あちらは logits を「次のトークンを選ぶ確率」に変えていた。次章の attention では、これが「どのトークンにどれだけ注目するか」の重みになる。同じ道具が、モデルの中と出口の両方に出てくる。

LayerNorm と GELU:

// LayerNorm は各行を平均0・分散1に正規化する(行ごと)。
// Transformer の各層の入口で使い、値の大きさを揃えて学習・推論を安定させる。
// (実物はこの後に学習された gain/bias を掛けるが、ここでは正規化のみ。)
func LayerNorm(x *Tensor, eps float32) *Tensor {
	out := New(x.Rows, x.Cols)
	n := float32(x.Cols)
	for r := 0; r < x.Rows; r++ {
		var mean float32
		for c := 0; c < x.Cols; c++ {
			mean += x.At(r, c)
		}
		mean /= n
		var variance float32
		for c := 0; c < x.Cols; c++ {
			d := x.At(r, c) - mean
			variance += d * d
		}
		variance /= n
		std := float32(math.Sqrt(float64(variance) + float64(eps)))
		for c := 0; c < x.Cols; c++ {
			out.Set(r, c, (x.At(r, c)-mean)/std)
		}
	}
	return out
}

// GELU は活性化関数。ReLU の滑らかな版で、GPT や BERT が使う。
// tanh 近似式: 0.5x(1 + tanh(√(2/π)(x + 0.044715x³)))。
func GELU(x *Tensor) *Tensor {
	out := New(x.Rows, x.Cols)
	const c = 0.7978845608 // √(2/π)
	for i, v := range x.Data {
		v64 := float64(v)
		inner := c * (v64 + 0.044715*v64*v64*v64)
		out.Data[i] = float32(0.5 * v64 * (1 + math.Tanh(inner)))
	}
	return out
}

LayerNorm は各行を平均 0・分散 1 に揃える。層を深く重ねると値が発散しがちなのを、各層の入口で揃えて抑える。

GELU は、入力を素通しさせずに負の側をなめらかに押しつぶす関数だ。負を 0 で切り落とす ReLU の、角を丸めた版だと思えばよい。この「まっすぐでない」ことが、この章でいちばん大事な役をしている。これが無いと、何層重ねても1層と同じになる。

行列積は結合的なので、(x·W₁)·W₂ と x·(W₁·W₂) は同じ値になる。つまり2層ぶんの重みは、前もって1つの行列にまとめられる。層を増やした意味が消えてしまう。テストで、この2つが一致すること、そして間に GELU を挟むと一致しなくなることを、両方固定した。

深さが効くのは、線形と非線形が交互に来るからだ。行列積が情報を混ぜ、非線形が折り曲げる。どちらか片方だけでは、いくら重ねても表せる形が広がらない。

設計の観点

• 数を絞る: 部品が4本なら、どこで何が起きているかを全部追える
• 形の管理を最初に決める: 行優先の1次元配列と決めてしまえば、添字の話は1か所で済む
• 計算量が同じでも速さは同じではない: 読む場所が連続するかどうかが効く
• 速くする努力を1点に集める: ほとんどが行列積なら、そこだけ速くすれば全体が速くなる
• 重ねる意味を確かめる: 線形だけを重ねても表せる形は広がらない。非線形が入って初めて深さが効く
• 同じ道具を使い回す: softmax は入口(注目度)と出口(次のトークン)の両方で働く

対照と実例

やり方	何を変えるか	速さの出どころ
素朴な3重ループ(内積)	何も	—
回る順を変える	読む場所の連続性	実測 約1.5 倍
ブロックに切る	キャッシュに乗る大きさで回す	段の数だけ効く
複数を1命令で(SIMD)	1命令あたりの掛け算の数	幅のぶん
GPU	同時に動く演算器の数	並列度のぶん
量子化	扱う数の型	帯域と整数演算(量子化)

裏どり:

• BLAS / GEMM: 行列積は GEMM という1つの関数に集約されていて、各社が自社の CPU 向けに徹底的に最適化している。ライブラリを差し替えるだけで速度が変わるのは、ここが効くから
• キャッシュのブロック化: 行列を小さな塊に切って回すと、塊がキャッシュに収まるあいだ再利用できる。この章の「回る順」は、その最も単純な段になる
• llm.c: GPT-2 を C で書く。この教科書の Go 版の精神的な元
• PyTorch の stride: N 次元テンソルも、内部は1本の配列と「各次元で何個ぶん飛ぶか」の組で表す。transpose がデータを動かさずに済むのはこのため
• GELU の由来: Hendrycks, Gimpel(2016)。ReLU の滑らかな版で、GPT と BERT が採った

簡略化したこと

• 2次元だけ: 実物はバッチ・ヘッド・系列長を含む4次元以上。ここは2次元に潰した
• 素朴なループだけ: ブロック化も SIMD も GPU も使わない。回る順の話までで止めている
• 自動微分なし: 学習には forward の逆(勾配計算)が要る。ここは推論だけ
• LayerNorm の学習パラメータなし: 実物は正規化のあとに学習された gain と bias を掛ける
• メモリの使い回しなし: 毎回新しい行列を作る。実物は確保済みの領域に書き戻す
• 形のチェックが最小: 掛けられるかどうかだけ見る

参考資料

• llm.c — GPT-2 を C で書く
• Neural Networks: Zero to Hero (Karpathy) — 行列演算から Transformer まで手を動かす講義
• Hendrycks, Gimpel, Gaussian Error Linear Units (GELUs)(2016)
• 実装: llm/tensor