# 大規模言語モデル事前学習の安定化

高瀬翔 1,2 清野舜 1,2 小林颯介 ${ }^{3}$ 鈴木潤 3
${ }^{1}$ LINE ヤフー株式会社 ${ }^{2}$ SB Intuitions 株式会社 3 東北大学
\{sho. takase, shun.kiyono\}@lycorp.co.jp
sosk@preferred.jp jun.suzuki@tohoku.ac.jp

## 概要

大規模言語モデルの事前学習では損失スパイクがしばしば起きることが知られている。損失スパイクはモデルの性能を低下させ，学習が失敗してしまうこともある.この損失スパイクの原因を探るため，本研究では中間層の勾配に着目する. 理論的な分析を通じて，大規模言語モデルの事前学習では LN 層が勾配爆発を引き起こすこと，および，その解決法を示す. 実験を通じて LN 層の勾配爆発を抑制することで損失スパイクも抑制されることを示す。

## 1 はじめに

GPT をはじめ大規模言語モデル（Large Language Model，以降本稿では LLM と記す）は様々なアプリケーションの根幹となっている $[1,2,3]$. LLM は大規模な訓練データで事前学習した大量のパラメー タを持つニューラル言語モデルであり，性能はパラメータ数と訓練データ量に対数比例することが知られている [4]. 事前学習は大量の計算資源を必要とするため学習が失敗するリスクを極力低減することが求められている.

Transformer [5] は多くの LLM に採用されているニューラルモデルであるが，その挙動の包括的な理論的理解には至っていない. 例えば，図 1 の Vanilla のように，Transformerを用いた LLM の事前学習では損失関数の値が跳放上がり（損失スパイク），また，しばしば発散してしまう現象が知られているが [2]，なぜ発生するのかは明らかにされていない.損失スパイクへの対処として様々な手法が提案されているが $[2,6,7]$, これらの理論的な正当化もなされていないため，既存研究の報告と異なるパラメー タ数のモデルなど，状況が異なる場合でも適用できるかには疑問が残る。

本研究では LLM の事前学習における損失スパイクの理論的な分析を行う. 特に初期化手法が学習に

図 117 億パラメータの，Transformer を用いた LLM の事前学習における訓練データでの損失関数の值の例. Vanilla は学習初期の段階で損失スパイクが何度か発生し, 13000 ステップで発散してしまっている. 開発データでの損失関数の値は付録 A の図 4 (a) に示す.

与える影響について議論する。Transformer を LLM に用いる場合，Self-attention 層と FFN 層の出力べクトルの標準偏差が小さくなる初期化を行うと学習が安定することが経験的にも理論的にも示されている $[8,6,9,10]$. 一方で，本研究では，このような初期化を行った場合には Layer Normalization (LN)層 [11] が勾配爆発を引き起こしてしまうことを示す.さらに，LN 層での勾配爆発を防ぐための要件を示し，モデルの簡易な修正でこの要件を達成できることを示す.

理論的な分析を検証するため LLM の事前学習の実験を行う．上記の，LN 層での勾配爆発を防ぐ要件を満たしている場合には損失スパイクが発生せず，満たしていない場合には，仮に既存研究で損失スパイクへの対処として提案された手法であっても損失スパイクが発生してしまうことを示す.

## 2 議論の準備

## 2.1 Pre-LN Transformer

本研究では，GPT $[12,13,1]$ などでも用いられている，Pre-LN Transformer [14] を対象に分析を行う. Pre-LN Transformer はオリジナルの Transformer より

も学習が安定することが示されている $[14,15,16]$. Transformer のある層への入力を $x \in \mathbb{R}^{d}$ としたとき， その層の出力 $y$ は以下の式で得られる：

$
\begin{aligned}
y & =x^{\prime}+\operatorname{FFN}\left(\operatorname{LN}\left(x^{\prime}\right)\right) \\
x^{\prime} & =x+\operatorname{Attention}(\operatorname{LN}(x))
\end{aligned}
$

以降では，式 (1)，(2)の第 1 項，すなわち， $x$ と $x^{\prime} を$ residual branch と呼ぶ. また, FFN と Attention の定義は以下のとおりである1):

$
\begin{aligned}
\operatorname{FFN}(x) & =\mathscr{F}\left(x W_{1}\right) W_{2}, \\
\operatorname{Attention}(x) & =\operatorname{concat}\left(\operatorname{head}_{1}(x), \ldots, \operatorname{head}_{h}(x)\right) W_{O}, \\
\operatorname{head}_{i}(x) & =\operatorname{softmax}\left(\frac{x W_{Q i}\left(X W_{K i}\right)^{\mathrm{T}}}{\sqrt{d_{\text {head }}}}\right) X W_{V i},
\end{aligned}
$

ここで F は活性化関数2)，concat は大力されたべクトルを結合する関数，softmax はソフトマックス関数, $W_{1} \in \mathbb{R}^{d \times d_{f f n}}, W_{2} \in \mathbb{R}^{d_{f f n} \times d}, W_{Q i} \in \mathbb{R}^{d \times d_{\text {head }}}$, $W_{K i} \in \mathbb{R}^{d \times d_{\text {head }}}, W_{V i} \in \mathbb{R}^{d \times d_{\text {head }}}, W_{O} \in \mathbb{R}^{d \times d}$ は重み行列である. また， $X$ は Self-attention 計算のために，系列分の層への入力ベクトルをまとめた行列である.

## 2.2 初期化手法

初期化手法はニューラルモデルの学習の挙動に多大な影響を及ぼすことが知られており，Transformer の初期化についても様々な手法が提案されている $[17,18,19,20]$. 本研究では, LLM の事前学習において広く使われている初期化手法 $[8,6,9]$ を用いる前提で議論を行う.この初期化手法では $\sigma=\sqrt{2 / 5 d}$ とし $[18], W_{2}$ と $W_{O}$ 以外の重み行列を $\mathcal{N}(0, \sigma)$ の正規分布を用いて初期化する. $W_{2}$ と $W_{O}$ については標準偏差がより小さな値となるよう， $\mathcal{N}(0, \sigma / \sqrt{2 N})$ を用いて初期化を行う。なお， $N$ は層の数である.

## 3 事前学習の不安定性の分析

我々は勾配爆発が学習を不安定にし，損失スパイクを引き起こしていると考えている. 本節では, $\mathrm{LN}$ 層が勾配爆発を引き起こすこと，また，勾配爆発を防ぐ手法を紹介する。

## 3.1 LN 層での勾配爆発

Xiong らは下記式のように，LNの勾配のノルムは入力ベクトルのノルムに依存することを証明し

1）簡略化するために本稿ではバイアス項を抜いている.

2）議論を単純にするために，本稿では恒等写像を仮定する.
た $[14]^{3)}:$

$
\left.\|\frac{\delta \operatorname{LN}(x)}{\delta x}\right.\|_{2}=\sigma\left(\frac{\sqrt{d}}{\|x\|_{2}}\right),
$

ここで， $x$ が正規分布に従うとすると,$x$ の平均は 0 であり $x$ のノルムは標準偏差に比例するため，標準偏差を考えればノルムが得られる。具体的には， $\operatorname{std}(x)$ を $x$ の標準偏差とすると, $\|x\|_{2}=\sqrt{d} \operatorname{std}(x)$ となる. よって上記の式は下記のように変形できる：

$
\left.\|\frac{\delta \mathrm{LN}(x)}{\delta x}\right.\|_{2}=\sigma\left(\frac{1}{\operatorname{std}(x)}\right)
$

すなわち, $\operatorname{std}(x) \ll 1$ の場合に $\mathrm{LN}$ 層は勾配のノルムを著しく増加させる．勾配がどこで増加しているかを推定するため，LNに入力されるべクトルの標準偏差を調べていく.

2.1 節で記したように, 各中間層の出力が $\mathrm{LN}$ 層への入力となる. すなわち，式 (1)，(2)のように， residual branch と $\operatorname{FFN}(x)$ もしくは Attention $(x)$ の和が次の層の $\mathrm{LN}$ 層への入力となる. ここで， $\operatorname{var}(x)$ をべクトル $x$ の分散とすると，正規分布に従うべクトル同士の和については下記が成り立つ：

$
\operatorname{var}(x+y)=\operatorname{var}(x)+\operatorname{var}(y) .
$

従って, $\operatorname{FFN}(x)$ および Attention $(x)$ で得られるべクトルもまた正規分布に従うと仮定した場合 ${ }^{4}$,各層の出力の分散については, residual branch, $\operatorname{FFN}(x), \operatorname{Attention}(x)$ の分散を考えれば良い. 第 1 層目の residual branch は埋め込み表現であり，埋め込み表現は $\mathcal{N}(0, \sigma)$ で初期化されている. 2.2 節より, $\sigma=\sqrt{2 / 5 d} \ll 1$ であるため, $\operatorname{std}(\mathrm{FFN}(x))$ と $\operatorname{std}(\operatorname{Attention}(x))$ が 1 よりも十分小さければ，式 (8) より各層の出力の標準偏差も 1 より十分小さく, LN 層によって勾配が爆発する。

学習初期の $\operatorname{var}(\mathrm{FFN}(x))$ は以下の式で得られる：

$
\operatorname{var}(\mathrm{FFN}(x))=d_{f f n} d \operatorname{var}(x) \operatorname{var}\left(W_{1}\right) \operatorname{var}\left(W_{2}\right)
$

なお，式 (1) にあるように FFN の計算前には LN が適用されるので $\operatorname{var}(x)=1$ である. 多くの研究で採用されている值と同様に $d_{f f n}=4 d$ とし, $\operatorname{var}\left(W_{1}\right)$ と $\operatorname{var}\left(W_{2}\right)$ については 2.2 節での値を代入すると，

$
\begin{aligned}
\operatorname{var}(\operatorname{FFN}(x)) & =4 d d \frac{2}{5 d} \frac{2}{5 d} \frac{1}{2 N} \\
& =\frac{8}{25 N}
\end{aligned}
$

3）Xiong らの証明は RMSNorm [21] にも適用できるため，本稿での議論は LNを RMSNorm に代えても成立する.

4）各パラメータを正規分布で初期化しているため，F゙恒等写像と置き， Attention $(x)$ のヘッドが 1 の場合にはこの仮定は成立する．実際のモデルにおいてどの程度現実的な仮定となっているかの検証は今後の課題である.

図 2 事前学習初期の Transformer の各層の勾配のノルム.一般にLLM においては $N>10$ であるため, $\operatorname{std}(\mathrm{FFN}(x))=\sqrt{8 / 25 N} \ll 1$ となる.

次に $\operatorname{var}(\operatorname{Attention}(x))$ について考える。まず $\operatorname{var}\left(\operatorname{head}_{i}(x)\right)$ について考える. $Z=\frac{x W_{Q i}\left(X W_{K i}\right)^{\mathrm{T}}}{\sqrt{d_{\text {head }}}}$ とすると $\operatorname{var}\left(\operatorname{softmax}\left(\frac{x W_{Q i}\left(X W_{K i}\right)^{\mathrm{T}}}{\sqrt{d_{\text {head }}}}\right) X\right)=\operatorname{var}(\operatorname{softmax}(Z) X)$ と書ける． $X$ は LN を適用されたべクトルをまとめたものであるため, $\operatorname{var}(\operatorname{softmax}(Z) X)$ は $\operatorname{softmax}(Z)$ が 1-hot ベクトルである場合に最大値である 1 となる5). 従って，式 (5) の分散は以下で得られる：

$
\begin{aligned}
\operatorname{var}\left(\operatorname{head}_{i}(x)\right) & =\operatorname{var}(\operatorname{softmax}(Z) X) d \operatorname{var}\left(W_{V i}\right) \\
& =\operatorname{var}(\operatorname{softmax}(Z) X) d \frac{2}{5 d} \\
& <\frac{2}{5}
\end{aligned}
$

従って, $\operatorname{var}(\operatorname{Attention}(x))$ は以下の式で得られる：

$
\begin{aligned}
\operatorname{var}(\operatorname{Attention}(x)) & =\operatorname{var}(\operatorname{head}(x)) d \operatorname{var}\left(W_{O}\right), \\
& =\operatorname{var}(\operatorname{head}(x)) d \frac{2}{5 d} \frac{1}{2 N}, \\
& =\operatorname{var}(\operatorname{head}(x)) \frac{1}{5 N} \\
& <\frac{2}{25 N}
\end{aligned}
$

$\operatorname{std}(\mathrm{FFN}(x))$ の際と同様に $N>10$ であるため, $\operatorname{std}(\operatorname{Attention}(x))<\sqrt{2 / 25 N} \ll 1$ となる. これらから，特に浅い層において ${ }^{6)}$, 式 (1)，(2) の出力べクトルの標準偏差は 1 より十分小さく, LN 層によって勾配が爆発する。

## 3.2 LN 層の勾配の安定化手法

$\mathrm{LN}$ 層での勾配爆発を防ぐためには $\mathrm{LN}$ 層へ入力するべクトルの標準偏差が 1 に近い値である必要が

5）なお, $\operatorname{softmax}(Z)$ が一様分布となる場合に最小值となり,入力系列が長いほど値が小さくなる. Li らは学習初期の入力系列は短い方が安定すると報告しているが [22], 本稿の議論から, $\operatorname{var}(\operatorname{softmax}(Z) X)$ の值が小さくなりすぎないように抑制する効果があるためと推測される.
6) $\mathrm{FFN}(x)$ と $\operatorname{Attention}(x)$ は層が深くなるごとに足されていくため，出力の標準偏差は層が深くなるほど大きくなっていく. 例えば，最終層では $N$ 回足されているため, 式 (10) と (12) の分母の $N$ を除去できる.
ある.これを達成するためには次のいずれかの方法がある：1. residual branch の標準偏差を大きくする， 2. $\operatorname{FFN}(x)$ と Attention(x) の標準偏差を大きくする. しかしながら，2 の場合には層が深くなるに連れての各層の出力のノルムの増加が著しく，これによって LN 層に起因しない勾配爆発が生じてしまう．悪影響を及ぼすことなくLN 層の勾配爆発を防ぐために，1を選択する。

第 1 層目の residual branch は埋め込み表現であるため，埋め込み表現の標準偏差を 1 に近づければ residual branch の標準偏差も 1 に近づく.これを達成する手法はいくつか考えられるが7)，本稿では2 種，Scaled Embed と Embed LN を紹介する． Scaled Embed は適切な值で埋め込み表現をスケーリングするものである．例えば，埋め込み表現に $\sqrt{d}$ を掛け合わせると ${ }^{8)}$ ，埋め込み表現の標準偏差は $\sqrt{2 / 5}$ となる. Embed LN はLNを埋め込み表現に適用するもので，既存研究でも損失スパイクを防いだと報告されている [6]. 学習初期での各層の勾配のノルムを図 2 に示す.この図から，何も対処をしていない場合（Vanilla）は層が浅くなるに連れて勾配が大きく増大しているが， Scaled Embed と Embed LN は勾配のノルムが一定であることが分かる. すなわち， Scaled Embed と Embed LN は LN 層による勾配爆発を防いでいることが分かる。

## 4 実験

本稿での理論分析について，実際に LLM の事前学習を通して検証する。すなわち，勾配爆発を抑える要件を満たす手法は損失スパイクが発生しないこと，要件を満たしていない手法は既存研究では損失スパイクの抑制に効果があるとされていてもスパイクを引き起こしてしまうことを示す.

## 4.1 データセット

事前学習の訓練・開発データとしては Common $\mathrm{Crawl}^{9)}$ から抽出した英語コーパスである，C4 [25] を用いる. 語彙は Byte Pair Encoding（BPE）[26] で構築されている，GPT-2 [13] の語彙を用いる. 事前学習したモデルの評価には WikiText [27] と

7）例えば埋め込み表現と出力層のパラメータ共有 $[23,24]$ を行わない場合，埋め込み表現の初期化における $\mathcal{N}(0, \sigma)$ について $\sigma=1$ とすれば良い. 本稿では初期化は 2.2 節で紹介したものを用いる前提で議論を進める.

8）この手続きはオリジナルの Transformer では行われていたが [5]，近年の実装では削除されている。

9) https://commoncrawl.org/

(a) 開発データでの損失関数の値.

(b) 勾配のノルム.

図 3 各手法の更新回数に対する開発データでの損失関数の值と勾配のノルム.

LAMBADA [28] データセットを用いてパープレキシティを計算する.

## 4.2 モデル設定

2 節で記したように，Pre-LN Transformer を用い， LLM の事前学習で広く使われている手法 $[8,6]$ で初期化する. モデルパラメータは 17 億とし, 層数 $N=24$, 各層の次元数 $d=2304$ とする. 入力の系列長は 2048 とし, 勾配ノルムを 1.0 でクリッピングする. 学習率 (lr) は $5.0 \times 10^{-4}$ とする ${ }^{10)}$ ．詳細な八イパーパラメータは付録 B に記す.

## 4.3 比較手法

3 節で言及した Vanilla，Embed LN，Scaled Embed に加え, Embed Detachを比較手法とする. Embed Detach では下記式のように，埋め込み表現の勾配を部分的に計算グラフから切り離すことで勾配を縮小する [29]：

Embed $=\gamma$ Embed $+(1-\gamma)$ Detach $($ Embed $)$, ここで $\gamma$ はイパーパラメータであり，Detach は入力を計算グラフから切り離す関数である. Zeng らは Embed Detach を LLM の事前学習に適用することで学習が安定したと報告しているが，この手法は本稿で説明した勾配爆発を防げてはいないため，本質的には損失スパイクを抑制できていないはずである. 本実験では Zeng らと同様に $\gamma=0.1$ とする.

## 4.4 結果

図 3 に更新回数に対する各手法の開発データでの損失関数の值と勾配のノルムを示す.この図から分かるように, Vanilla と Embed Detach では損失スパイ


表 1 各手法のパープレキシティ.

クおよび勾配のノルムのスパイクが起きているが, Embed LN と Scaled Embed では起こっていない. すなわち，3 節で示した勾配爆発を防ぐ要件を満たす手法では損失スパイクが起きておらず，勾配爆発の抑制が LLM の事前学習の安定させることが分かる.

表 1 に各手法の WikiText と LAMBADA でのパー プレキシティを示す.この表から, Embed LN と Scaled Embed は Vanilla や Embed Detach よりも高い性能を達成していることが分かる. この結果は，損失スパイクを防ぐことは学習したモデルの性能向上にも貢献すると示唆している。また, Embed LN と Scaled Embed は同等の性能を達成していることから，損失スパイクを抑制していれば，手法間には性能に有意な差はないと考えられる.

## 5 おわりに

本研究では LLM の事前学習においてしばしば発生する，損失スパイクの原因を分析した. 勾配に着目し，LLM の事前学習では LN 層が勾配を増大させていること， residual branch の標準偏差を 1 に近づけることで勾配の増加を抑制できることを理論的に示した. 実験を通して, residual branch の標準偏差を 1 に近づける手法，すなわち， Embed LN と Scaled Embed は実際に損失スパイクや勾配のノルムのスパイクを抑制し, LLM の事前学習を安定化させることを示した.

## 謝辞

本稿の執筆にあたり理研 AIP の Benjamin Heinzerling 氏に有用な助言をいただきました。感謝いたします.

## 参考文献

[1] T. Brown, B. Mann, N. Ryder, et al. Language models are fewshot learners. In Advances in Neural Information Processing Systems 33 (NeurIPS), pages 1877-1901, 2020.

[2] A. Chowdhery, S. Narang, J. Devlin, et al. PaLM: Scaling language modeling with pathways, 2022.

[3] H. Touvron, L. Martin, K. Stone, et al. Llama 2: Open foundation and fine-tuned chat models, 2023

[4] J. Kaplan, S. McCandlish, T. Henighan, et al. Scaling laws for neural language models, 2020.

[5] A. Vaswani, N. Shazeer, N. Parmar, et al. Attention is all you need. In Advances in Neural Information Processing Systems 30 (NIPS), pages 5998-6008. 2017.

[6] T. Le Scao, T. Wang, D. Hesslow, et al. What language model to train if you have one million GPU hours? In Findings of the Association for Computational Linguistics: EMNLP 2022, pages 765-782, 2022.

[7] A. Zeng, X. Liu, Z. Du, et al. GLM-130b: An open bilingual pre-trained model. In The Eleventh International Conference on Learning Representations (ICLR), 2023.

[8] M. Shoeybi, M. Patwary, R. Puri, et al. Megatron-lm: Training multi-billion parameter language models using model parallelism, 2020.

[9] S. Biderman, H. Schoelkopf, Q. Anthony, et al. Pythia: A suite for analyzing large language models across training and scaling. In Proceedings of the 40th International Conference on Machine Learning (ICML), pages 2397-2430, 2023.

[10] S. Takase, S. Kiyono, S. Kobayashi, and J. Suzuki. Spike no more: Stabilizing the pre-training of large language models, 2023.

[11] J. L. Ba, J. R. Kiros, and G. E. Hinton. Layer normalization, 2016.

[12] A. Radford, K. Narasimhan, T. Salimans, and I. Sutskever. Improving language understanding by generative pre-training. 2018.

[13] A. Radford, J. Wu, R. Child, et al. Language models are unsupervised multitask learners. 2019.

[14] R. Xiong, Y. Yang, D. He, et al. On layer normalization in the transformer architecture. In Proceedings of the 37th International Conference on Machine Learning (ICML), pages 10524-10533, 2020.

[15] L. Liu, X. Liu, J. Gao, W. Chen, and J. Han. Understanding the difficulty of training transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 5747-5763, 2020.

[16] S. Takase, S. Kiyono, S. Kobayashi, and J. Suzuki. B2T connection: Serving stability and performance in deep transformers. In Findings of the Association for Computational Linguistics: ACL 2023, pages 3078-3095, 2023.

[17] H. Zhang, Y. N. Dauphin, and T. Ma. Fixup initialization: Residual learning without normalization. In Proceedings of the 7th International Conference on Learning Representations (ICLR), 2019.

[18] T. Q. Nguyen and J. Salazar. Transformers without tears: Improving the normalization of self-attention. In Proceedings of the 16th International Conference on Spoken Language Translation (IWSLT), 2019.

[19] B. Zhang, I. Titov, and R. Sennrich. Improving deep transformer with depth-scaled initialization and merged attention. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLPIJCNLP), pages 898-909, 2019.

[20] H. Wang, S. Ma, L. Dong, et al. Deepnet: Scaling transformers to 1,000 layers, 2022.

[21] B. Zhang and R. Sennrich. Root mean square layer normalization. In Advances in Neural Information Processing Systems 32 (NeurIPS), 2019.

[22] C. Li, M. Zhang, and Y. He. The stability-efficiency dilemma: Investigating sequence length warmup for training gpt models. In Advances in Neural Information Processing Systems 35 (NeurIPS), pages 26736-26750, 2022.

[23] O. Press and L. Wolf. Using the output embedding to improve language models. In Proceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics (EACL), pages 157-163, 2017.

[24] H. Inan, K. Khosravi, and R. Socher. Tying word vectors and word classifiers: A loss framework for language modeling. In Proceedings of the 5th International Conference on Learning Representations (ICLR), 2017.

[25] C. Raffel, N. Shazeer, A. Roberts, et al. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1-67, 2020.

[26] R. Sennrich, B. Haddow, and A. Birch. Neural machine translation of rare words with subword units. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (ACL), pages 1715-1725, 2016.

[27] S. Merity, C. Xiong, J. Bradbury, and R. Socher. Pointer Sentinel Mixture Models. In Proceedings of the 5th International Conference on Learning Representations (ICLR), 2017.

[28] D. Paperno, G. Kruszewski, A. Lazaridou, et al. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (ACL), pages 1525-1534, 2016.

[29] M. Ding, Z. Yang, W. Hong, et al. Cogview: Mastering textto-image generation via transformers. In Advances in Neural Information Processing Systems 34 (NeurIPS), pages 1982219835, 2021.

(a) 17 億パラメータ, $\mathrm{lr}=1.0 \times 10^{-3}$.

(b) 17 億パラメータ, $1 r=5.0 \times 10^{-4}$.

(c) 17 億パラメータ, $\mathrm{lr}=1.0 \times 10^{-4}$.

(d) 130 億パラメータ, $1 r=3.0 \times 10^{-4}$.

(e) 130 億パラメータ, $\operatorname{lr}=1.0 \times 10^{-4}$.

図 417 億パラメータおよび 130 億パラメータで学習率を変えた際の各手法の開発データでの損失関数の値.

表 217 億パラメータの各手法の各学習率におけるパープレキシティ。

表 3130 億パラメータの各手法の各学習率におけるパープレキシティ。

## A パラメータ数・学習率を変えた実験

本節では 4 節での 17 億パラメータ，学習率 $5.0 \times 10^{-4}$ に加え，17 億パラメー タで学習率を $1.0 \times 10^{-3}, 1.0 \times 10^{-4}$ とした場合，および， 130 億パラメータでの実験を行う. 130 億パラメータは実際の LLM の事前学習の状況により近くなるよう，バッチサイズを大きな值とし，Adamの $\beta_{2}=0.95$ とした. $\beta_{2}=0.95$ は事前学習を安定にすると報告されている値である。また，学習率は Llama 2 [3] で用いられている $3.0 \times 10^{-4}$ と，より低い值である $1.0 \times 10^{-4}$ を用いた。 4 節と同様に，訓練・開発データは C4を，語彙には GPT-2のものを用いた．計算資源の都合上，ここでは Vanilla と Scaled Embed のみを対象に実験を行う.

図 4 に各手法の各設定での更新回数に対する開発データ上での損失関数の値を示す. 17 億パラメータについてである，図 (a), (b), (c) から明らかなように， Vanilla では学習率を大きくするほど損失スパイクが頻繁に発生するようになっている. 特に， 17 億パラメータのモデルを学習率 $1.0 \times 10^{-3}$ で学習した際には，

表 4 実験に用いたハイパーパラ

Vanilla では損失関数の値が発散してしまい，学習が失敗してしまっている。一方で，Scaled Embed については学習が安定しており，損失関数の値は更新回数に応じて一貫して下がっている. 130 億パラメータについては， $\beta_{2}=0.95$ としているからか，Vanillaでも学習率 $1.0 \times 10^{-4}$ で損失スパイクは見られない. しかしながら，高い学習率，すなわち， $\mathrm{lr}=3.0 \times 10^{-4}$ では損失関数の値が発散してしまっており，学習が失敗してしまっている. この結果から， $\beta_{2}=0.95$ のような学習を安定化させると報告されている状況においても，より安定した学習のために， residual branch の標準偏差を 1 に近づけることは必要と言える. 表 2, 表 3 に 17 億パラメータと 130 億パラメータの WikiText と LAMBADA データセットでのパープレキシティを示す. これらの表から， 17 億パラメータ，130 億パラメータ共に，学習率を変えた場合でも Scaled Embed は Vanilla と同等か，より高い性能を達成できていることが分かる.

## B ハイパーパラメータ

表 4 に実験に用いたハイパーパラメータを示す.