Transformer ブロック：残差接続・層正規化・MLP とパラメータ数

0. この記事の要点

ChatGPT のような文章生成モデルの中身は、同じ形の「ブロック」を何十段も積み重ねたものです。この記事では、1 つのブロックが何でできているのか、なぜ深く積んでも学習が崩れにくいのか、モデル全体にパラメータが何個あるのかを、式と小さな数値例で確かめます。読み終えると、ブロックの図を自分で描け、パラメータ数を手計算で数えられるようになります。

• Transformer ブロックは、自己注意と位置ごとの MLP という 2 つの部品を、それぞれ残差接続 $\boldsymbol{x} \mapsto \boldsymbol{x} + F(\boldsymbol{x})$ と層正規化で包んだものです。これを $L$ 個重ね、前後に埋め込みと出力層を付けると、デコーダのみの Transformer になります。
• 層正規化は、入力の正の定数倍と全成分への定数の加算で出力が変わりません（命題 3.1）。その微分形として、ヤコビ行列（出力の各成分が入力の各成分でどれだけ変わるかを並べた行列）は $\boldsymbol{1}$ と $\boldsymbol{x} - \mu\boldsymbol{1}$ の 2 方向を消し、大きさは $1/\sigma(\boldsymbol{x})$ に比例します（命題 3.3）。
• Pre-LN の残差接続を重ねると、出力から各層へのヤコビ行列に恒等写像の項が必ず残ります（命題 4.1）。どの配置が学習しやすいかは経験的な事柄で、定理ではありません。
• パラメータ数は $|V|d + L_{\mathrm{ctx}}d + L(4d^2 + 4d + 2d\,d_{\mathrm{ff}} + d_{\mathrm{ff}} + d + 4d) + 2d$ です（定理 5.1）。12 層・$d = 768$・語彙 50,257・文脈長 1,024 では 124,439,808 になります。
• 1 層の乗算の回数は $12nd^2 + 2n^2d$（$d_{\mathrm{ff}} = 4d$ のとき）で、$n < 6d$ では MLP と射影の項が注意の項より大きくなります（命題 6.1）。
• 位置の間で情報を運ぶのは注意だけで、MLP と層正規化は行ごとに働きます。したがって前章の同変性（入力の行を並べ替えると出力も同じように並べ替わる性質）と因果性（ある位置の出力が後ろの位置に依らない性質）は、ブロックとモデル全体に引き継がれます（命題 7.2）。

この記事で使う事前知識

• 自己注意: 各位置のベクトルを、ほかの位置のベクトルの重み付き平均で置き換える演算。→ 自己注意
• 同変: 入力の行を並べ替えると、出力も同じ並べ替えになる性質。→ 同変
• 埋め込み行列: トークンの番号を $d$ 次元のベクトルに写す表。→ 埋め込み行列
• 多層パーセプトロン: 線形変換と成分ごとの非線形関数を交互に重ねた関数。→ 多層パーセプトロン
• 分散: 平均からのずれの 2 乗の平均。→ 分散
• 正規直交基底: 長さが 1 で互いに直交するベクトルからなる基底。→ 正規直交基底

1. 動機：注意だけでは足りないもの

前章の自己注意は、各位置のベクトルをほかの位置のベクトルの凸結合で置き換える演算でした（命題 4.2［自己注意］）。これだけを重ねても、モデルとしては 3 つのものが足りません。

1. 各位置での非線形な加工。 注意は位置の間で情報を混ぜますが、混ぜたあとの 1 本のベクトルを作り替える働きは、射影の線形変換しか持っていません。そこで行ごとに MLP を施します。
2. 深く重ねても学習できる構造。 層を単純に合成すると、勾配は層ごとのヤコビ行列の積になり、その大きさの上界は層の数について指数的に小さくなりえます（命題 8.1［ニューラルネットワークと逆伝播］）。残差接続は、この積の中に恒等写像の経路を作ります。
3. 値の大きさをそろえる仕組み。 残差接続で足し続けると、ベクトルの大きさが層ごとに変わっていきます。層正規化は、各行を平均 0・分散 1 に直します。

残差接続は He らが画像認識の深いネットワークのために導入し（[3]）、層正規化は Ba らが提案しました（[2]）。Vaswani らはこの 2 つで注意と MLP を包んで Transformer の層にしました（[1]）。この章では、これらの部品について厳密に言えることを命題にし、最後にモデル全体を NumPy で組んでパラメータ数と因果性を確かめます。

2. 準備：部品の定義

前章と同じく、$X \in \mathbb{R}^{n \times d}$ の第 $i$ 行を列ベクトルとみなしたものを $\boldsymbol{x}_i$ と書きます。

定義 2.1（層正規化）↗

$\boldsymbol{x} \in \mathbb{R}^{d}$ の成分の平均と分散を

$$\mu(\boldsymbol{x}) := \frac{1}{d}\sum_{j=1}^{d} x_j, \qquad \sigma^2(\boldsymbol{x}) := \frac{1}{d}\sum_{j=1}^{d}\big(x_j - \mu(\boldsymbol{x})\big)^2$$

とする。パラメータ $\boldsymbol{\gamma}, \boldsymbol{\beta} \in \mathbb{R}^{d}$ と定数 $\varepsilon \ge 0$ に対し

$$\mathrm{LN}(\boldsymbol{x}) := \boldsymbol{\gamma} \odot \frac{\boldsymbol{x} - \mu(\boldsymbol{x})\boldsymbol{1}}{\sqrt{\sigma^2(\boldsymbol{x}) + \varepsilon}} + \boldsymbol{\beta}$$

を層正規化という（$\odot$ は成分ごとの積）。$\varepsilon = 0$ のときは $\sigma^2(\boldsymbol{x}) > 0$、すなわち $\boldsymbol{x}$ の成分がすべて等しくはない場合に定義される。行列 $X$ に対しては各行に施す。

$\varepsilon$ は 0 で割ることを避けるための小さな正の数です（実装では $10^{-5}$ 程度）。以下の命題は、式が簡単になる $\varepsilon = 0$ の場合について述べます。

定義 2.2（残差接続・位置ごとの MLP・ブロック）↗

（残差接続） 写像 $F : \mathbb{R}^{n\times d} \to \mathbb{R}^{n \times d}$ に対し、$X \mapsto X + F(X)$ を $F$ の残差接続という。

（位置ごとの MLP） パラメータ $W_1 \in \mathbb{R}^{d \times d_{\mathrm{ff}}}$、$\boldsymbol{b}_1 \in \mathbb{R}^{d_{\mathrm{ff}}}$、$W_2 \in \mathbb{R}^{d_{\mathrm{ff}} \times d}$、$\boldsymbol{b}_2 \in \mathbb{R}^{d}$ と活性化関数 $\phi : \mathbb{R} \to \mathbb{R}$ に対し

$$\mathrm{MLP}(X) := \phi\big(XW_1 + \boldsymbol{1}\boldsymbol{b}_1^{\mathsf{T}}\big)W_2 + \boldsymbol{1}\boldsymbol{b}_2^{\mathsf{T}}$$

と定める（$\phi$ は成分ごと）。第 $i$ 行は $g(\boldsymbol{x}_i) := W_2^{\mathsf{T}}\phi(W_1^{\mathsf{T}}\boldsymbol{x}_i + \boldsymbol{b}_1) + \boldsymbol{b}_2$ で、1 層の隠れ層を持つ多層パーセプトロン（定義 2.1［ニューラルネットワークと逆伝播］ で層数を 2、2 層目の活性化関数を恒等写像としたもの。同定義は活性化関数を $C^1$ 級と仮定しており、$\mathrm{GELU}$ はこれを満たすが、$\mathrm{ReLU}$ は 0 で微分可能でないので満たさない）を各行に施したものである。$d_{\mathrm{ff}}$ を MLP の幅という。活性化関数は $\mathrm{ReLU}(t) := \max(0, t)$ か、標準正規分布の分布関数 $\Phi$ を使った $\mathrm{GELU}(t) := t\,\Phi(t)$ を取る。

（ブロック） マルチヘッド自己注意 $\mathrm{MHA}$（定義 4.1［自己注意］）、位置ごとの $\mathrm{MLP}$、2 つの層正規化 $\mathrm{LN}_1, \mathrm{LN}_2$ から、次の 2 通りのブロックを定める。

$$\begin{aligned} \textbf{Post-LN：}\quad & Y = \mathrm{LN}_1\big(X + \mathrm{MHA}(X)\big), & X' &= \mathrm{LN}_2\big(Y + \mathrm{MLP}(Y)\big) \\ \textbf{Pre-LN：}\quad & Y = X + \mathrm{MHA}\big(\mathrm{LN}_1(X)\big), & X' &= Y + \mathrm{MLP}\big(\mathrm{LN}_2(Y)\big) \end{aligned}$$

Pre-LN のブロックで、層から層へ足し込まれながら受け渡される $X, Y, X', \ldots$ を残差流という。

Post-LN は Vaswani らの配置で、各部品の出力を「残差に足してから正規化」します（[1] 3.1 節）。Pre-LN は「正規化してから部品に入れ、残差に足す」配置で、Radford らは層正規化を各部品の入力側に移し、最後の自己注意ブロックのあとに正規化を 1 つ追加したと述べています（[5] 2.3 節）。Vaswani らの MLP は ReLU で $d = 512$、$d_{\mathrm{ff}} = 2048 = 4d$ です（[1] 3.3 節）。GELU は Hendrycks と Gimpel が提案した活性化関数です（[7]）。

Pre-LN のブロック。上の線が残差流で、注意と MLP の出力が順に足し込まれる。正規化は各部品の入口にあり、残差流そのものには掛からない。

定義 2.3（デコーダのみの Transformer）↗

語彙 $V$、次元 $d$、層数 $L$、MLP の幅 $d_{\mathrm{ff}}$、文脈長 $L_{\mathrm{ctx}}$ を定める。パラメータは、埋め込み行列 $E \in \mathbb{R}^{|V| \times d}$、学習する位置埋め込み $P \in \mathbb{R}^{L_{\mathrm{ctx}} \times d}$（命題 4.6［埋め込みと位置符号化］）、$L$ 個の Pre-LN ブロック（注意は因果マスクつき）のパラメータ、最後の層正規化 $\mathrm{LN}_f$ のパラメータである。トークン列 $(t_1, \ldots, t_n)$（$n \le L_{\mathrm{ctx}}$）に対し

1. $X^{(0)}$ を、第 $i$ 行が $E_{t_i} + P_i$ である $n \times d$ 行列とする。
2. $l = 1, \ldots, L$ の順に、$X^{(l)}$ を $X^{(l-1)}$ に第 $l$ ブロックを施したものとする。
3. $Z := \mathrm{LN}_f\big(X^{(L)}\big)\,E^{\mathsf{T}} \in \mathbb{R}^{n \times |V|}$ とする。

$Z$ の第 $i$ 行 $\boldsymbol{z}_i$ を位置 $i$ のロジットといい、$\sigma(\boldsymbol{z}_i)$（softmax）を、$t_1, \ldots, t_i$ の次に来るトークンの予測分布とする。この写像を（入出力の重み共有ありの）デコーダのみの Transformer という。

手順 3 で出力側の行列に $E$ そのものを使っているのが重み共有です（定義 2.1［埋め込みと位置符号化］）。$Z_{iv} = \langle \mathrm{LN}_f(\boldsymbol{x}^{(L)}_i), E_v\rangle$ なので、ロジットは「最後のベクトルと各トークンの埋め込みの内積」です。ロジットは softmax にかける前の、各トークンの得点にあたります。入力側と出力側で同じ表 $E$ を使うので、出力側のために新しく $|V|\,d$ 個のパラメータを持たずに済みます。位置の与え方を RoPE（クエリとキーを位置に応じた角度で回す方式。第 2 章で扱いました）にする場合は、$P$ を無くし、各層の注意の中でクエリとキーを回します。

3. 層正規化の性質

命題 3.1（層正規化はスケールと平行移動に不変）✓

$\varepsilon = 0$ とし、$\boldsymbol{x} \in \mathbb{R}^d$ は成分がすべて等しくはないとする。

1. 任意の $a > 0$ と $b \in \mathbb{R}$ について $\mathrm{LN}(a\boldsymbol{x} + b\boldsymbol{1}) = \mathrm{LN}(\boldsymbol{x})$。
2. $\boldsymbol{\gamma} = \boldsymbol{1}$、$\boldsymbol{\beta} = \boldsymbol{0}$ のとき、$\boldsymbol{y} := \mathrm{LN}(\boldsymbol{x})$ は $\mu(\boldsymbol{y}) = 0$、$\sigma^2(\boldsymbol{y}) = 1$、$\|\boldsymbol{y}\| = \sqrt{d}$ を満たす。

証明（命題 3.1）

$\mu := \mu(\boldsymbol{x})$、$\sigma := \sqrt{\sigma^2(\boldsymbol{x})} > 0$ と書きます。

1. $\boldsymbol{x}' := a\boldsymbol{x} + b\boldsymbol{1}$ について、$\mu(\boldsymbol{x}') = \frac1d\sum_j(ax_j + b) = a\mu + b$。よって $x'_j - \mu(\boldsymbol{x}') = a(x_j - \mu)$ で、$\sigma^2(\boldsymbol{x}') = \frac1d\sum_j a^2(x_j - \mu)^2 = a^2\sigma^2$。$a > 0$ なので $\sqrt{\sigma^2(\boldsymbol{x}')} = a\sigma$。したがって

$$\frac{\boldsymbol{x}' - \mu(\boldsymbol{x}')\boldsymbol{1}}{\sqrt{\sigma^2(\boldsymbol{x}')}} = \frac{a(\boldsymbol{x} - \mu\boldsymbol{1})}{a\sigma} = \frac{\boldsymbol{x} - \mu\boldsymbol{1}}{\sigma}$$

で、これに $\boldsymbol{\gamma}\odot$ と $+\boldsymbol{\beta}$ を施したものが $\mathrm{LN}$ なので、両者は等しくなります。

2. $y_j = (x_j - \mu)/\sigma$ です。$\sum_j y_j = \frac1\sigma\big(\sum_j x_j - d\mu\big) = 0$ より $\mu(\boldsymbol{y}) = 0$。$\sigma^2(\boldsymbol{y}) = \frac1d\sum_j y_j^2 = \frac{1}{d\sigma^2}\sum_j(x_j - \mu)^2 = \frac{d\sigma^2}{d\sigma^2} = 1$。$\|\boldsymbol{y}\|^2 = \sum_j y_j^2 = d\,\sigma^2(\boldsymbol{y}) = d$。

∎

2 は、層正規化の出力が「$\boldsymbol{1}$ に直交する、半径 $\sqrt{d}$ の球面」の上にあることを言っています。入力の大きさの情報は捨てられ、向きだけが残ります。$\boldsymbol{\gamma}, \boldsymbol{\beta}$ は、そのあとで成分ごとの倍率とずれを学習で付け直すためのパラメータです。

例 3.2（3 倍しても出力は変わらない）

$\boldsymbol{x} = (1, 2, 3, 6)$ とします。$\mu = 3$、$\sigma^2 = \frac14(4 + 1 + 0 + 9) = 3.5$、$\sigma = 1.8708$ なので、$\boldsymbol{\gamma} = \boldsymbol{1}$、$\boldsymbol{\beta} = \boldsymbol{0}$、$\varepsilon = 0$ で

$$\mathrm{LN}(\boldsymbol{x}) = \frac{(-2, -1, 0, 3)}{1.8708} = (-1.0690,\ -0.5345,\ 0,\ 1.6036)$$

です。$3\boldsymbol{x} + 5\cdot\boldsymbol{1} = (8, 11, 14, 23)$ では $\mu = 14$、$\sigma^2 = \frac14(36 + 9 + 0 + 81) = 31.5 = 9 \times 3.5$、$\sigma = 5.6125$ で、$(-6, -3, 0, 9)/5.6125$ は同じベクトルになります。ノルムは $\sqrt{1.1428 + 0.2857 + 0 + 2.5714} = \sqrt{4} = 2 = \sqrt{d}$ です。

次の命題は、層正規化を通して勾配がどう戻るかを調べます。入力に定数を足しても定数倍しても出力は変わらないので、その 2 方向への変化率は 0 になるはずです。証明の要は、ヤコビ行列を「その 2 方向の成分を取り除く操作（直交射影）の $1/\sigma$ 倍」と読み替えることで、そのために $\boldsymbol{1}$ と $\boldsymbol{y}$ から正規直交基底を作ります。

命題 3.3（層正規化のヤコビ行列）✓

$\varepsilon = 0$、$\boldsymbol{\gamma} = \boldsymbol{1}$、$\boldsymbol{\beta} = \boldsymbol{0}$ とし、$\boldsymbol{x} \in \mathbb{R}^d$ は成分がすべて等しくはないとする。$\sigma := \sqrt{\sigma^2(\boldsymbol{x})}$、$\boldsymbol{y} := \mathrm{LN}(\boldsymbol{x})$ と書くと、$\mathrm{LN}$ は $\boldsymbol{x}$ で微分可能で、ヤコビ行列は

$$J(\boldsymbol{x}) = \frac{1}{\sigma}\Big( I - \frac{1}{d}\boldsymbol{1}\boldsymbol{1}^{\mathsf{T}} - \frac{1}{d}\boldsymbol{y}\boldsymbol{y}^{\mathsf{T}} \Big)$$

である。さらに次が成り立つ。

1. $J(\boldsymbol{x})\boldsymbol{1} = \boldsymbol{0}$、$J(\boldsymbol{x})\big(\boldsymbol{x} - \mu(\boldsymbol{x})\boldsymbol{1}\big) = \boldsymbol{0}$。
2. $W := \mathrm{span}\{\boldsymbol{1},\ \boldsymbol{x} - \mu(\boldsymbol{x})\boldsymbol{1}\}$ とすると、$\sigma J(\boldsymbol{x})$ は $W$ の直交補空間 $W^{\perp}$ への直交射影である。すなわち $\boldsymbol{v} \in W^{\perp}$ ならば $J(\boldsymbol{x})\boldsymbol{v} = \boldsymbol{v}/\sigma$。
3. 任意の $a > 0$ について $J(a\boldsymbol{x}) = J(\boldsymbol{x})/a$。

証明（命題 3.3）

$c_i := x_i - \mu(\boldsymbol{x})$ と置きます。$\partial\mu/\partial x_j = 1/d$ なので $\partial c_i/\partial x_j = \delta_{ij} - 1/d$ です。$\sigma^2 = \frac1d\sum_i c_i^2$ を微分すると

$$\frac{\partial \sigma^2}{\partial x_j} = \frac{2}{d}\sum_i c_i\Big(\delta_{ij} - \frac1d\Big) = \frac2d\Big(c_j - \frac1d\sum_i c_i\Big) = \frac{2}{d}c_j$$

（$\sum_i c_i = 0$ を使いました）。$\sigma = \sqrt{\sigma^2}$、$\sigma > 0$ なので $\partial\sigma/\partial x_j = \dfrac{1}{2\sigma}\cdot\dfrac{2}{d}c_j = \dfrac{c_j}{d\sigma}$。$y_i = c_i/\sigma$ に商の微分を使うと

$$\frac{\partial y_i}{\partial x_j} = \frac{1}{\sigma}\Big(\delta_{ij} - \frac1d\Big) - \frac{c_i}{\sigma^2}\cdot\frac{c_j}{d\sigma} = \frac{1}{\sigma}\Big(\delta_{ij} - \frac1d - \frac{1}{d}\,y_iy_j\Big)$$

で、これが主張の行列の $(i, j)$ 成分です。偏導関数が連続なので微分可能です。

2 を先に示します。$\boldsymbol{u} := \boldsymbol{1}/\sqrt{d}$、$\boldsymbol{w} := \boldsymbol{y}/\sqrt{d}$ と置くと、$\|\boldsymbol{u}\| = 1$、命題 3.1 の 2 より $\|\boldsymbol{w}\| = \|\boldsymbol{y}\|/\sqrt{d} = 1$、$\langle\boldsymbol{u}, \boldsymbol{w}\rangle = \frac1d\sum_j y_j = 0$ です。つまり $\boldsymbol{u}, \boldsymbol{w}$ は $W = \mathrm{span}\{\boldsymbol{1}, \boldsymbol{y}\} = \mathrm{span}\{\boldsymbol{1}, \boldsymbol{x} - \mu\boldsymbol{1}\}$ の正規直交基底で、

$$\sigma J(\boldsymbol{x}) = I - \boldsymbol{u}\boldsymbol{u}^{\mathsf{T}} - \boldsymbol{w}\boldsymbol{w}^{\mathsf{T}}$$

です。直交分解定理（定理 7.1［内積空間とグラム・シュミット直交化］ の 2）より、$W$ への直交射影は $P_W\boldsymbol{v} = \langle\boldsymbol{v},\boldsymbol{u}\rangle\boldsymbol{u} + \langle\boldsymbol{v},\boldsymbol{w}\rangle\boldsymbol{w} = (\boldsymbol{u}\boldsymbol{u}^{\mathsf{T}} + \boldsymbol{w}\boldsymbol{w}^{\mathsf{T}})\boldsymbol{v}$ なので、$\sigma J(\boldsymbol{x}) = I - P_W$ です。同じ定理の 1 より任意の $\boldsymbol{v}$ は $\boldsymbol{v} = P_W\boldsymbol{v} + \boldsymbol{v}'$（$\boldsymbol{v}' \in W^{\perp}$）と一意に書けるので、$I - P_W$ は $W^{\perp}$ への直交射影です。$\boldsymbol{v} \in W^{\perp}$ なら $P_W\boldsymbol{v} = \boldsymbol{0}$ で、$J(\boldsymbol{x})\boldsymbol{v} = \boldsymbol{v}/\sigma$ です。

1. $\boldsymbol{1}$ と $\boldsymbol{x} - \mu\boldsymbol{1}$ は $W$ の元なので、$(I - P_W)$ で $\boldsymbol{0}$ に写ります。

2. 命題 3.1 の 1 より $\mathrm{LN}(a\boldsymbol{x}) = \boldsymbol{y}$ で、その証明より $\sqrt{\sigma^2(a\boldsymbol{x})} = a\sigma$。ヤコビ行列の式で $\boldsymbol{y}$ は同じ、$\sigma$ が $a\sigma$ に替わるので $J(a\boldsymbol{x}) = J(\boldsymbol{x})/a$ です。

∎

1 は 命題 3.1 の微分形です。$\boldsymbol{1}$ の向き（全成分に同じ数を足す）と $\boldsymbol{x} - \mu\boldsymbol{1}$ の向き（平均からのずれを定数倍する）に動かしても出力が変わらないので、その 2 方向の微分が 0 になります。3 は、層正規化の手前の値が大きいほど、そこを通って戻る勾配が小さくなることを言っています。入力を $a$ 倍すると勾配は $1/a$ 倍です。なお $d = 2$ では $W = \mathbb{R}^2$ となって $J = O$ です。実際、$d = 2$ の層正規化の出力は $(1, -1)$ か $(-1, 1)$ のどちらかしか取りません。

4. 残差接続と恒等経路

この節では、$n \times d$ 行列を $nd$ 次元のベクトルとみなし、ブロックの中の部品（注意と MLP）を区別せずに通し番号 $l = 1, \ldots, M$ を付けます（$L$ ブロックなら $M = 2L$）。

次の命題のねらいは、残差接続があると勾配に「層を通らない道」が必ず残ることを、式で示すことです。Pre-LN では各部品が前の値に変化分を足すだけなので、1 段のヤコビ行列は $I + J_l$ の形になります。この積を展開すると、各段で「迂回する（$I$ を選ぶ）」か「通る（$J_l$ を選ぶ）」かの 2 択の組み合わせになり、3 段なら下の図のように 8 通りです。Post-LN では足したあとに正規化が掛かるので、1 段が $K_l(I + D_l)$ の形になります。$K_l$ は正規化のヤコビ行列、$D_l$ は部品 $F_l$ のヤコビ行列で、どの段でも $K_l$ が掛かるため、単独の $I$ の項は現れません。

3 段の Pre-LN で、積 (I + J₃)(I + J₂)(I + J₁) を展開したときの 8 つの項。各段で I（迂回）か J（通る）を選ぶ。すべて I を選ぶ道（太線）が恒等経路で、残差接続が無ければ一番下の J₃J₂J₁ だけが残る。

命題 4.1（残差接続の恒等経路）✓

$F_1, \ldots, F_M$ と正規化 $N_1, \ldots, N_M$ を $\mathbb{R}^{nd}$ から $\mathbb{R}^{nd}$ への微分可能な写像とする。

（Pre-LN） $\boldsymbol{x}_l := \boldsymbol{x}_{l-1} + F_l\big(N_l(\boldsymbol{x}_{l-1})\big)$（$l = 1, \ldots, M$）とする。$G_l := F_l \circ N_l$、$J_l :=$（$G_l$ の $\boldsymbol{x}_{l-1}$ でのヤコビ行列）と書くと

$$\boldsymbol{x}_M = \boldsymbol{x}_0 + \sum_{l=1}^{M} G_l(\boldsymbol{x}_{l-1}), \qquad \frac{\partial \boldsymbol{x}_M}{\partial \boldsymbol{x}_l} = (I + J_M)(I + J_{M-1})\cdots(I + J_{l+1}) = I + \sum_{\emptyset \ne S \subset \{l+1, \ldots, M\}} \ \prod_{k \in S}^{\leftarrow} J_k$$

が成り立つ。ここで $\prod^{\leftarrow}$ は添字の大きい順に左から並べた積である。とくに右辺には、どの $J_k$ も含まない項として恒等行列 $I$ がちょうど 1 つある。

（Post-LN） $\boldsymbol{x}_l := N_l\big(\boldsymbol{x}_{l-1} + F_l(\boldsymbol{x}_{l-1})\big)$ とする。$D_l :=$（$F_l$ の $\boldsymbol{x}_{l-1}$ でのヤコビ行列）、$K_l :=$（$N_l$ の $\boldsymbol{x}_{l-1} + F_l(\boldsymbol{x}_{l-1})$ でのヤコビ行列）と書くと

$$\frac{\partial \boldsymbol{x}_M}{\partial \boldsymbol{x}_l} = K_M(I + D_M)\,K_{M-1}(I + D_{M-1})\cdots K_{l+1}(I + D_{l+1})$$

で、どの $D_k$ も含まない項は $K_M K_{M-1}\cdots K_{l+1}$ である。

証明（命題 4.1）

Pre-LN。 定義より $\boldsymbol{x}_l - \boldsymbol{x}_{l-1} = G_l(\boldsymbol{x}_{l-1})$ で、$l = 1, \ldots, M$ について足すと左辺は $\boldsymbol{x}_M - \boldsymbol{x}_0$ になり、1 つ目の式を得ます。

写像 $\boldsymbol{x}_{k-1} \mapsto \boldsymbol{x}_k = \boldsymbol{x}_{k-1} + G_k(\boldsymbol{x}_{k-1})$ のヤコビ行列は $I + J_k$ です。$\boldsymbol{x}_M$ は $\boldsymbol{x}_l$ にこの写像を $k = l+1, \ldots, M$ の順に施したものなので、連鎖律（定理 3.1［ニューラルネットワークと逆伝播］）を繰り返し使って、ヤコビ行列はあとから施す写像のものを左に置いた積 $(I + J_M)\cdots(I + J_{l+1})$ です。

この積を展開します。各因子 $(I + J_k)$ から $I$ か $J_k$ のどちらかを選んで掛け合わせた項が、選び方の数 $2^{M-l}$ だけ出ます。$J_k$ を選んだ添字の集合を $S$ とすると、その項は $\prod^{\leftarrow}_{k\in S}J_k$ です（$I$ は積の中で消えます）。$S = \emptyset$ の項が $I$ で、これは 1 つだけです。

Post-LN。 写像 $\boldsymbol{x}_{k-1} \mapsto \boldsymbol{x}_{k-1} + F_k(\boldsymbol{x}_{k-1})$ のヤコビ行列は $I + D_k$、それに $N_k$ を合成した写像のヤコビ行列は連鎖律より $K_k(I + D_k)$ です。これを $k = l+1, \ldots, M$ について左に掛けていけば主張の式になります。各 $(I + D_k)$ から $I$ を選んだ項が $K_M\cdots K_{l+1}$ です。

∎

損失 $\mathcal{L}$ が $\boldsymbol{x}_M$ の関数のとき、連鎖律より $\nabla_{\boldsymbol{x}_l}\mathcal{L} = \big(\partial\boldsymbol{x}_M/\partial\boldsymbol{x}_l\big)^{\mathsf{T}}\nabla_{\boldsymbol{x}_M}\mathcal{L}$ です。Pre-LN では、右辺に $\nabla_{\boldsymbol{x}_M}\mathcal{L}$ そのもの（$I$ の項）が、途中の層を 1 つも通らずに含まれます。残差接続のない合成では、勾配は層ごとのヤコビ行列の積 1 つだけで、その大きさの上界は因子の数だけ掛け合わさって指数的に小さくなりえました（命題 8.1［ニューラルネットワークと逆伝播］）。恒等経路は、その積を経由しない道を 1 本用意します。

Post-LN では、同じ位置にある項が層正規化のヤコビ行列の積 $K_M\cdots K_{l+1}$ になります。命題 3.3 より各 $K_k$ は、行ごとに 2 方向を消して $1/\sigma$ 倍する行列なので、この項は恒等行列ではありません。

注意 4.2（どちらが学習しやすいかは経験的な事柄）

命題 4.1 が言っているのは、ヤコビ行列の展開に $I$ の項があるという式の上の事実だけです。ほかの項との和がどういう大きさになるか、学習がうまく進むかは、この命題からは出ません。

経験的な報告としては、Xiong らが、Post-LN の Transformer は初期化の時点で出力に近い層の勾配が大きく、学習率のウォームアップ（最初に学習率を小さく始めて徐々に上げること）が必要になること、Pre-LN では初期化の時点の勾配の振る舞いが良く、ウォームアップなしでも同程度の結果が得られることを、理論的な分析と実験で論じています（[6]）。He らは画像認識の残差ネットワークで、残差の経路を恒等写像のまま保つ配置を調べています（[4]）。これらは特定の初期化・設定のもとでの結果で、「Pre-LN が常に良い」という定理ではありません。

5. パラメータ数

定理 5.1（デコーダのみの Transformer のパラメータ数）✓

定義 2.3 のモデル（語彙 $|V|$、次元 $d$、層数 $L$、MLP の幅 $d_{\mathrm{ff}}$、長さ $L_{\mathrm{ctx}}$ の学習する位置埋め込み、入出力の重み共有あり、注意・MLP にバイアスあり）のパラメータ数は

$$|V|\,d + L_{\mathrm{ctx}}\,d + L\big(\underbrace{4d^2 + 4d}_{\text{注意}} + \underbrace{2d\,d_{\mathrm{ff}} + d_{\mathrm{ff}} + d}_{\text{MLP}} + \underbrace{4d}_{\text{層正規化 2 つ}}\big) + \underbrace{2d}_{\text{最後の層正規化}}$$

である。$d_{\mathrm{ff}} = 4d$ のとき、1 層あたりは $12d^2 + 13d$ である。位置を RoPE や正弦波で与える場合は $L_{\mathrm{ctx}}\,d$ の項が無くなり、重み共有をしない場合は出力側の行列の $|V|\,d$ が加わる。

証明（定理 5.1）

構成要素ごとに成分の個数を数えます。

• 埋め込み行列 $E \in \mathbb{R}^{|V|\times d}$：$|V|\,d$ 個。出力層は同じ $E$ を使うので、新しいパラメータはありません。
• 位置埋め込み $P \in \mathbb{R}^{L_{\mathrm{ctx}} \times d}$：$L_{\mathrm{ctx}}\,d$ 個（命題 4.6［埋め込みと位置符号化］）。
• 各ブロックのマルチヘッド自己注意：$4d^2 + 4d$ 個（命題 7.1［自己注意］）。
• 各ブロックの MLP：$W_1$ が $d\,d_{\mathrm{ff}}$、$\boldsymbol{b}_1$ が $d_{\mathrm{ff}}$、$W_2$ が $d_{\mathrm{ff}}\,d$、$\boldsymbol{b}_2$ が $d$ で、計 $2d\,d_{\mathrm{ff}} + d_{\mathrm{ff}} + d$ 個。
• 各ブロックの層正規化 2 つ：それぞれ $\boldsymbol{\gamma}, \boldsymbol{\beta} \in \mathbb{R}^d$ で $2d$ 個、計 $4d$ 個。
• 最後の層正規化：$2d$ 個。

ブロックは $L$ 個あるので、合計が主張の式です。$d_{\mathrm{ff}} = 4d$ を代入すると、1 層あたり $4d^2 + 4d + 8d^2 + 4d + d + 4d = 12d^2 + 13d$ です。RoPE と正弦波はパラメータを持たないので $P$ の分が無くなります。重み共有をしなければ、出力側に別の $|V| \times d$ 行列が要ります。

∎

例 5.2（小さなモデルで数える）

$|V| = 10$、$d = 4$、$d_{\mathrm{ff}} = 16$、$L = 2$、$L_{\mathrm{ctx}} = 8$ とします。

• 埋め込み：$10 \times 4 = 40$。位置：$8 \times 4 = 32$。
• 1 層の注意：$4\times16 + 4\times4 = 80$。MLP：$2\times4\times16 + 16 + 4 = 148$。層正規化：$16$。1 層で $244$（$12d^2 + 13d = 192 + 52 = 244$ とも合います）。
• 最後の層正規化：$8$。

合計は $40 + 32 + 2\times244 + 8 = 568$ です。9 節の実装で配列の成分を数えると、同じ 568 になります。

例 5.3（12 層・d = 768 の構成での検算）

Radford らの最小の構成は 12 層・$d = 768$ で、語彙は 50,257、文脈長は 1,024 です（[5] 表 2、2.3 節）。層正規化の配置は Pre-LN で、最後に正規化が 1 つあります（同 2.3 節）。この値を 定理 5.1 の式に入れます。MLP の幅は $d_{\mathrm{ff}} = 4d = 3072$ とします。

構成要素	式	個数
埋め込み	$50257 \times 768$	38,597,376
位置埋め込み	$1024 \times 768$	786,432
ブロック 1 つ	$12\times768^2 + 13\times768$	7,087,872
ブロック 12 個	$12 \times 7{,}087{,}872$	85,054,464
最後の層正規化	$2\times768$	1,536
合計		124,439,808

注意 5.4（論文の表の値との違い）

例 5.3 の数え上げの結果は 124,439,808 です。一方、Radford らの論文の表 2 は、この構成（12 層・$d_{\mathrm{model}} = 768$）のパラメータ数を 117M と記載しています（[5]）。2 つの数は一致しません。

確かめられる事実として、著者らが公開したコードのリポジトリの説明には、当初のパラメータ数（ブログ記事と論文に書いたもの）は誤りによるものだったという注記があります（[8]）。同じリポジトリのモデルのコードでは、MLP の幅が $4d$ であること、層正規化が各部品の入力側と最後にあること、出力のロジットを埋め込み行列の転置との積で作っていること（重み共有）、位置埋め込みが文脈長 $\times$ $d$ の表であることが読み取れ、この章の数え上げの前提と合っています。117M という値がどのように計算されたかについては、論文にも上の注記にも説明が無いので、ここでは推測を書きません。

パラメータ数の式から読み取れることを 2 つ挙げます。1 つ目に、ブロックの部分は $d$ の 2 乗と $L$ に比例し、語彙や文脈長には依りません。2 つ目に、小さいモデルでは埋め込みの占める割合が大きくなります。上の構成では $38{,}597{,}376 / 124{,}439{,}808 \approx 31\%$ が埋め込みです。語彙の大きさ $|V|$ は第 1 章のマージ回数で決まるので、トークン化の設計がモデルの大きさに直接効きます。

6. 順伝播の計算量

命題 6.1（1 層の順伝播の計算量）✓

長さ $n$ の列に対する Pre-LN（または Post-LN）のブロック 1 つの順伝播で、スカラーの乗算の回数は

$$\underbrace{4nd^2 + 2n^2 d}_{\text{注意}} + \underbrace{2nd\,d_{\mathrm{ff}}}_{\text{MLP}} + O(nd)$$

である。$d_{\mathrm{ff}} = 4d$ のとき $12nd^2 + 2n^2d + O(nd)$ で、時間計算量は $O(nd^2 + n^2d)$ である。$n < 6d$ のとき、$n$ について 1 次の項 $12nd^2$ が 2 次の項 $2n^2d$ より大きい。モデル全体では、これが $L$ 層分と、出力層の $nd\,|V|$ 回になる。

証明（命題 6.1）

注意の $4nd^2 + 2n^2d$ は 命題 7.2［自己注意］ です。MLP は、$XW_1$（$n \times d$ と $d \times d_{\mathrm{ff}}$ の積）が $nd\,d_{\mathrm{ff}}$ 回、その結果に $W_2$ を掛けるのが $n\,d_{\mathrm{ff}}\,d$ 回で、計 $2nd\,d_{\mathrm{ff}}$ 回です。層正規化は 1 行につき平均・分散・除算・$\boldsymbol{\gamma}$ との積で $d$ の定数倍、残差の加算は乗算なしなので、合わせて $O(nd)$ です。活性化関数と softmax の評価は乗算とは別に数えますが、回数は $O(n\,d_{\mathrm{ff}} + hn^2)$ で、上の主要項を超えません。

$d_{\mathrm{ff}} = 4d$ なら $4nd^2 + 8nd^2 = 12nd^2$ です。$O$ 記法の定義（定義 3.1［計算量と O 記法］）より全体は $O(nd^2 + n^2d)$。$12nd^2 > 2n^2d$ は、両辺を $2nd > 0$ で割ると $6d > n$ と同値です。出力層 $\mathrm{LN}_f(X^{(L)})E^{\mathsf{T}}$ は $n \times d$ と $d \times |V|$ の積で $nd\,|V|$ 回です。

これは 1 回の順伝播についての評価です。勾配の計算（逆伝播）は順伝播の定数倍の演算で済みます（命題 6.1［ニューラルネットワークと逆伝播］）。

∎

たとえば $d = 768$、$n = 1024$ では、$12nd^2 = 7{,}247{,}757{,}312$、$2n^2d = 1{,}610{,}612{,}736$ で、注意の 2 次の項は MLP と射影の項の $0.222$ 倍（$= n/(6d)$）です。「注意は列の長さの 2 乗で重くなる」のはそのとおりですが、列が $6d$ より短いあいだは、計算の大半は $n$ に比例する部分です。2 乗の項が支配的になるのは $n$ が $6d = 4608$ を超えてからで、長い文脈を扱うときの工夫は第 7 章で扱います。

7. 同変性と因果性はブロックに引き継がれる

命題 7.1（MLP と層正規化は位置の間で情報を混ぜない）✓

位置ごとの MLP と層正規化は行ごとに働く写像である。すなわち出力の第 $i$ 行は入力の第 $i$ 行だけの関数で、$i \ne j$ のとき、出力の第 $i$ 行を入力の第 $j$ 行で微分したヤコビ行列は 0 である。したがって、自己注意を含まず、行ごとに働く写像・位置ベクトルの加算・それらの合成と和（残差接続）だけからなるモデルでは、位置 $i$ の出力は位置 $i$ のトークンと $i$ だけで決まる。

証明（命題 7.1）

定義 2.2 より $\mathrm{MLP}(X)$ の第 $i$ 行は $g(\boldsymbol{x}_i)$、定義 2.1 より $\mathrm{LN}(X)$ の第 $i$ 行は $\mathrm{LN}(\boldsymbol{x}_i)$ で、どちらも $\boldsymbol{x}_i$ だけの関数です。$\boldsymbol{x}_j$（$j \ne i$）を含まない式を $\boldsymbol{x}_j$ で微分すれば 0 です。

後半は組み立ての段数についての帰納法です。「第 $i$ 行が、入力の第 $i$ 行と $i$ だけの関数である」という性質は、入力の埋め込み（第 $i$ 行は $E_{t_i} + P_i$）で成り立ち、行ごとに働く写像との合成、2 つのそのような写像の和で保たれます（合成 $g(f_i(\boldsymbol{x}_i))$ も和 $f_i(\boldsymbol{x}_i) + f'_i(\boldsymbol{x}_i)$ も、$\boldsymbol{x}_i$ と $i$ だけの関数）。

∎

つまり、ある位置の予測が別の位置のトークンに依存できるのは、注意があるからです。ブロックの中での役割分担は、「注意がほかの位置から情報を集め、MLP が集めた結果をその位置の中で加工する」と整理できます。

命題 7.2（ブロックとモデル全体の同変性・因果性）✓

1. マスクのない注意を使った Pre-LN・Post-LN のブロックは置換同変である。
2. 因果マスクつきの注意を使った Pre-LN・Post-LN のブロックは因果的（定義 6.1［自己注意］）である。
3. 定義 2.3 のモデルについて、埋め込み後の行列 $X^{(0)}$ からロジット $Z$ への写像は因果的である。したがって位置 $i$ のロジット $\boldsymbol{z}_i$ は $t_1, \ldots, t_i$ だけで決まり、長さ $n$ の列を入力したときの $\boldsymbol{z}_i$ は、接頭辞 $(t_1, \ldots, t_i)$ だけを入力したときの最後の位置のロジットに等しい。

証明（命題 7.2）

層正規化と MLP は行ごとに働く写像です（命題 7.1）。Pre-LN のブロックは $X \mapsto X + \mathrm{MHA}(\mathrm{LN}_1(X))$ と $Y \mapsto Y + \mathrm{MLP}(\mathrm{LN}_2(Y))$ の合成、Post-LN のブロックは $X \mapsto \mathrm{LN}_1(X + \mathrm{MHA}(X))$ と $Y \mapsto \mathrm{LN}_2(Y + \mathrm{MLP}(Y))$ の合成で、どちらも、自己注意・行ごとに働く写像・合成・和だけからできています。

1. マスクのない自己注意は置換同変で、行ごとに働く写像との合成と和で同変性は保たれます（定理 5.2［自己注意］）。
2. 因果マスクつきの自己注意は因果的で、行ごとに働く写像との合成と和で因果性は保たれます（定理 6.3［自己注意］）。
3. $X^{(0)} \mapsto X^{(L)}$ は因果的なブロック $L$ 個の合成、$X^{(L)} \mapsto \mathrm{LN}_f(X^{(L)})E^{\mathsf{T}}$ は行ごとに働く写像（第 $i$ 行は $E\,\mathrm{LN}_f(\boldsymbol{x}^{(L)}_i)$）なので、定理 6.3［自己注意］ より全体は因果的です。$X^{(0)}$ の最初の $i$ 行は $t_1, \ldots, t_i$ だけで決まり（第 $j$ 行は $E_{t_j} + P_j$）、接頭辞 $(t_1, \ldots, t_i)$ の埋め込みは $X^{(0)}_{1:i}$ そのものなので、後半の主張は因果的であることの定義と 系 6.4［自己注意］ です。

∎

3 のモデルは位置埋め込み $P_i$ を足しているので、置換同変ではありません。トークンを並べ替えると、同じトークンに違う位置ベクトルが足されるからです。これが、語順を区別できるようにするという位置埋め込みの目的でした。

8. MLP は何を表せるか：区分線形性

この節では、MLP が入力空間をどう切り分けるかを見ます。ReLU は負の値を 0 にし、正の値はそのまま通す関数です。MLP の隠れ層の各成分は、入力がある境目より大きいか小さいかで「効いている」か「0 になっている」かが決まります。どの成分が効いているかの組（活性パターン）が同じ領域の中では、MLP 全体が 1 次式になります。これが次の命題の 1 です。2 は、GELU には同じ性質が無いことを確かめます。

命題 8.1（ReLU の MLP は区分線形）✓

$g(\boldsymbol{x}) = W_2^{\mathsf{T}}\,\mathrm{ReLU}(W_1^{\mathsf{T}}\boldsymbol{x} + \boldsymbol{b}_1) + \boldsymbol{b}_2$（$\boldsymbol{x} \in \mathbb{R}^d$、幅 $d_{\mathrm{ff}}$）とする。

1. $\mathbb{R}^d$ は、高々 $2^{d_{\mathrm{ff}}}$ 個の互いに交わらない凸集合に分かれ、そのそれぞれの上で $g$ は 1 次式（アフィン写像）に一致する。
2. $\mathrm{GELU}(t) = t\,\Phi(t)$ は $\mathbb{R}$ 上で何回でも微分可能で、どの開区間の上でも 1 次式に一致しない。

証明（命題 8.1）

1. $\boldsymbol{a}(\boldsymbol{x}) := W_1^{\mathsf{T}}\boldsymbol{x} + \boldsymbol{b}_1 \in \mathbb{R}^{d_{\mathrm{ff}}}$ と置きます。$\boldsymbol{s} \in \{0, 1\}^{d_{\mathrm{ff}}}$（活性パターン）ごとに

$$R_{\boldsymbol{s}} := \big\{\boldsymbol{x} \in \mathbb{R}^d : s_k = 1 \text{ なら } a_k(\boldsymbol{x}) > 0,\ s_k = 0 \text{ なら } a_k(\boldsymbol{x}) \le 0 \quad (k = 1, \ldots, d_{\mathrm{ff}})\big\}$$

と定めます。各 $\boldsymbol{x}$ について、$a_k(\boldsymbol{x}) > 0$ かどうかで $\boldsymbol{s}$ がちょうど 1 つ決まるので、$R_{\boldsymbol{s}}$ は互いに交わらず、合わせると $\mathbb{R}^d$ 全体です。パターンは $2^{d_{\mathrm{ff}}}$ 通りなので、空でない $R_{\boldsymbol{s}}$ は高々 $2^{d_{\mathrm{ff}}}$ 個です。$a_k$ は 1 次式なので、$\{a_k > 0\}$ と $\{a_k \le 0\}$ は半空間で凸、その共通部分 $R_{\boldsymbol{s}}$ も凸です。

$\boldsymbol{x} \in R_{\boldsymbol{s}}$ のとき、$\mathrm{ReLU}(a_k(\boldsymbol{x})) = s_k a_k(\boldsymbol{x})$ です（$s_k = 1$ なら $a_k > 0$ でそのまま、$s_k = 0$ なら $a_k \le 0$ で 0）。$D_{\boldsymbol{s}} := \mathrm{diag}(\boldsymbol{s})$ と書けば

$$g(\boldsymbol{x}) = W_2^{\mathsf{T}} D_{\boldsymbol{s}}\big(W_1^{\mathsf{T}}\boldsymbol{x} + \boldsymbol{b}_1\big) + \boldsymbol{b}_2 \qquad (\boldsymbol{x} \in R_{\boldsymbol{s}})$$

で、右辺は $\boldsymbol{x}$ の 1 次式です。

2. 標準正規分布の密度を $\varphi(t) = \frac{1}{\sqrt{2\pi}}e^{-t^2/2}$ とすると、$\Phi(t) = \int_{-\infty}^{t}\varphi(u)\,du$ で、$\varphi$ は連続なので $\Phi' = \varphi$ です（任意の $a < t$ について $\Phi(t) = \Phi(a) + \int_a^t \varphi(u)\,du$ と書き、有界閉区間の上で 定理 5.2［積分の基本定理と定積分］ の (2) を使う）。$\varphi'(t) = -t\,\varphi(t)$ で、$\varphi$ は何回でも微分可能なので、$\Phi$ も $\mathrm{GELU}$ もそうです。微分すると

$$\mathrm{GELU}'(t) = \Phi(t) + t\,\varphi(t), \qquad \mathrm{GELU}''(t) = \varphi(t) + \varphi(t) + t\,\varphi'(t) = (2 - t^2)\,\varphi(t)$$

です。$\varphi(t) > 0$ なので、$\mathrm{GELU}''(t) = 0$ となるのは $t = \pm\sqrt2$ の 2 点だけです。ある開区間の上で $\mathrm{GELU}$ が 1 次式に一致すれば、その区間の上で 2 階微分は恒等的に 0 ですが、開区間は無限個の点を含むので矛盾します。

∎

$2^{d_{\mathrm{ff}}}$ は上界で、実際の個数はもっと少なくなります。たとえば入力が 2 次元で幅が 3 のとき、境界は平面上の 3 本の直線で、直線 3 本が平面を分ける領域は多くても 7 個です。9 節の実験でも、8 通りのパターンのうち現れるのは 7 通りでした。それでも幅を増やせば領域は増え、区分線形な関数で複雑な関数を細かく近似できるようになります。

注意 8.2（この章で扱わなかったこと）

• 万能近似。 十分に幅の広い 1 層の MLP が、有界閉集合の上の連続関数を任意の精度で近似できるという種類の定理があります。この章では証明も利用もしていません。
• ドロップアウトと初期化。 学習時に成分を確率的に 0 にする正則化や、初期値の分散の選び方は、学習の安定性に関わる実践上の工夫で、この章の命題には現れません。
• エンコーダ・デコーダ型。 Vaswani らのもとのモデルは、入力の列を処理するエンコーダと、出力の列を生成するデコーダからなります（[1] 3.1 節）。このテーマで扱うデコーダのみのモデルとの違いは次のとおりです。

	エンコーダの層	デコーダの層（エンコーダ・デコーダ型）	デコーダのみのモデルの層
自己注意のマスク	なし	因果マスク	因果マスク
交差注意	なし	あり（キーとバリューはエンコーダの出力）	なし
MLP	あり	あり	あり
1 層の部品の数	2	3	2

9. 最小実装と検算

デコーダのみの Transformer の順伝播を NumPy で実装します。自己注意は前章の attention.py の self_attention をそのまま使います。

# Python 3.12.14, NumPy 2.5.2。attention.py は第 3 章のコード
import math
import numpy as np
from attention import init_params, self_attention

def layer_norm(X, gamma, beta, eps=1e-5):
    mu = X.mean(axis=-1, keepdims=True)
    var = X.var(axis=-1, keepdims=True)          # 成分の分散（d で割る）
    return gamma * (X - mu) / np.sqrt(var + eps) + beta

def gelu(X):
    return X * 0.5 * (1.0 + np.vectorize(math.erf)(X / math.sqrt(2.0)))

def init_model(V, d, L, dff, Lctx, h, rng):
    model = {"wte": 0.5 * rng.standard_normal((V, d)),       # トークンの埋め込み（出力層と共有）
             "wpe": 0.5 * rng.standard_normal((Lctx, d)),    # 学習する位置埋め込み
             "blocks": [], "lnf": [np.ones(d), np.zeros(d)]}
    for _ in range(L):
        model["blocks"].append({
            "ln1": [np.ones(d), np.zeros(d)], "attn": init_params(d, h, rng),
            "ln2": [np.ones(d), np.zeros(d)],
            "W1": rng.standard_normal((d, dff)) / np.sqrt(d), "b1": 0.1 * rng.standard_normal(dff),
            "W2": rng.standard_normal((dff, d)) / np.sqrt(dff), "b2": 0.1 * rng.standard_normal(d)})
    return model

def forward(tokens, model):
    """トークン id の列（長さ n）からロジット (n, V) を返す。Pre-LN、因果マスクつき。"""
    n = len(tokens)
    X = model["wte"][tokens] + model["wpe"][:n]
    for b in model["blocks"]:
        X = X + self_attention(layer_norm(X, *b["ln1"]), b["attn"], causal=True)[0]
        X = X + gelu(layer_norm(X, *b["ln2"]) @ b["W1"] + b["b1"]) @ b["W2"] + b["b2"]
    return layer_norm(X, *model["lnf"]) @ model["wte"].T

def count(obj):
    if isinstance(obj, np.ndarray):
        return obj.size
    return sum(count(v) for v in (obj.values() if isinstance(obj, dict) else obj))

def formula(V, d, L, dff, Lctx):
    per_layer = (4 * d * d + 4 * d) + (2 * d * dff + dff + d) + 4 * d
    return V * d + Lctx * d + L * per_layer + 2 * d

if __name__ == "__main__":
    rng = np.random.default_rng(0)
    V, d, L, dff, Lctx, h = 10, 4, 2, 16, 8, 2
    model = init_model(V, d, L, dff, Lctx, h, rng)
    print("count:", count(model), " formula:", formula(V, d, L, dff, Lctx))

    tokens = np.array([3, 1, 4, 1, 5, 9])
    logits = forward(tokens, model)
    changed = tokens.copy(); changed[4:] = [2, 6]            # 5・6 番目のトークンを変える
    logits2 = forward(changed, model)
    print("logits shape:", logits.shape)
    print("rows 1-4 change: %.1e   rows 5-6 change: %.1e"
          % (np.abs(logits2[:4] - logits[:4]).max(), np.abs(logits2[4:] - logits[4:]).max()))
    pre = max(np.abs(forward(tokens[:t], model)[-1] - logits[t - 1]).max() for t in range(1, 7))
    print("prefix consistency: %.1e" % pre)

    V, d, L, Lctx = 50257, 768, 12, 1024
    per_layer = 12 * d * d + 13 * d
    print("embeddings:", V * d, " positions:", Lctx * d, " per layer:", per_layer,
          " layers:", L * per_layer, " final LN:", 2 * d)
    print("total:", formula(V, d, L, 4 * d, Lctx))

count: 568  formula: 568
logits shape: (6, 10)
rows 1-4 change: 0.0e+00   rows 5-6 change: 1.4e+00
prefix consistency: 4.4e-16
embeddings: 38597376  positions: 786432  per layer: 7087872  layers: 85054464  final LN: 1536
total: 124439808

• 1 行目：実装の配列の成分の総数と 定理 5.1 の式が、例 5.2 の 568 で一致します。
• 3 行目：5・6 番目のトークンを変えても、1〜4 番目の位置のロジットはまったく変わりません（命題 7.2 の 3）。
• 4 行目：接頭辞だけを入力したときの最後の位置のロジットが、全体を入力したときの対応する位置のロジットに丸め誤差の範囲で一致します。
• 5・6 行目：例 5.3 の表の値です。この構成は配列を実際には確保せず、式で数えています。

次に、層正規化の命題と、ReLU の活性パターンの個数、6 節の数値を確かめます。

# Python 3.12.14, NumPy 2.5.2。transformer.py は前のコード
import numpy as np
from transformer import layer_norm

np.set_printoptions(precision=4, suppress=True)
one, zero = np.ones(4), np.zeros(4)
x = np.array([1.0, 2.0, 3.0, 6.0])
y = layer_norm(x, one, zero, eps=0.0)
print("LN(x)      =", y, " mean %.1e  var %.4f  norm %.4f" % (y.mean(), y.var(), np.linalg.norm(y)))
print("LN(3x + 5) =", layer_norm(3 * x + 5, one, zero, eps=0.0))

def jacobian(x):
    d, sigma = len(x), x.std()
    y = (x - x.mean()) / sigma
    return (np.eye(d) - np.ones((d, d)) / d - np.outer(y, y) / d) / sigma

rng = np.random.default_rng(0)
x = rng.standard_normal(6)
J, h = jacobian(x), 1e-6
one, zero = np.ones(6), np.zeros(6)
Jnum = np.stack([(layer_norm(x + h * e, one, zero, 0.0) - layer_norm(x - h * e, one, zero, 0.0)) / (2 * h)
                 for e in np.eye(6)], axis=1)
print("formula vs finite difference: %.1e" % np.abs(J - Jnum).max())
print("|J 1| = %.1e   |J (x - mu)| = %.1e" % (np.abs(J @ one).max(), np.abs(J @ (x - x.mean())).max()))
print("singular values * sigma:", np.linalg.svd(J, compute_uv=False) * x.std())
print("|J(3x) - J(x)/3| = %.1e" % np.abs(jacobian(3 * x) - J / 3).max())

# ReLU の MLP の線形領域：入力 2 次元・幅 3 で、活性パターンを数える
W1 = np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]); b1 = np.array([0.0, 0.0, -1.0])
g = np.linspace(-3, 3, 601)
pts = np.stack(np.meshgrid(g, g), axis=-1).reshape(-1, 2)
patterns = {"".join(map(str, row)) for row in (pts @ W1.T + b1 > 0).astype(int)}
print("activation patterns:", len(patterns), "of", 2 ** 3, sorted(patterns))

d, n = 768, 1024
print("MLP+proj 12nd^2 = %d   attention 2n^2d = %d   ratio = %.3f   n/(6d) = %.3f"
      % (12 * n * d * d, 2 * n * n * d, 2 * n * n * d / (12 * n * d * d), n / (6 * d)))

LN(x)      = [-1.069  -0.5345  0.      1.6036]  mean 0.0e+00  var 1.0000  norm 2.0000
LN(3x + 5) = [-1.069  -0.5345  0.      1.6036]
formula vs finite difference: 2.2e-10
|J 1| = 4.4e-16   |J (x - mu)| = 2.2e-16
singular values * sigma: [1. 1. 1. 1. 0. 0.]
|J(3x) - J(x)/3| = 1.9e-16
activation patterns: 7 of 8 ['000', '010', '011', '100', '101', '110', '111']
MLP+proj 12nd^2 = 7247757312   attention 2n^2d = 1610612736   ratio = 0.222   n/(6d) = 0.222

最初の 2 行が 例 3.2、3〜6 行目が 命題 3.3 です。ヤコビ行列の式は中心差分と一致し、$\boldsymbol{1}$ と $\boldsymbol{x} - \mu\boldsymbol{1}$ を 0 に写します。$\sigma J$ の特異値は 1 が 4 個、0 が 2 個で、$d = 6$ 次元のうち 2 方向を消す直交射影であることが見えています。7 行目では、3 つの境界 $x_1 = 0$、$x_2 = 0$、$x_1 + x_2 = 1$ に対し、「$x_1 \le 0$ かつ $x_2 \le 0$ かつ $x_1 + x_2 > 1$」にあたるパターン 001 だけが現れません。

10. 次章へ

この章で、トークン id の列からロジットまでの関数が 1 つ定まりました。パラメータは 定理 5.1 の個数だけあり、まだ乱数で初期化しただけです。次章では、このパラメータを「次のトークンの予測」の損失で学習します。そこでは、命題 7.2 の 3 によって 1 本の列から $n$ 個の予測の損失が同時に得られることと、命題 6.1 の計算量が、学習の 1 ステップの費用を決めます。第 6 章の生成では、同じ因果性から、すでに計算した位置の結果を使い回せることを示します。

11. 演習

演習 11.1易

$|V| = 1000$、$d = 64$、$L = 4$、$d_{\mathrm{ff}} = 4d$、$L_{\mathrm{ctx}} = 128$ のデコーダのみの Transformer（重み共有あり）のパラメータ数を求めよ。そのうち埋め込み行列の占める割合はいくらか。

解答

定理 5.1 より、埋め込み $1000 \times 64 = 64{,}000$、位置 $128 \times 64 = 8{,}192$、1 層 $12\times64^2 + 13\times64 = 49{,}152 + 832 = 49{,}984$、4 層で $199{,}936$、最後の層正規化 $128$。合計は $64{,}000 + 8{,}192 + 199{,}936 + 128 = 272{,}256$ です。埋め込みの割合は $64{,}000/272{,}256 \approx 23.5\%$ です。

演習 11.2標準

$\varepsilon = 0$ とし、$\boldsymbol{\gamma} = g\boldsymbol{1}$（$g > 0$ は定数）、$\boldsymbol{\beta} = b\boldsymbol{1}$ の層正規化を考える。出力の成分の平均、分散、ノルムを求めよ。また、$\boldsymbol{\gamma}$ が一般のベクトルのとき、命題 3.1 の 1 は成り立つか。

解答

$\hat{\boldsymbol{y}} := (\boldsymbol{x} - \mu\boldsymbol{1})/\sigma$ と書くと、出力は $g\hat{\boldsymbol{y}} + b\boldsymbol{1}$ です。命題 3.1 の 2 より $\mu(\hat{\boldsymbol{y}}) = 0$、$\sigma^2(\hat{\boldsymbol{y}}) = 1$ なので、同じ命題の 1 の証明の計算（$a = g$）から、平均は $g\cdot0 + b = b$、分散は $g^2$ です。ノルムの 2 乗は $\sum_j (g\hat y_j + b)^2 = g^2\sum_j\hat y_j^2 + 2gb\sum_j\hat y_j + db^2 = g^2d + db^2$ なので、ノルムは $\sqrt{d(g^2 + b^2)}$ です。

1 は一般の $\boldsymbol{\gamma}, \boldsymbol{\beta}$ でも成り立ちます。証明で示したのは $\hat{\boldsymbol{y}}$ が $a\boldsymbol{x} + b'\boldsymbol{1}$ に対しても同じになることで、そのあとに施す $\boldsymbol{\gamma}\odot$ と $+\boldsymbol{\beta}$ は入力に依らない同じ操作だからです。成り立たなくなるのは 2 のほうで、$\boldsymbol{\gamma}$ が一般のとき出力の平均が 0、分散が 1 とは限りません。

演習 11.3標準

命題 4.1 の Pre-LN の場合で $M = 3$、$l = 0$ とする。$\partial\boldsymbol{x}_3/\partial\boldsymbol{x}_0$ を展開した 8 個の項をすべて書け。残差接続が無い（$\boldsymbol{x}_l = G_l(\boldsymbol{x}_{l-1})$）場合には、このうちどの項だけが残るか。

解答

$(I + J_3)(I + J_2)(I + J_1)$ を展開すると

$$I + J_1 + J_2 + J_3 + J_2J_1 + J_3J_1 + J_3J_2 + J_3J_2J_1$$

の 8 項です（$S$ が $\emptyset$、1 元集合 3 つ、2 元集合 3 つ、全体）。各項は「どの層を通り、どの層を迂回したか」に対応します。残差接続が無い場合のヤコビ行列は連鎖律より $J_3J_2J_1$ で、すべての層を通る最後の 1 項だけです（この場合の $J_l$ は $G_l$ の $\boldsymbol{x}_{l-1}$ でのヤコビ行列で、$\boldsymbol{x}_l$ の値そのものも残差ありの場合とは違います）。

演習 11.4易

$d_{\mathrm{ff}} = 4d$ のブロックで、乗算の回数の主要項 $12nd^2 + 2n^2d$ について、$n$ を 2 倍にしたときの比を $n = d$、$n = 6d$、$n = 60d$ のそれぞれで求めよ。

解答

$T(n) = 12nd^2 + 2n^2d$ とすると $T(2n)/T(n) = (24nd^2 + 8n^2d)/(12nd^2 + 2n^2d) = (12d + 4n)/(6d + n)$ です。

• $n = d$：$16d/7d \approx 2.29$。
• $n = 6d$：$36d/12d = 3$。
• $n = 60d$：$252d/66d \approx 3.82$。

列が短いあいだは 2 倍に近く（$n$ に比例する項が支配的）、長くなるにつれて 4 倍に近づきます（$n^2$ に比例する項が支配的）。命題 6.1 の $n = 6d$ がちょうどその中間です。

参考文献

• [1] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, I. Polosukhin, “Attention Is All You Need”, 2017. arXiv:1706.03762 — 3 節（モデルの構成。3.1 節が残差接続と層正規化の配置、3.3 節が位置ごとの MLP）。
• [2] J. L. Ba, J. R. Kiros, G. E. Hinton, “Layer Normalization”, 2016. arXiv:1607.06450 — 層正規化の提案。
• [3] K. He, X. Zhang, S. Ren, J. Sun, “Deep Residual Learning for Image Recognition”, 2015. arXiv:1512.03385 — 残差接続の提案。
• [4] K. He, X. Zhang, S. Ren, J. Sun, “Identity Mappings in Deep Residual Networks”, ECCV 2016. arXiv:1603.05027 — 残差の経路を恒等写像に保つ配置。
• [5] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, “Language Models are Unsupervised Multitask Learners”, OpenAI, 2019 — 2.3 節（層正規化の配置、語彙の大きさ、文脈長）、表 2（4 つの大きさの構成）。
• [6] R. Xiong, Y. Yang, D. He, K. Zheng, S. Zheng, C. Xing, H. Zhang, Y. Lan, L. Wang, T.-Y. Liu, “On Layer Normalization in the Transformer Architecture”, ICML 2020. arXiv:2002.04745 — Pre-LN と Post-LN の比較。
• [7] D. Hendrycks, K. Gimpel, “Gaussian Error Linear Units (GELUs)”, 2016. arXiv:1606.08415 — GELU の定義。
• [8] OpenAI, “gpt-2”（[5] のコードと学習済みモデルの公開リポジトリ）. https://github.com/openai/gpt-2 — README のパラメータ数についての注記と、src/model.py のモデルの定義。