\documentclass[japanese,other]{jnlp_3.8}
\usepackage{jnlpbbl_2.2}
\usepackage[dvipdfmx]{graphicx}
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{url}
\usepackage{array}
\allowdisplaybreaks

\usepackage{hangcaption_jnlp}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% author's definition

%%%%%\usepackage{amsthm}
%%%%%\usepackage{algorithm}
%%%%%\usepackage{amsfonts}
%%%%%\usepackage{mleftright}
%%%%%\usepackage{multirow}
\usepackage{bm}
\usepackage{booktabs}
%%%%%\usepackage{wrapfig}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

\Volume{32}
\Number{1}
\Month{March}
\Year{2025}
%%%%\Issuetitle{「　　　　」特集号}


\setcounter{page}{366}



%%%%   タイトル部分                                      %%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\jsubject{学会記事}
\jtitle{Initialization of Large Language Models via Reparameterization to Mitigate Loss Spikes}
\jauthor{西田　光甫\affiref{Author_1}}


\etitle{Initialization of Large Language Models via Reparameterization to Mitigate Loss Spikes}
\eauthor{Kosuke Nishida\affiref{Author_1}} 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\headauthor{西田}
\headtitle{Initialization of Large Language Models via Reparameterization to Mitigate Loss Spikes}



%%%% 所属を指定して下さい．
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\affilabel{Author_1}{日本電信電話株式会社 NTT人間情報研究所}{NTT Human Informatics Laboratories, NTT Corporation}




%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{document}
\maketitle



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S1
\section{はじめに}

本稿はEMNLP2024に採択された論文 ``Initialization of Large Language Models via Reparameterization to Mitigate Loss Spikes'' \cite{nishida-etal-2024-initialization}の解説記事である．
今回はあえて通常の論文の流れではなく，時系列の順に記述することで，論文執筆の経緯に焦点を当てようと思う．

筆者が所属するNTT人間情報研究所は大規模言語モデルtsuzumiの研究開発に取り組んでいる．筆者はそのプロジェクトで事前学習を遂行しており，本研究はプロジェクトを進める中で生まれたものである．大規模言語モデルの事前学習では，損失関数値が突然発散してしまう損失スパイク現象が知られている．損失スパイク現象は様々な問題を引き起こす：(1) ハイパーパラメータへの頑健さを失い，丁寧なチューニングを要する．(2) 小さな学習率を設定する必要が生じ，最終的な性能が抑制される．(3) 時に完全に発散し，それまでの学習を無駄にしてしまう．(4) 何より，開発者の精神衛生に悪い．このような背景から，損失スパイク現象の抑制に取り組んだのが本研究の始まりであった．



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S2
\section{Transformerの初期化の謎}
Transformer \cite{transformer}の初期化には勾配爆発・消失を避けるための工夫が施されている．Fine-Tuningでは気にすることがない初期化手法だが，事前学習の場合は初期値の決定が重要な要素になる．
まずは，初期化の既存研究とその背景について調査する必要があった．

Transformerの重みパラメータを，$\bm{W}_e$が埋め込み層，$\bm{W}_k, \bm{W}_q, \bm{W}_v, \bm{W}_o$がそれぞれSelf-{\linebreak}Attention層のkey-, query-, value-, output-projection層，$\bm{W}_u, \bm{W}_d$がMLP層のup-, down-projection層，$\bm{W}_p$が予測層とする．また，$d$をTransformerの隠れ次元数，$N$をTransformer層の数とする．

深層学習モデルの初期化でよく使われる手法はHe Initialization \cite{He_2015_ICCV}である．
この手法は，層の前後で勾配のノルムが不変になるように設計されており，モデル全体での勾配爆発・消失を避ける効果がある．しかし，大規模言語モデルの初期化では，より小さい標準偏差で初期化するSmall Initialization \cite{nguyen-salazar-2019-transformers}が経験的に安定すると言われる．これらの初期化手法を表\ref{tab:init_list}にまとめる．
ここで，ゲートとは文献 \cite{takase2023spike}が提案した埋め込み層$\bm{W}_e$の出力を$\bar{\bm{W}_e} =\alpha \bm{W}_e$の式によって標準偏差1に変換する手法であり，Layer Normalization前後での勾配のノルムの変化を安定化する役割がある．



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%T1
\begin{table}[b]
\input{36table01.tex}
\caption{各初期化手法が設定する標準偏差（Weight列）と，ゲートを通過した$\bar{\bm{W}}$の標準偏差（Gate列）．}
\label{tab:init_list}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%



本研究では，$\bm{W}_o, \bm{W}_d$の二つの層で標準偏差に$\sqrt{\frac{1}{2N}}$などの極端に小さな値がかけられている点に着目する．この理由を，ResNetにおける議論 \cite{taki2017deep}に基づいて説明する．Self-Attention変換を$f$を書く．このとき，Self-Attention層では
%%%%%
\begin{align}
\bm{y} =f(\mathrm{LN}(\bm{x})) + \bm{x}
\label{eq:residual_forward}
\end{align}
%%%%%
の計算が行われる．なお，LNはLayer Noramalizationである．式\ref{eq:residual_forward}における勾配は
%%%%%
\begin{align}
\frac{\partial \mathcal{L}}{\partial \bm{x}} = \frac{\partial \mathcal{L}}{\partial \bm{y}} \frac{\partial \bm{y}}{\partial \bm{x}} =\frac{\partial \mathcal{L}}{\partial \bm{y}} \left(\frac{\partial f(\mathrm{LN}(\bm{x}))}{\partial \bm{x}} +\bm{I}\right)
\label{eq:residual_backward}
\end{align}
%%%%%
である．よって，$\left|\frac{\partial f(\mathrm{LN}(\bm{x}))}{\partial \bm{x}}\right|^2$のオーダーを$s$としたとき，
勾配のノルムの2乗はSelf-Attention層で$s+1$倍されて下の層に伝播する．
式\ref{eq:residual_backward}内の$+\bm{I}$はResidual結合に起因する項であり，同じ事象はMLP層でも起こる．そのためモデル全体では，$(s+1)^{2N}$倍の勾配爆発が起こる．なお，簡便のためMLP層の勾配のノルムの2乗も$s$で表した．
勾配が$N$に関して指数的に増大することは，$N$が大きくなる大規模言語モデルでは許容できない．そこで，$\bm{W}_o, \bm{W}_d$層の初期値の標準偏差を通常の$1/\sqrt{2N}$倍とすることで，$E[s] =\mathcal{O}\left(\frac{1}{2N}\right)$とする．すると，モデル全体での勾配爆発$(s+1)^{2N}$が$N\rightarrow \infty$で$e$に収束し，$N$に関する指数的な増大を回避できる．





%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S3
\section{本研究の取り組み}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S3.1
\subsection{提案手法の着想}
前節の初期化手法を踏まえて，本研究では，全てのパラメータを同じ学習率で更新するにもかかわらず$\bm{W}_d, \bm{W}_o$だけ小さい標準偏差を設定することには害もあるのでは，と考えた．この考えを発展させ，本研究では，ゲートパラメータを全てのパラメータに導入することで，全てのパラメータを均一に初期化することを提案した．

パラメータ$\bm{W}_i$の初期化を確率分布$\mathcal{N} (0, \sigma_i)$で行う場合を考える．
このとき，提案手法では，全パラメータで共通の標準偏差を$\sigma$として，
%%%%%
\[
\bm{W}_i \sim \mathcal{N} (0, \sigma), \quad
\bar{\bm{W}}_i =\frac{\sigma_i}{\sigma} \bm{W}_i
\]
%%%%%
とする．つまり，$\bm{W}_i$を$\mathcal{N} (0, \sigma_i)$からサンプリングする代わりに，全パラメータで共通のスケールを持つ$\bm{W}_i$をパラメータとしてモデルに登録する．そして，パラメータをモデル内で使う際は常にスカラ$\frac{\sigma_i}{\sigma}$をかけた$\bar{\bm{W}}_i$として使用する．よって，$\bar{\bm{W}}_i$の標準偏差は$\sigma_i$となる．ここで，$\frac{\sigma_i}{\sigma}$は実際には学習可能なゲートパラメータ$\alpha_i$とし，その初期値を$\frac{\sigma_i}{\sigma}$とする．

この操作により，モデルの挙動を変えることなく均一な標準偏差で初期化することが可能になる．
全てのパラメータにゲートパラメータを導入する利点は，どんなに極端な初期化手法であっても，その初期化手法を採用しながら，実際に更新されるパラメータを均一な標準偏差に設定できる点である．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S3.2
\subsection{提案手法の高度化}
\label{ssec:propose_init}
しかし，均一的初期化だけでは，既存手法のSmall Initializationに損失関数値の減衰速度に関して劣っていた．当時は，ハイパーパラメータの$\sigma$を$\frac{1}{\sqrt{d}}$で固定していたためである．散歩の折に$\sigma$は任意の値に設定できることに気がつき，Small Initializationに倣い$\sigma$を小さな値に設定することで損失関数の減衰が速くなること，Small Initializationの規定値よりもさらに小さい$\sigma^2 =\textrm{4e-5}$に設定することで安定かつ高速な学習を実現できることを発見した．

深層学習モデルの慣習では$\sigma_i$を$d$を引数とする関数値として決定している．結果として，$d$の大きなモデルほど小さな学習率を設定する必要がある．均一的初期化では，$\sigma_i$を$d$を引数とする関数値としながら$\sigma$を任意の値に定めることができる．例えば全ての$d$に対して固定の$\sigma$と学習率を採用することも可能であり，提案手法は深層学習の慣習を打破する可能性を秘めていると考えられる．



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S3.3
\subsection{提案手法の解釈}
提案手法は経験的に安定した学習を実現したが，「なぜ提案手法が安定するのか」，あるいは「初期値の標準偏差の不均一性は真に害をもたらすのか」を調査することは重要である．
そこで，パラメータ$\bm{W}$とその1時刻における更新量$\Delta \bm{W}$を用いて，
パラメータの更新比率$\|\Delta \bm{W}\| / \|\bm{W}\|$を定義した．更新比率はパラメータが相対的にどの程度激しく更新されるかを示す値である．
つまり，$\bm{W}_i$の更新比率が極端に大きい場合，$\bm{W}_i$で更新が安定するように学習率を小さい値に設定して他のパラメータの更新を犠牲にするか，他のパラメータの更新を優先して$\bm{W}_i$の急激な更新を許容するか，の2択を迫られる，よって，更新比率はモデル全体である程度の範囲に収まっていることが好ましいと考えられる．なお，パラメータの初期値の平均は0であるため，標準偏差は$\|\bm{W}\|$の初期値の期待値に比例する．



%%%%%%%%%%%%%%%%%%%%%%%%%%%%F2
\begin{figure}[b]
\begin{center}
\includegraphics{32-1ia36f2.eps}
\end{center}
\hangcaption{13Bモデルの学習初期の損失関数値（上）と最終層の$\bm{W}_d, \bm{W}_u$における更新比率（下） \protect\cite{nishida-etal-2024-initialization}．}
\label{fig:13B_init_loss}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%



学習初期の損失関数値と，最終層の$\bm{W}_d, \bm{W}_u$における更新比率を図\ref{fig:13B_init_loss}に示す．
実際に，学習の最初期は標準偏差の小さい$\bm{W}_d$の更新比率が極端に大きいこと，小規模なスパイクを起こしながら$\bm{W}_d$の更新比率が小さくなっていくこと，最後のスパイクの前後で$\bm{W}_d$の更新比率の値が大きく減少し，$\bm{W}_d, \bm{W}_u$の更新比率がある程度の範囲に収まること，を確認できた．本発見は，「初期値の標準偏差の不均一性は真に害をもたらすのか」の疑問に対して，「損失スパイクの前後でパラメータの更新比率が大きく変化することが確認できる」「パラメータの更新比率の不均一性は損失スパイクの一因であり（または，少なくとも関連はあり），標準偏差の不均一性は害になりうる」の回答を提示する．

上記の内容は原論文中ではIntroductionで研究の動機として述べられているが，発見の順序としては最後であった，
本来は論文の順序通りに最初に更新比率に着目して研究を始めるのが理想ではあったが，得られた結果を説明する現象を後から発見することも研究の面白さの一つである，と改めて感じるいい機会であった．



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S4
\section{評価}
パラメータ数130億 (13B) 個のTransformerモデルの事前学習により評価した．モデル構造はLLaMA \cite{touvron2023llama}に準拠したが，MLP層の活性化関数にはgelu関数を用いた．
RefinedWeb \cite{penedo2023refinedweb}データセットから30Bトークンをサンプリングして事前学習を行ない，損失関数値の減衰を評価した．
ベースラインモデルとして，Small Initializationによって初期化したモデルを採用した．
学習率とバッチサイズとして，5e-4と4Mトークンを採用するStable設定と，1e-3と1Mトークンを採用するRapid設定を採用した．Stable設定は学習の安定を重視した慣習的な設定であり，Rapid設定は高速な減衰を意図した設定である．
図\ref{fig:13B_init_loss}はRapid設定の2モデルの比較である．提案手法は学習初期の損失関数値の減衰を安定化し，かつ更新比率を一定の小さな範囲に抑制している．図\ref{fig:13B_loss}はベースラインのStable設定を加え，30Bトークン全てでの学習の様子を示す．Stable設定であっても損失スパイクが発生すること，提案手法は30Bトークンまで損失スパイクを起こさないことを確認できる．



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%F1
\begin{figure}[b]
\begin{center}
\includegraphics{32-1ia36f1.eps}
\end{center}
\hangcaption{13Bモデルの学習中の損失関数値．ベースラインモデルでは学習の初期に損失スパイクが発生 \protect\cite{nishida-etal-2024-initialization}．}
\label{fig:13B_loss}
\vspace{-25pt}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%





%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S5
\section{おわりに}
大規模言語モデルに著しい注目が集まっているが，その学習原理や動作原理は謎のままである．大規模言語モデルの研究開発を通じて，これらの謎の探究を続けていく所存である．



%%% Bibliography %%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begingroup
\addtolength{\baselineskip}{-0.75pt}
\bibliographystyle{jnlpbbl_1.7}
\bibliography{36refs}
\endgroup


%%% Biography %%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{biography}
\bioauthor{西田　光甫}{2015年東京大学工学部計数工学科卒業．2017年東京大学情報理工学系研究科数理情報学専攻修士課程修了．2024年東京大学情報理工学系研究科電子情報学専攻博士課程修了．博士（情報理工学）．現在，日本電信電話株式会社NTT人間情報研究所准特別研究員．}

\end{biography}


\end{document}





