DLLG: Dynamic Logit-Level Gating of LLM Experts

技術要約：DLLG: LLMエキスパートの動的なロジットレベル・ゲーティング

問題提起
特定のドメイン（数学的推論、コード生成など）に最適化された特化型大規模言語モデル（LLM）の普及は、一つの課題を提示している。それは、いかにして個別に学習された「ブラックボックス」的なエキスパートを、膨大な再学習コストをかけることなく、それらの相補的な強みを活用できる統一されたシステムへと統合するかという点である。既存のアプローチは、適応性、堅牢性、および実用性の間で、永続的な緊張関係に直面している：

• ルーティングベースの手法は、ハードな初期エキスパート選択に依存している。これは早期のコミットメント（決定）を招く。つまり、単一のレスポンス内のあるサブタスクにおいて選択されたエキスパートが失敗した場合、そのエラーは回復不可能となる。さらに、多くのルーターはベンチマークごとの広範なキャリブレーションを必要とする。
• ヒューリスティックなアンサンブル（パープレキシティ重み付け、トークン投票など）は、推論時のプロキシ信号を使用する。これらの指標は経験的なものが多く、真のタスクの正解性と乖離しているため、サブオプティマル（最適ではない）な出力につながる。
• パラメータ・マージング（モデル・スーピング、タスク・アーキテクティックなど）は、重みを静的に融合させる。これはモジュール性を犠牲にし、しばしば相反する勾配方向を持つエキスパートを組み合わせる際に破壊的な干渉を引き起こし、性能低下を招く。
• **混合エキスパート（MoE）**アーキテクチャは、通常、ゲートとエキスパートのエンドツーエンドの共同学習を必要とするため、オフザシェルフ（既製品）の事前学習済みスペシャリストを活用することに限界がある。

手法：DLLG
著者らは、疎なレスポンスレベルの監督から学習される軽量なゲーティングモデルを用いて、動的なトークンレベルのロジット融合を行うフレームワークである**DLLG（Dynamic Logit-Level Gating）**を提案している。

1. アーキテクチャ:

   • 凍結されたエキスパート: $K$ 個の特化型LLMは凍結された状態に保たれる。各デコーディングステップ $t$ において、プロンプト $x$ と接頭辞 $y_{<t}$ が与えられると、各エキスパート $E_i$ は隠れ状態 $h^{(i)}_t$ と次トークンのロジット $\ell^{(i)}_t$ を生成する。
   • ゲーティングモデル: 軽量なモデル ($g_\phi$) が、全エキスパートの連結された隠れ状態を処理する。このモデルは、低ランク投影を用いてこれらの表現を共有埋め込み空間に投影し、自己回帰的なステム（例：Qwen2.5-0.5B-Instruct に LoRA アダプターを適用したもの）を通じて処理を行う。
   • 重み予測: 予測ヘッドは、シグモイド変換を介してトークンごとの融合重み $w_{t,i}$ を出力する。これらの重みはエキスパートのロジットに適用される：$\ell_t = \sum_{i=1}^K w_{t,i} \ell^{(i)}_t$。重みは合計が1になるよう制約されておらず、ソフトで柔軟な寄与を可能にしている。

2. 学習（レスポンスレベルの監督）:

   • 本フレームワークは、トークンレベルのアノテーションの欠如に対処するため、自動検証器（例：コードのユニットテスト、数学の回答チェック）から得られるレスポンスレベルの正解信号 ($s_i \in \{0, 1\}$) を使用する。
   • ティーチャー・フォーシング: 学習中、エキスパートは正解の接頭辞 $y_{<t}$ に条件付けられた隠れ状態を生成する。ゲーティングモデルは、各トークン位置に対して重み $w_{t,i}$ を予測する。
   • 目的関数: モデルは、予測された重みとブロードキャストされたレスポンスレベルの正解信号との間の平均二乗誤差（MSE）を最小化するように学習される：$\mathcal{L} = \frac{1}{T} \sum_{t=1}^T \sum_{i=1}^K (w_{t,i} - s_i)^2$。これにより、ゲートはトークンレベルの融合を、軌跡全体の成功と一致させるように学習される。

3. 推論:

   • DLLGは自己回帰的に動作する。各ステップにおいて、ゲーティングモデルは現在のエキスパートの隠れ状態に基づいて融合重みを予測し、ロジットが結合され、次トークンが選択される。これにより、回復可能なエキスパート利用が可能になる。例えば、あるエキスパートが早い段階でミスをしたとしても、ゲーティングモデルは後続のステップでより有能なエキスパートへと動的に重みをシフトさせることができる。

主な貢献

1. フレームワークの提案: トークンレベルのアノテーションやオンライン強化学習を必要とせずに、トークンレベルの生成と疎なレスポンスレベルの監督を橋渡しする、動的なロジットレベルのアンサンブル・フレームワークであるDLLGを導入した。
2. メカニズムの分析: 学習されたゲーティング・メカニズムが、ロールアウト中にエキスパートの寄与を動的に適応させ、進化するコンテキスト（例：単一のレスポンス内で数学的推論からコード生成への移行）に応じて重みを調整することを実証した。
3. 実証的検証: 多様な推論およびコードのベンチマーク（GSM8K, MATH, HumanEvalなど）における評価を通じて、ルーティング、ヒューリスティックなアンサンブル、およびパラメータ・マージングのベースラインを一貫して上回る性能を示した。

結果

• 性能: 0.5Bおよび1.5Bの両方のモデルスケールにおいて、DLLGはRouterDC、EmbedLLM、各種ロジット・アンサンブル手法（GaC, UniTe）、およびパラメータ・マージング技術（Linear, SLERP, Task Arithmetic）を含む強力なベースラインと比較して、一貫して最高の平均性能を達成している。
• 堅牢性: 個別のエキスパートが弱い低容量（0.5B）の設定において、DLLGは脆弱なルーティングやヒューリスティックな手法を凌駕しており、ソフトな融合によって初期のミスを軽減できる能力を示している。
• 汎用性: 手法は、アウトオブドメイン（例：HumanEval, MBPP, BBH）のベンチマークにおいても強い性能を示しており、学習された融合ルールが特定の学習分布を超えて汎用することを示唆している。
• アブレーション研究: 実験により、効果的なトークンレベルの融合には、補助的なLLMと低ランク投影を組み込んだ十分に表現力のあるゲーティング・アーキテクチャが必要であることが確認され、単純なMLPヘッドよりも優れた性能を示すことが証明された。

意義と主張
本論文は、DLLGが特化したエキスパートを統合するための堅牢かつスケーラブルなパラダイムを提供すると主張している。弱監督からロジットレベルの融合を学習することで、エキスパートの再学習コストやパラメータ干渉のリスクを負うことなく、MoEのような適応性を実現している。著者らは、DLLGがエキスパートのモジュール性を維持し、きめ細かく、回復可能なエキスパート利用を可能にすることで、ルーティングにおける早期決定の限界や、パラメータ・マージングの静的な性質を克服することを強調している。この研究は、弱監督から詳細なエキスパート利用を学習することが、ますます専門化する言語モデルを活用するための実用的な道筋であることを示唆している。