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

GPTをはじめ大規模言語モデル（LargeLanguageModel，以降本稿ではLLMと記す）は様々なアプリケーションの根幹となっている[1,2,3]．LLMは大規模な訓練データで事前学習した大量のパラメータを持つニューラル言語モデルであり，性能はパラメータ数と訓練データ量に対数比例することが知られている[4]．事前学習は大量の計算資源を必要とするため学習が失敗するリスクを極力低減することが求められている．Transformer[5]は多くのLLMに採用されているニューラルモデルであるが，その挙動の包括的な理論的理解には至っていない．例えば，図1のVanillaのように，Transformerを用いたLLMの事前学習では損失関数の値が跳ね上がり（損失スパイク），また，しばしば発散してしまう現象が知られているが[2]，なぜ発生するのかは明らかにされていない．損失スパイクへの対処として様々な手法が提案されているが[2,6,7]，これらの理論的な正当化もなされていないため，既存研究の報告と異なるパラメータ数のモデルなど，状況が異なる場合でも適用できるかには疑問が残る．本研究ではLLMの事前学習における損失スパイクの理論的な分析を行う．特に初期化手法が学習に事前学習における訓練データでの損失関数の値の例．Vanillaは学習初期の段階で損失スパイクが何度か発生し，13000ステップで発散してしまっている．開発データでの損失関数の値は付録Aの図4(a)に示す．与える影響について議論する．TransformerをLLMに用いる場合，Self-attention層とFFN層の出力ベクトルの標準偏差が小さくなる初期化を行うと学習が安定することが経験的にも理論的にも示されている[8,6,9,10]．一方で，本研究では，このような初期化を行った場合にはLayerNormalization（LN）層[11]が勾配爆発を引き起こしてしまうことを示す．さらに，LN層での勾配爆発を防ぐための要件を示し，モデルの簡易な修正でこの要件を達成できることを示す．理論的な分析を検証するためLLMの事前学習の実験を行う．上記の，LN層での勾配爆発を防ぐ要件を満たしている場合には損失スパイクが発生せず，満たしていない場合には，仮に既存研究で損失スパイクへの対処として提案された手法であっても損失スパイクが発生してしまうことを示す．