技術要約：選択的レイヤー復元による多様性の回復

問題提起

インストラクションチューニングや人間からのフィードバックによる強化学習（RLHF）といったポストトレーニング手法は、大規模言語モデル（LLM）の指示追従能力と有用性を大幅に向上させる。しかし、これらの手法は**モード崩壊（mode collapse）**を引き起こすことが多く、これはモデルが狭い範囲の応答に対して不釣り合いに高い確率質量を割り当て、多様ではあるが妥当な代替案を抑制してしまう現象である。これにより、創造的な執筆、オープンエンドなQA、推論などの設定において、出力が反復的になり、独特な応答を必要とするアプリケーションにおけるモデルの有用性が制限される。

既存の緩和戦略には3つのカテゴリがあるが、それぞれに限界がある：

1. デコーディングベースの手法（例：温度スケーリング、min-pサンプリング）：確率性を高めるが、ポストトレーニング中のパラメータ変化によって失われた多様性を回復することはできない。
2. プロンプトベースの手法（例：複数の回答を要求する）：推論コストが増大し、プロンプトの言い回しに敏感である。
3. トレーニングベースの手法（例：目的関数の再設計）：高価な再学習とデータ収集を必要とする。

著者らは、多様性を効果的に回復し、適応時の計算効率が高く、推論時の追加コストが発生しない手法を求めている。

手法

提案手法である**選択的レイヤー復元（Selective Layer Restoration: SLR）**は、2つの経験的な観察に基づく、トレーニングフリーの重み空間介入である：

1. LLMのレイヤーは、明確に異なる機能的役割を持つ（例：事実的な関連性は、しばしば中間層の深さに局在している）。
2. ポストトレーニングは、主にどの事前学習済み行動を表現するかを制御しており、新しい知識を完全に獲得するというよりは、既存の行動を抑制（破滅的忘却）することが多い。

著者らは、モード崩壊は特定のレイヤーに局在しており、ポストトレーニング済みモデル（$M_{post}$）の特定のレイヤー範囲を事前学習済みモデル（$M_{pre}$）の重みに復元することで、品質を維持しながら多様性を回復できるという仮説を立てている。

SLRフレームワーク

$N$個のレイヤーを持つモデルに対し、SLRは、$M_{post}$内の連続したレイヤー区間 $[i, j]$ を対応する$M_{pre}$のレイヤーで置き換えることで、ハイブリッドモデル $M_{i:j}$ を構築する。目標は、最小品質閾値（$Q \ge q_{min}$）を満たしつつ、多様性（$D$）を最大化する区間 $(i, j)$ を見つけることである。

CRCプロキシ・タスク

複雑なダウンストリーム・タスク上で復元区間を直接最適化することは、計算コストが高く、「品質」の主観性ゆえに評価も困難である。これに対処するため、著者らは**制約付きランダム文字（Constrained Random Character: CRC）**をプロキシ・タスクとして導入した。

• メカニズム： CRCは、厳密に定義された有効な集合内（例：「[0, 5]の範囲の整数を生成せよ」）でランダムな文字（数字または文字）を生成することを含む。
• 指標：
  • 品質（$Q_{CRC}$）： 有効な出力に割り当てられた確率質量（指示への遵守度を測定）。
  • 多様性（$D_{CRC}$）： 有効な出力に対する分布のエントロピー（多様性を測定）。
• 選択プロセス： 著者らは、すべての可能な連続レイヤー区間に対してCRCを用いたグリッドサーチを実行する。$Q_{CRC} \ge 0.9 \times Q_{CRC}(M_{post})$ を満たしつつ、$D_{CRC}$を最大化する区間を選択する。選択された区間は、さらなるチューニングなしにダウンストリーム・タスクに適用される。

主な貢献

1. 選択的レイヤー復元（SLR）： ポストトレーニング済みLLMの特定の連続したレイヤーの重みを、事前学習済みのパラメータに復元する、トレーニングフリーの手法。元のアーキテクチャとパラメータ数を保持し、追加の推論コストが発生しない。デコーディングやプロンプトベースの多様性手法と補完的であることが示されている。
2. 制約付きランダム文字（CRC）： 明示的な有効集合と自然な多様性目的を持つプロキシ・タスク。CRCは、効率的かつ低コストな復元区間の探索を可能にし、レイヤー局在化仮説の経験的な証拠を提供する。
3. 経験的検証： 3つのモデルファミリー（Llama, Qwen, Gemma）および3つのダウンストリーム・タスク（創造的な執筆、オープンエンドなQA、多段階推論）にわたる包括的な評価。

実験結果

著者らは、SLRをポストトレーニング済みのベースラインおよび「モデル・スープ（Model Soup）」ベースライン（事前学習済みモデルとポストトレーニング済みモデルの重み平均）と比較評価した。

• 創造的な執筆： SLRは、最小限の品質低下（LLMジャッジ・スコアで平均 -2.1%）で、一貫して多様性を向上させた（埋め込みベースの非類似度で平均 +42.5% の増加）。対照的に、Model Soupは多様性の利得が小さく（+27.3%）、品質の低下が大きかった（-6.0%）。
• オープンエンドなQA： SLRは回答の多様性を大幅に向上させ、エントロピーを +112.7% 、カバレッジ（ユニークな正解の割合）を +169.5% 増加させた一方で、高い精度を維持した（平均 -1.8% の低下）。
• 推論（GSM8K）： グリーディ・デコーディング下でポストトレーニング済みモデルが失敗する問題において、SLRはすべてのサンプリング予算（$k \in \{1, \dots, 64\}$）においてPass@kを改善した。これは、確率質量が多様だが正しい推論パスへとシフトしたことを示している。
• 結合可能性： SLRは、高い温度のデコーディング（$T=1.5$）や言語化サンプリング（プロンプトベースの手法）と組み合わせても効果を維持しており、これはSLRがサンプリングプロセスや入力を操作する他の手法とは異なり、基礎となるモデルの分布自体を修正していることを示唆している。
• アブレーション研究： CRCプロキシによって導かれたレイヤーを復元することは、任意の区間（例：最初または最後のレイヤー）を復元することよりも優れた結果をもたらした。任意のレイヤーを復元すると、多様性が向上しないか品質が低下することが多く、特定の区間選択が極めて重要であることを裏付けている。

意義と主張

本論文は、再学習の計算オーバーヘッドや複雑なプロンプティングによる推論コストなしに、モード崩壊を緩和するためのシンプルかつ効果的な道筋をSLRが提供すると主張している。多様性をターゲットとした重み空間の介入を通じて回復できることを示すことで、本研究はポストトレーニング済みLLMのモジュール的視点を支持している：

1. 機能や振る舞い（モード崩壊を含む）は、一様に分布しているのではなく、特定のレイヤーに局在している。
2. ポストトレーニングは、モデルの核となる知識ベースを根本的に変えるのではなく、主に事前学習された振る舞いの表現を制御している。

著者らは、微細な重み空間の介入（SLRなど）が、品質と多様性のトレードオフを制御するための有望な手段を提供すると結論づけている。これにより、ポストトレーニングの利点（指示への遵守など）を保持しつつ、望ましい事前学習の特性（多様性など）を回復することが可能になる。本論文はスケーラビリティについては控えめな姿勢を見せており、現在の結果は〜7〜9Bパラメータのモデルに限定されており、これらの知見がより大規模なスケールや、より粒度の細かい介入（例：サブレイヤーの復元）においても成立するかどうかを判断するには今後の研究が必要であると述べている。