NESTONE、判断の難しさに応じて推論深度を変えるモデル

導入

大規模言語モデルを、素早く直接的に判断するSystem 1型のモデルへ転用する動きがある。一般的な方法は、言語モデル用の出力ヘッドを判断用ヘッドに置き換え、トランスフォーマーのバックボーンはそのまま利用するものだ。この方法は既存能力を活用できる一方、簡単な入力であっても全層を通過させるため、推論コストが固定されてしまう。Hugging Face Daily Papersで紹介された研究は、判断に必要な深さはどの程度なのかを検討した。

主なポイント

• 判断情報は層ごとに均等ではない。 初期層は基本的な特徴を構築し、中間層には判断に必要な情報が多く集まるという。報告された実験では、最後の6層を取り除いても性能低下は確認されなかった。
• 小型化だけでは専門化につながらない。 独立して学習した小型モデルは、判断に特化した能力を獲得するというより、完全モデルの挙動を大きく模倣する傾向を示した。単純にバックボーンを短くするだけでは、効率化の効果が限られる可能性がある。
• NESTONEは異なる深さで計算を共有する。 完全モデルと複数の前置き部分モデルを同時に学習し、後段の層と判断ヘッドを共有する。自己蒸留によって、浅い経路にも完全モデルの判断能力を学習させる。
• ゲートが早期終了を選ぶ。 軽量なゲートが各出口の隠れ状態を読み、追加計算が必要かを判断する。容易な入力は浅い地点で終了し、複雑な入力だけが深い層へ進む。

意義と影響

この研究の重要な点は、効率化を「より小さなモデルを作ること」だけで捉えていないことだ。深さの異なるモデルを別々に保持する代わりに、1つのモデルが複数の推論経路を持ち、パラメータと知識を共有する。これにより、訓練や保存における重複を減らしながら、難しい入力には高い計算能力を提供できる可能性がある。

低遅延が求められる判断システムでは、簡単なケースを早く処理し、曖昧なケースだけに追加計算を割り当てる構成が有効になり得る。また、モデルの最終層だけでなく、中間表現が実際に利用可能な判断能力を持つかを評価する必要性も示している。

ただし、提供された資料には、7つのベンチマークの名称、詳細な精度、実機での計算コストは記載されていない。そのため、NESTONEの実運用上の優位性は、異なるタスクやハードウェアで追加検証する必要がある。

出典：Hugging Face Daily Papers