条件付き深さルーティングにおける補助損失の再検討：実証研究

arXiv cs.LG / 2026/4/21

📰 ニュースDeveloper Stack & InfrastructureIdeas & Deep AnalysisModels & Research

共有:

要点

この論文は、軽量なFFN（cheap）と通常のFFN（full）をトークンごとに切り替える条件付き深さルーティングの学習安定化において、補助損失がどのように効くかを検討しています。
157.5Mパラメータのデコーダのみモデルで（コントローラのみ学習、fullパス予算50%、3シード）、ゲート設計の2案（G1：MLPのユーティリティスコア、G3：JEPAガイドの行動条件付き予測）を比較し、標準レシピ（util/rank補助損失）下ではG3が3/3回の試行で初期〜中盤の最適化を改善すると報告されています。
アブレーションでは、util/rank補助監督を除去すると両ゲートで最良/平均のLMと閾値到達速度が改善し、G3のG1に対する初期〜中盤の優位性も消えることが示されています。
その原因として、util/rankのオラクルラベルが「以降の層はすべてfullを実行する」と仮定するオフポリシー性を持つ一方、実際のゲート実行では一部トークンしかfullパスに流れないため、現行レシピではutil/rankが負の効果を持つ点が挙げられています。
util/rankを外すことで学習FLOPsの代理指標も（約1.53×から1.07×へ）低下し、検証された範囲では計算効率の実利も示唆されています。

Abstract

条件付きの深さ実行ルーティングでは、トークンの一部を軽量で低コストなFFNに通し、残りは制御された各層で通常の完全FFNを実行します。中心となる困難はゲート学習です。ゲートの判断は、言語モデル（LM）の損失に影響するまで多くの層をまたいで伝播する必要があるため、得られる勾配は弱くノイズが多くなります。訓練を安定化するために補助損失を積み重ねることは一般的ですが、補助損失同士の相互作用、特に予測型の補助と明示的なスコア監督との間の相互作用は、制御された条件下で体系的に比較されていません。本研究では、157.5Mパラメータのデコーダのみモデルに対して、コントローラのみの学習、50%のフルパス予算、fineweb-eduサブセットで3シードの実験を行い、2つのゲート設計を評価します。MLPゲート（G1）は現在の隠れ状態をユーティリティスコアに写像します。一方でJEPAに導かれたゲート（G3）は、低次元の潜在空間で、トークンごとに「完全を実行する」か「安価を実行する」かを実行した結果を予測する、アクション条件付きの予測器を追加し、固定されたターゲットヘッドに照合します。util/rank（utilの回帰＋ペアごとの順位付け監督）のオラクル風ユーティリティ回帰を用いた標準レシピでは、G3は3/3のシードすべてでG1よりも初期〜中期の最適化が改善します（平均LMが低い、閾値到達が速い、勾配ノルムが約10.3倍低い）。また、20kステップのエンドポイントLMは、0.005というヒューリスティック参照値の範囲内にあります。重要な発見（アブレーションA3）：util/rankをゲートに対して同時に除去すると、両方のゲートで3/3のシードすべてにおいて最良／平均LMおよび閾値到達速度が改善し、G3がG1に対して示していた初期〜中期の優位性は消失します。これを、次のように追跡します。util/rankに用いるオラクルラベルはオフポリシーであり、以降のすべての層が完全を実行すると仮定していますが、条件付き実行では完全を通るのは一部の分率に過ぎません。そのため、現在のレシピではutil/rankが負の効果（net-negative）になります。さらにutil/rankを除去すると、訓練FLOPsの代理指標が約1.53xから約1.07xへと削減されます（V100-32GBで2.87hから1.75h、約39%）。結論は、検討したレジームに限定されます。