モデル拡張によるパラメータ効率的なLLMの事前学習

大規模言語モデル（LLM）は，人工知能分野の基盤的な技術となり，実用システムにおいても用いられるようになった．性能の高いLLMを獲得するためには，スケーリング則[1]に基づき，膨大なパラメータを持つTransformerモデル[2]を大規模コーパスで事前学習する流れが主流である[3]．結果として，事前学習を現実的な時間で終えるには膨大な数のGPUを搭載した計算資源が要求される[4]．この巨大な計算資源の要求は，LLMの事前学習研究への重大な参入障壁となっている．計算資源量の課題を受け，本研究では，LLM事前学習の計算資源量削減について議論する．計算資源量削減には様々な観点があり得るが，本稿では，標準的な事前学習と同等の性能を維持しながら，事前学習に必要となるGPUの最大メモリ要求量を一定量に抑制可能な事前学習法を提案する．概要としては，層追加によるモデル拡張法と，微調整学習でよく用いられるパラメータ効率のよい学習法を組合せた新たな事前学習法を構築する．提案法の具体的な説明として，図1に提案法の手順に関する概要図を示す．提案法では，逐次モデル拡張手順での各ステージの最大メモリ要求量をモデル設定値を変数とした整数計画問題として定式化し，その解を得て実段階1段階2🔥❄Layer5Layer5Layer3🔥🔥⼿順4:訓練⼿順1:⼩規模モデルの訓練⼿順2:層拡張⼿順3:PETの適⽤Layer4Layer4🔥❄Layer3Layer3Layer2🔥🔥Layer2Layer2🔥❄Layer1Layer1Layer1🔥な事前学習を実施する（手順1）．その後，学習済みモデルに新たに層を追加しモデルを拡張する（手順2）．学習済みの層はパラメータを固定し，代替の学習用にPET（Parameter-EﬃcientTraining）を適用（手順3），拡張後のモデルを再度学習する（手順4）．手順4では，層拡張で追加したパラメータとPETによって導入された小規模なパラメータのみが学習対象となる．際の各ステージでのモデル設定に用いることで，事前学習の実行前に実行中の最大メモリ要求量が最小になるようにモデル拡張の設定を制御する．これにより，最大メモリ要求量を一定量に抑えた事前学習が実現できる．以降，本提案手法をSTEP（STagedparameterEﬃcientPre-training）と記述する．本稿では，一定のFLOPsの条件下で，既存手法との性能比較を通じてSTEPの有効性を検証する．その結果，従来の事前学習法と比べ，事前学習モデルの標準的な評価タスクにおいて同等の性能を維持しながら，最大メモリ要求量を53.9%削減可能であることが示された．さらに，標準的な微調整学習を実施した場合でも，STEPで学習したモデルは，標準的な事前学習によるモデルと下流タスクにおいて同等の性能を達成することを示す．