Transformer は Pre-LN 構成が安定することを解析・検証した研究 [1] の主定理 (定理 1) は、初期化直後の最終エンコーダ層内の FFN 層の 2 番目の線形変換の重み $W^{2,L}$ の勾配の上界のオーダーが、Post-LN では $\mathcal{O}(d\sqrt{\ln d})$、Pre-LN では $\mathcal{O}(d\sqrt{\ln d / L})$ になることを示しています ($L$ はエンコーダ層数、$d$ はトークンの埋め込み次元)。Pre-LN は因子 $\sqrt{1/L}$ の寄与でこの上界が小さくなり、そのため「Pre-LN はウォームアップなしでも安定して学習できる」と主張されています。
この記事では、別記事で追った補題 2 (中間出力のノルムのスケール) と補題 3 (LayerNorm のヤコビ行列のノルムのオーダー) を踏まえて定理 1を追います。
参考文献
- Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie-Yan Liu. On Layer Normalization in the Transformer Architecture. Proceedings of the 37th International Conference on Machine Learning (ICML 2020), 2020.
https://proceedings.mlr.press/v119/xiong20b.html
関連記事：
- 【文献メモ】現代 LLM の正規化 Pre-LN + RMSNorm の成り立ち ： [1] の概要や、よく併用する手法の概要や、その他の発展や、トイコードがあります。
- LayerNorm の入力に関するヤコビ行列のスペクトルノルムのオーダー ： 補題 3 (LayerNorm のヤコビ行列のノルムのオーダー) を追いました。 $\|\mathbf{J}_{LN}(x)\|_{2}=\mathcal{O}(\dfrac{\sqrt{d}}{\|x\|_{2}})$ になるとわかりました。
- 初期化直後の Transformer における中間出力のノルムのスケール ： 補題 2 (初期化直後の中間出力のノルムのスケール) を追いました。Post-LN Transformer は任意の $l, i$ で $\mathbb{E}(\|x_{l,i}^{{\rm post},5}\|_{2}^{2})=\dfrac{3}{2}d$ を満たし、Pre-LN Transformer は任意の $l, i$ で $(1+\dfrac{l}{2})d\leqq\mathbb{E}(\|x^{{\rm pre}}_{l,i}\|_{2}^{2})\leqq(1+\dfrac{3l}{2})d$ を満たすとわかりました。
以下、ベクトルは横ベクトルとします (入力を重みに左から掛ける流儀 $y=xW$)。また、パラメータの初期化については、原則として Xavier 初期化でバイアスはゼロ、ただし LayerNorm は $\gamma=1,$ $\beta=0$ とし、$W^{Q,l}, W^{K,l}$ はゼロ初期化とします。
表記法 (一部再掲)
補題 2 の記事と同様に、初期化直後の $l$ 番目のエンコーダ層の内部で、第 $i$ トークンがたどる中間出力を次のように書きます。
- Post-LN Transformer では、入力 $x^{{\rm post}}_{l,i}$ から順に以下をたどります。
- $x^{{\rm post},1}_{l,i}$：MHA 後
- $x^{{\rm post},2}_{l,i}$：1 回目の残差接続後
- $x^{{\rm post},3}_{l,i}$：1 回目の LayerNorm 後
- $x^{{\rm post},4}_{l,i}$：FFN 後
- $x^{{\rm post},5}_{l,i}$：2 回目の残差接続後
- $x^{{\rm post}}_{l+1,i}$：2 回目の LayerNorm 後 (次のエンコーダ層への入力)
- Pre-LN Transformer では、入力 $x^{{\rm pre}}_{l,i}$ から順に以下をたどります。
- $x^{{\rm pre},1}_{l,i}$：1 回目の LayerNorm 後
- $x^{{\rm pre},2}_{l,i}$：MHA 後
- $x^{{\rm pre},3}_{l,i}$：1 回目の残差接続後
- $x^{{\rm pre},4}_{l,i}$：2 回目の LayerNorm 後
- $x^{{\rm pre},5}_{l,i}$：FFN 後
- $x^{{\rm pre}}_{l+1,i}$：2 回目の残差接続後 (次のエンコーダ層への入力)
その他、以下の表記法を導入します。
- $\mathcal{L}(x)=-\log\text{softmax}_{y_i}(xW^{{\rm emb}})$ ： 1 トークン分の損失
- $\tilde{\mathcal{L}}=\dfrac{1}{n}\sum_{i=1}^{n}\mathcal{L}$ ： 系列全体の損失 ($n$ は 1 系列あたりのトークン数)
- $\text{FFN}(x)=(\text{ReLU}(xW^{1,l}))W^{2,l}$ ： $l$ 層目の FFN サブ層 ($W^{1,l}, W^{2,l}$ は 1, 2 番目の線形変換の重み)
この記事でいう「最終層の重み」とは、最終 (第 $L$) エンコーダ層の $W^{2,L}$ を指します。
(ε, δ)-有界の定義
最終層の出力のノルムが「ほとんどいつも期待値の定数倍で上から抑えられる」ことを表す概念を用意します。
定義 1. 確率変数 $Z\geqq 0$ が、少なくとも確率 $1-\delta$ で $\dfrac{Z-\mathbb{E}Z}{\mathbb{E}Z}\leqq\epsilon$ を満たすとき、$Z$ は $(\epsilon, \delta)$-有界であるという ($\epsilon>0,\ 0<\delta<1$)。
これは言い換えると、確率 $1-\delta$ で上界 $Z\leqq(1+\epsilon)\mathbb{E}Z$ が成り立つということです。下界への条件はありません (証明で必要になる下界は補足1で別途導きます)。
定理 1
任意の $i$ で $\|x^{{\rm post},5}_{L,i}\|_{2}^{2}$ および $\|x^{{\rm pre}}_{L+1,i}\|_{2}^{2}$ が $(\epsilon, \delta)$-有界であるとする ($\epsilon,\ \delta=\delta(\epsilon)$ は小さい数)。このとき、少なくとも確率 $0.99-\delta-\dfrac{\epsilon}{0.9+\epsilon}$ で、Post-LN Transformer の最終層の重みの勾配は、以下を満たす。
\left\|\frac{\partial\tilde{\mathcal{L}}}{\partial W^{2,L}}\right\|_{F}\leqq\mathcal{O}(d\sqrt{\ln d})
Pre-LN Transformer の最終層の重みの勾配は、以下を満たす。
\left\|\frac{\partial\tilde{\mathcal{L}}}{\partial W^{2,L}}\right\|_{F}\leqq\mathcal{O}\left(d\sqrt{\frac{\ln d}{L}}\right)
定理 1 の証明
以下の証明は FFN の中間層の次元を $d_F=d$ として書いています ($W^{1,l}, W^{2,l}$ をいずれも $d\times d$ 行列とみなしています)。実際によく使われる $d_F=4d$ でも以下の議論は成り立ちます。変わるのは、フロベニウスノルムをとる際の成分数が $d^2$ 個から $4d^2$ 個になる点、補足2 で劣加法性を使う事象が $d$ 個から $4d$ 個になる点、Xavier 初期化での $W^{1,L}$ の各要素の分散が $\dfrac{1}{d}$ から $\dfrac{2}{5d}$ になり $[x^{{\rm post},3}_{L,i}W^{1,L}]_{p}$ が $N(0,1)$ でなく $N(0,\frac{2}{5})$ になる点の 3 つで、いずれも定数倍の違いなので $\mathcal{O}$ に吸収され、結論のオーダーは変わりません。
Post-LN の場合
損失は $\tilde{\mathcal{L}}=\dfrac{1}{n}\sum_{i=1}^{n}\mathcal{L}(x^{{\rm post}}_{L+1,i})$ である。まず 1 トークン分 $\mathcal{L}(x^{{\rm post}}_{L+1,i})$ の、$W^{2,L}$ の 1 成分 $W^{2,L}_{pq}$ ($p,q\in{1,\dots,d}$) に関する勾配を評価する。連鎖律より、以下となる。
\frac{\partial\mathcal{L}(x_{L+1,i}^{{\rm post}})}{\partial W^{2,L}_{pq}}
=\frac{\partial\mathcal{L}(x_{L+1,i}^{{\rm post}})}{\partial x_{L+1,i}^{{\rm post}}}\frac{\partial x_{L+1,i}^{{\rm post}}}{\partial x_{L,i}^{{\rm post},5}}\frac{\partial x_{L,i}^{{\rm post},5}}{\partial x_{L,i}^{{\rm post},4}}\frac{\partial x_{L,i}^{{\rm post},4}}{\partial W^{2,L}_{pq}}
ここで、$x^{{\rm post}}_{L+1,i}=\text{LN}(x^{{\rm post},5}_{L,i})$ なので $\dfrac{\partial x^{{\rm post}}_{L+1,i}}{\partial x^{{\rm post},5}_{L,i}}=\mathbf{J}_{LN}(x^{{\rm post},5}_{L,i})$ であり、$x^{{\rm post},5}_{L,i}=x^{{\rm post},3}_{L,i}+x^{{\rm post},4}_{L,i}$ なので $\dfrac{\partial x^{{\rm post},5}_{L,i}}{\partial x^{{\rm post},4}_{L,i}}=I$ である。また、$x^{{\rm post},4}_{L,i}=\text{ReLU}(x^{{\rm post},3}_{L,i}W^{1,L})W^{2,L}$ なので、この $W^{2,L}_{pq}$ に関する偏微分は第 $q$ 成分だけが $[\text{ReLU}(x^{{\rm post},3}_{L,i}W^{1,L})]_{p}$ で他が $0$ のベクトルになる。よって、以下となる。
\frac{\partial\mathcal{L}(x_{L+1,i}^{{\rm post}})}{\partial W^{2,L}_{pq}}
=\frac{\partial\mathcal{L}(x_{L+1,i}^{{\rm post}})}{\partial x_{L+1,i}^{{\rm post}}}\mathbf{J}_{LN}(x_{L,i}^{{\rm post},5})(0,\dots,[\text{ReLU}(x_{L,i}^{{\rm post},3}W^{1,L})]_{p},\dots,0)^{\top}
したがって、コーシー・シュワルツの不等式と作用素ノルムの定義から、以下となる。
\left|\frac{\partial\mathcal{L}(x_{L+1,i}^{{\rm post}})}{\partial W^{2,L}_{pq}}\right|
\leqq\left\|\frac{\partial\mathcal{L}(x_{L+1,i}^{{\rm post}})}{\partial x_{L+1,i}^{{\rm post}}}\right\|_{2}\|\mathbf{J}_{LN}(x_{L,i}^{{\rm post},5})\|_{2}\,\bigl|[\text{ReLU}(x_{L,i}^{{\rm post},3}W^{1,L})]_{p}\bigr|
右辺の右端の絶対値は、第 $q$ 成分だけ非ゼロなベクトルのノルムなのでその成分の絶対値である。ここで、損失の各微分は有界なので $\|\frac{\partial\mathcal{L}}{\partial x^{{\rm post}}_{L+1,i}}\|_{2}^{2}=\mathcal{O}(1)$ となり、以下となる。
\left|\frac{\partial\mathcal{L}(x_{L+1,i}^{{\rm post}})}{\partial W^{2,L}_{pq}}\right|^{2}
=\mathcal{O}\left(\|\mathbf{J}_{LN}(x_{L,i}^{{\rm post},5})\|^{2}_{2}\,\bigl|[\text{ReLU}(x_{L,i}^{{\rm post},3}W^{1,L})]_{p}\bigr|^{2}\right)
右辺の因子を評価する。
- $\|\mathbf{J}_{LN}(x^{{\rm post},5}_{L,i})\|_{2}^{2}$ は、補題 3 より $\mathcal{O}(d/\|x^{{\rm post},5}_{L,i}\|_{2}^{2})$。
- $\bigl|[\text{ReLU}(x^{{\rm post},3}_{L,i}W^{1,L})]_{p}\bigr|^{2}$ は、$\|x^{{\rm post},3}_{L,i}\|_{2}^{2}=d$ かつ $W^{1,L}$ の各要素が $N(0,1/d)$ なので $[x^{{\rm post},3}_{L,i}W^{1,L}]_{p}\sim N(0,1)$。したがって、確率 $0.99$ 以上で、すべての $p=1,\dots,d$ について $\text{ReLU}([x^{{\rm post},3}_{L,i}W^{1,L}]_{p})^{2}\leqq 2\ln 100d$ が成り立つ (補足 2)。
さらに分母の $\|x^{{\rm post},5}_{L,i}\|_{2}^{2}$ を下から抑える。仮定よりこれは $(\epsilon, \delta)$-有界なので、確率 $1-\delta(\epsilon)$ で上界 $\|x^{{\rm post},5}_{L,i}\|_{2}^{2}\leqq(1+\epsilon)\mathbb{E}\|x^{{\rm post},5}_{L,i}\|_{2}^{2}$ が成り立つ。この上界と、上に有界な確率変数のマルコフの不等式 (補足 1) から、$\alpha_0>0$ に対して以下を得る。
\operatorname{Pr}\left[\|x^{{\rm post},5}_{L,i}\|_{2}^{2}\geqq\alpha_{0}\mathbb{E}\|x^{{\rm post},5}_{L,i}\|_{2}^{2}\right]\geqq 1-\frac{\epsilon}{1+\epsilon-\alpha_{0}}
以上より、確率 $0.99-\delta(\epsilon)-\dfrac{\epsilon}{1+\epsilon-\alpha_0}$ 以上で、以下となる。
\left|\frac{\partial\mathcal{L}(x_{L+1,i}^{{\rm post}})}{\partial W^{2,L}_{pq}}\right|^{2}
\leqq\mathcal{O}\left(\frac{2d\ln 100d}{\|x_{L,i}^{{\rm post},5}\|_{2}^{2}}\right)
\leqq\mathcal{O}\left(\frac{d\ln d}{\alpha_{0}\mathbb{E}\|x_{L,i}^{{\rm post},5}\|_{2}^{2}}\right)
=\mathcal{O}\left(\frac{\ln d}{\alpha_{0}}\right)
ただし、補題 2 の $\mathbb{E}\|x^{{\rm post},5}_{L,i}\|_{2}^{2}=\dfrac{3}{2}d$ より $d/\mathbb{E}\|x^{{\rm post},5}_{L,i}\|_{2}^{2}=\mathcal{O}(1)$ となることを使った。
これを全トークン平均 $\tilde{\mathcal{L}}$ に戻す。$|\cdot|^2$ の凸性より、以下となる。
\left|\frac{\partial\tilde{\mathcal{L}}}{\partial W^{2,L}_{pq}}\right|^{2}=\left|\frac{1}{n}\sum_{i=1}^{n}\frac{\partial\mathcal{L}(x_{L+1,i}^{{\rm post}})}{\partial W^{2,L}_{pq}}\right|^{2}\leqq\frac{1}{n}\sum_{i=1}^{n}\left|\frac{\partial\mathcal{L}(x_{L+1,i}^{{\rm post}})}{\partial W^{2,L}_{pq}}\right|^{2}=\mathcal{O}\left(\frac{\ln d}{\alpha_{0}}\right)
したがって、$W^{2,L}$ の全成分の和は以下となる。
\left\|\frac{\partial\tilde{\mathcal{L}}}{\partial W^{2,L}}\right\|_{F}=\sqrt{\sum_{p,q=1}^{d}\left|\frac{\partial\tilde{\mathcal{L}}}{\partial W^{2,L}_{pq}}\right|^{2}}=\mathcal{O}\left(\sqrt{\frac{d^{2}\ln d}{\alpha_{0}}}\right)
上式で $\alpha_0=\dfrac{1}{10}$ とおけば、確率 $0.99-\delta(\epsilon)-\dfrac{\epsilon}{0.9+\epsilon}$ 以上で $\left\|\dfrac{\partial\tilde{\mathcal{L}}}{\partial W^{2,L}}\right\|_{F}\leqq\mathcal{O}(d\sqrt{\ln d})$ となる。
Pre-LN の場合
手順は Post-LN とほぼ同じで、違いは次の 2 点だけである。
(1) 損失の直前に最終 LayerNorm が入る ： Pre-LN では損失の直前に最終 LayerNorm が入り、損失は $\mathcal{L}(x^{\rm pre}_{{\rm Final},i})$ となる。連鎖律は以下となる。
\frac{\partial\mathcal{L}(x_{{\rm Final},i}^{{\rm pre}})}{\partial W^{2,L}_{pq}}
=\frac{\partial\mathcal{L}(x_{{\rm Final},i}^{{\rm pre}})}{\partial x_{{\rm Final},i}^{{\rm pre}}}\mathbf{J}_{LN}(x_{L+1,i}^{{\rm pre}})(0,\dots,[\text{ReLU}(x_{L,i}^{{\rm pre},4}W^{1,L})]_{p},\dots,0)^{\top}
$\mathbf{J}_{LN}$ の引数 $x^{\rm pre}_{L+1,i}$ はすなわち最終エンコーダ層の 2 回目の残差接続後の出力であるが、Post-LN のときの引数 $x^{{\rm post},5}_{L,i}$ も最終エンコーダ層の 2 回目の残差接続後の出力なので、以降は Post-LN と同様に評価できる。
(2) 分母のノルムが層とともに大きくなる ： 補題 2 より $\mathbb{E}\|x^{{\rm pre}}_{L+1,i}\|_{2}^{2}$ のオーダーは $\Theta(Ld)$ であり、$d/\mathbb{E}\|x^{{\rm pre}}_{L+1,i}\|_{2}^{2}=\mathcal{O}(1/L)$ となる。よって Post-LN の評価がそのまま $1/L$ 倍され、以下となる。
\left|\frac{\partial\mathcal{L}(x_{{\rm Final},i}^{{\rm pre}})}{\partial W^{2,L}_{pq}}\right|^{2}
\leqq\mathcal{O}\left(\frac{d\ln d}{\alpha_{0}\mathbb{E}\|x_{L+1,i}^{{\rm pre}}\|_{2}^{2}}\right)=\mathcal{O}\left(\frac{\ln d}{\alpha_{0}L}\right)
あとは Post-LN と同様に全成分の和を取ると $\mathcal{O}\left(\sqrt{\dfrac{d^{2}\ln d}{\alpha_0 L}}\right)$ となり、$\alpha_0=\dfrac{1}{10}$ とおけば確率 $0.99-\delta(\epsilon)-\dfrac{\epsilon}{0.9+\epsilon}$ 以上で $\left\|\dfrac{\partial\tilde{\mathcal{L}}}{\partial W^{2,L}}\right\|_{F}\leqq\mathcal{O}\left(d\sqrt{\dfrac{\ln d}{L}}\right)$ となる。
補足を書きます。
補足 1
上に有界な確率変数についてのマルコフの不等式を示します。確率変数 $Y$ が確率 1 で $Y\leqq B$ を満たすとき、$X=B-Y\geqq 0$ にマルコフの不等式を適用すると、$a<B$ に対して以下が成り立ちます。
\operatorname{Pr}[Y\leqq a]=\operatorname{Pr}[X\geqq B-a]\leqq\frac{\mathbb{E}[X]}{B-a}=\frac{\mathbb{E}[B-Y]}{B-a}
ここに、$(\epsilon, \delta)$-有界であることから得た上界 $B=(1+\epsilon)\mathbb{E}\|x^{{\rm post},5}_{L,i}\|_{2}^{2}$ と、$Y=\|x^{{\rm post},5}_{L,i}\|_{2}^{2}$、$a=\alpha_0\mathbb{E}\|x^{{\rm post},5}_{L,i}\|_{2}^{2}$ ($\alpha_0<1+\epsilon$ とします) を代入します。$\mathbb{E}[B-Y]=B-\mathbb{E}Y=\epsilon,\mathbb{E}\|x^{{\rm post},5}_{L,i}\|_{2}^{2}$ なので、以下となります。
\operatorname{Pr}\left[\|x^{{\rm post},5}_{L,i}\|_{2}^{2}\leqq\alpha_{0}\mathbb{E}\|x^{{\rm post},5}_{L,i}\|_{2}^{2}\right]\leqq\frac{\epsilon\,\mathbb{E}\|x^{{\rm post},5}_{L,i}\|_{2}^{2}}{(1+\epsilon-\alpha_{0})\mathbb{E}\|x^{{\rm post},5}_{L,i}\|_{2}^{2}}=\frac{\epsilon}{1+\epsilon-\alpha_{0}}
よって、$\operatorname{Pr}[\|x^{{\rm post},5}_{L,i}\|_{2}^{2}\geqq\alpha_0\mathbb{E}\|x^{{\rm post},5}_{L,i}\|_{2}^{2}]\geqq 1-\dfrac{\epsilon}{1+\epsilon-\alpha_0}$ を得ます。
なお、マルコフの不等式は $X=B-Y\geqq 0$ が確率 1 で成り立つことを要求しますが、ここでの $B$ は確率 $1-\delta(\epsilon)$ でしか成り立ちません。$X$ が負になりうる以上マルコフの不等式は使えず、$X\geqq 0$ の下では $\mathbb{E}[B-Y]$ が $\epsilon,\mathbb{E}\|x^{{\rm post},5}_{L,i}\|_{2}^{2}$ から変わってしまうため、$\delta(\epsilon)$ を差し引くだけでは正当化されません。原論文もこの箇所は同じ書き方をしているため、ここでは原論文にならい、$\delta(\epsilon)$ を差し引く形でまとめておきます。
補足 2
標準正規分布にしたがう確率変数 $Z\sim N(0,1)$ は $\operatorname{Pr}[|Z|\geqq a_0]\leqq\exp(-a_0^{2}/2)$ を満たします。$[x^{{\rm post},3}_{L,i}W^{1,L}]_{p}\sim N(0,1)$ なので、$a_0=\sqrt{2\ln 100d}$ とおくと以下となります。
\operatorname{Pr}\left[\bigl|[x^{{\rm post},3}_{L,i}W^{1,L}]_{p}\bigr|\geqq\sqrt{2\ln 100d}\right]\leqq\exp(-\ln 100d)=\frac{0.01}{d}
また、$\text{ReLU}(z)=\max(z,0)$ なので、$\text{ReLU}(z)^{2}\geqq 2\ln 100d$ となるのは $z\geqq\sqrt{2\ln 100d}$ のときに限られます。これは $|z|\geqq\sqrt{2\ln 100d}$ の片側にあたるので、以下となります。
\operatorname{Pr}\left[\text{ReLU}([x^{{\rm post},3}_{L,i}W^{1,L}]_{p})^{2}\geqq 2\ln 100d\right]\leqq\frac{0.005}{d}\leqq\frac{0.01}{d}
$p=1,\dots,d$ の少なくとも 1 つで不等式が破れる確率は、劣加法性より各 $p$ での確率の和 $d\cdot\dfrac{0.01}{d}=0.01$ 以下です。よって、$d$ 個すべてが $\text{ReLU}([x^{{\rm post},3}_{L,i}W^{1,L}]_{p})^{2}\leqq 2\ln 100d$ を満たす確率は $0.99$ 以上になります。
なお、本文では固定した $i$ について確率 $0.99$ を得ていますが、$\tilde{\mathcal{L}}$ の評価では $i=1,\dots,n$ の全てで同時に成り立つ必要があります。$100d$ を $100nd$ に取り替えて $n\times d$ 個の事象に劣加法性を使えば同じ議論ができ、$n$ が定数なら $\ln 100nd=\mathcal{O}(\ln d)$ なのでオーダーは変わりません。