Nexus：同一の事前学習損失で、共通ミニマにより下流タスクの汎化を改善

Abstract

事前学習は、大規模言語モデル（LLM）の基盤であり、その能力を駆動する主要エンジンとして、膨大な計算予算とデータの大部分を占めています。事前学習の間、LLMは類を見ないほど巨大で多様なデータソースから、基礎的な知識を獲得します。そこには、一般言語、数学、コード、複雑な推論など、幅広い領域が含まれます。本研究では、事前学習の収束した状態に関する興味深い幾何学的問いを検討します。すなわち、モデルは、すべてのデータソース（例：\cref{fig:cwa_illustration:close}）にわたって共通の最小値へ収束するのか、それとも、単に合計損失の最小値（例：\cref{fig:cwa_illustration:distant}）へ収束するだけなのか、という問いです。我々は、タスク固有の最小値同士の幾何学的な「近さ」が、下流の汎化に本質的に結びついていると仮説を立てます。標準的な最適化手法（例：AdamW）は、タスク固有の最小値が互いに離れているような点に収束することが多いことを示します。これに対処するために、Nexus最適化器を提案します。これは、最適化中に勾配の類似度を最大化することで、これらの最小値の近さを促進します。パラメータ数130Mから3Bまでのモデル、さまざまなデータ混合、および多様なハイパーパラメータのスケジュールにわたる実験により、Nexusは\textit{（事前学習損失が同じであるにもかかわらず）下流性能を大幅に向上させる}ことを示します（\cref{fig:demo:benchmark}参照）。特に3Bモデルでは、Nexusは分布外損失を0.012だけ低減し、複雑な推論タスク（例：GSM8k）で最大15.0\%の精度向上をもたらします。この発見は、モデル評価の唯一の代理指標として事前学習損失への依存に疑問を投げかけ、下流の汎化を解き放つうえで暗黙のバイアスが重要であることを示しています。