Weight Decay Improves Language Model Plasticity

技術要約：ウェイトディケイが言語モデルの可塑性を向上させる

問題提起
現在の大規模言語モデル（LLM）の学習における慣習は、通常、事前学習とポストトレーニング（ファインチューニング、アライメント）を切り離して考えている。事前学習におけるハイパーパラメータ最適化とスケーリング則は、主にベースモデルの検証損失（クロスエントロピー）に基づいて評価されており、これは「より低い事前学習損失が、より有能なダウンストリームモデルにつながる」という仮定に基づいている。しかし、このアプローチはモデルの可塑性（plasticity）、すなわち、再初期化なしにさらなる学習を通じて新しいデータやタスクに成功裏に適応できる能力を見落としている。本論文は、事前学習の損失を最小化するために事前学習のハイパーパラメータを最適化することが、その後のファインチューニングにおける適応力を不当に損なう可能性があるかどうかを調査している。具体的には、標準的な正則化ハイパーパラメータである**ウェイトディケイ（weight decay）**に焦点を当て、事前学習損失を最小化するための最適な値が、ダウンストリームの適応力を最大化するための最適な値と異なるかどうかを検証している。

手法
著者らは、2つのモデルファミリー（Llama-2およびOLMo-2）と複数のモデルサイズ（最大4Bパラメータ）にわたって系統的な実証研究を行っている。実験設計では、他の要因を制御しながら、事前学習におけるウェイトディケイ・ハイパーパラメータ（$\lambda$）を変化させている。

1. 学習レジーム: モデルは以下の2つの異なるレジームで学習される：
   • 計算最適（Compute-optimal）: 20 TPP（tokens-per-parameter）。
   • 過学習（Overtrained）: 140 TPP。
2. ファインチューニング評価: 事前学習済みモデルを、以下を含む多様なダウンストリームタスクでファインチューニングする：
   • 6種類のChain-of-Thought (CoT) 推論データセット（例：MetaMathQA, MMLUProCoT）。
   • 5種類の言語理解および常識推論ベンチマーク（例：HellaSwag, WinoGrande）。
   • 1種類の安全性アライメントタスク。
3. 指標: ダウンストリームの性能は、Greedy (Pass@1), Pass@16, Majority Vote (Maj@16), Reward Model scoring (RM@16), および Correct Ratioを含む複数の指標を用いて評価される。
4. メカニズム分析: ウェイトディケイがなぜ可塑性に影響を与えるのかを理解するために、著者らは以下を分析している：
   • 線形プロービング（Linear Probing）: 内部表現の線形分離可能性（感情およびトピック分類）を測定。
   • アテンションのランク（Attention Rank）: クエリ・キー（$W_{QK}$）行列およびバリュー投影（$W_{VP}$）行列の擬似ランクを分析。
   • 過学習（Overfitting）: 学習損失と検証損失の差（Train-Val Gap）を算出。

主な貢献

1. 可塑性の駆動因子としてのウェイトディケイ: 本論文は、ウェイトディケイがモデルの可塑性を形成する上で極めて重要な要因であることを示している。標準的なデフォルト設定とは対照的に、事前学習におけるより大きなウェイトディケイの値は、その後のファインチューニング後のダウンストリーム性能を一貫して向上させる。
2. 目的のデカップリング（分離）: 本研究は、事前学習の検証損失を最小化するハイパーパラメータの値が、必ずしもダウンストリームの性能を最大化する値ではないという証拠を提示している。いくつかのレジームにおいて、高い事前学習損失を持つモデル（大きなウェイトディケイによるもの）が、ファインチューニング後に、より低い事前学習損失を持つモデルよりも優れた性能を示す。
3. メカニズムに関する洞察: 著者らは、ウェイトディケイが可塑性を高める3つの具体的なメカニズムを特定している：
   • 線形分離可能な表現を促進し、内部特徴をダウンストリームタスクに対してよりアクセスしやすいものにする。
   • アテンション行列を正則化し、特にクエリ・キーの積（$W_{QK}$）のランクを減少させることで、より圧縮され堅牢な構造をもたらす。
   • 事前学習データに対する過学習を抑制し、効果的に特定の学習ノイズを「忘却」させることで、新しい情報を学習するためのモデルの容量を保持する。

主な結果

• 最適なウェイトディケイのシフト: 計算最適（20 TPP）の学習では、ダウンストリーム性能のための最適なウェイトディケイは、標準的なデフォルト値である0.1よりも大幅に大きい。例えば、このレジームにおけるLlama-2およびOLMo-2モデルでは、0.5〜1.0で最適な値が見出された。過学習（140 TPP）レジームにおいても、ダウンストリーム性能のための最適値（0.3）は、事前学習損失を最小化する値（0.1）よりも高いままである。
• 直感に反するトレードオフ: 本論文は、「劣った」ベースモデル（事前学習損失が高いモデル）が、ファインチューニング後に「より優れた」最終モデルになり得るというトレードオフを強調している。例えば、$\lambda=0.3$または$1.0$で事前学習されたOLMo-2-1B-140xモデルは、$\lambda=0.1$のベースラインよりも事前学習損失が高かったが、顕著に優れたダウンストリーム精度を達成した。
• 堅牢性: これらの知見は、異なるモデルファミリー、サイズ、および事前学習とファインチューニングの両方で他のハイパーパラメータ（学習率、バッチサイズ）が変化した場合でも維持される。
• メカニズムの相関:
  • 高いウェイトディケイは、より構造化された表現を示す、高い線形プロービング精度と相関している。
  • 高いウェイトディケイは、アテンション行列（特に$W_{QK}$）におけるより低い擬似ランクと相関しており、低ランクで堅牢な仮説への移行を示唆している。
  • 高いウェイトディケイは、より小さなTrain-Val Gapと相関しており、過学習の減少を示している。

意義と主張
本論文は、ハイパーパラメータ最適化の唯一の指標として事前学習の検証損失にのみ依存することを、業界は再考すべきであると主張している。著者らは、モデルの可塑性は、純粋に事前学習の損失を最適化しようとする際にしばしば犠牲になる、明確かつ重要な特性であると主張している。

本研究の意義は以下の点にある：

1. デフォルトの再評価: 標準的なハイパーパラメータの選択（例：$\lambda=0.1$）は、ダウンストリームの適応力が究極の目標である場合、LLM開発の全ライフサイクルにおいて最適ではない可能性があることを示唆している。
2. クロスエントロピーの限界: クロスエントロピー損失はダウンストリームの能力の不完全なプロキシ（代理指標）であることを浮き彫りにしている。なぜなら、損失が同程度であっても、モデルの可塑性は劇的に異なる可能性があるからである。
3. ハイパーパラメータの多面的な役割: 単一のハイパーパラメータであるウェイトディケイが、汎化性能だけでなく、将来の学習を促進する構造的特性（線形性、ランク）にも影響を与えるという、複雑で多面的な役割を果たしていることを強調している。

著者らは、結果がテストされたスケールにおいて堅牢であるものの、特定の最適値は、現在の実験範囲を超えた極めて大規模な、あるいは高度に過学習されたシナリオにおいては変化する可能性があることを認め、謙虚な姿勢を保っている。彼らは、安定性と可塑性のトレードオフ、およびこれらの知見の他の基盤モデルのモダリティへの適用可能性について、さらなる研究を求めている。