Dr.LLM: 大規模言語モデルにおける動的レイヤールーティング

Dr.LLM: Dynamic Layer Routing in LLMs

October 14, 2025

著者: Ahmed Heakl, Martin Gubri, Salman Khan, Sangdoo Yun, Seong Joon Oh

cs.AI

要旨

大規模言語モデル（LLMs）は、すべてのトークンをトランスフォーマーの全層を通じて処理するため、単純なクエリでは計算が無駄になり、より深い推論を必要とする難しいクエリには柔軟性が不足します。適応的な深さの手法は効率を向上させることができますが、従来のアプローチは高コストな推論時の検索、アーキテクチャの変更、または大規模な再学習に依存しており、実際には効率の向上にもかかわらず精度が低下することがよくあります。我々は、Dr.LLM（Dynamic routing of Layers for LLMs）を紹介します。これは、事前学習済みモデルに軽量な層ごとのルーターを装備し、ブロックをスキップ、実行、または繰り返すかを決定する後付け可能なフレームワークです。ルーターは明示的な監督で訓練されます：モンテカルロ木探索（MCTS）を使用して、計算予算の下で精度を維持または向上させる高品質な層構成を導出します。我々の設計、安定したルーティングのためのウィンドウプーリング、クラスバランスを考慮した焦点損失、およびボトルネックMLPルーターは、クラスの不均衡や長いシーケンスの下での堅牢性を確保します。ARC（論理）とDART（数学）では、Dr.LLMは精度を最大+3.4%p向上させながら、平均して1例あたり5層を節約します。ルーターは、ドメイン外のタスク（MMLU、GSM8k、AIME、TruthfulQA、SQuADv2、GPQA、PIQA、AGIEval）にも一般化し、効率を維持しながらわずか0.85%の精度低下で、従来のルーティング手法を最大+7.7%p上回ります。全体として、Dr.LLMは、明示的に監督されたルーターが、ベースの重みを変更することなく、予算を意識した精度主導の推論のために凍結されたLLMsを後付けできることを示しています。

English

Large Language Models (LLMs) process every token through all layers of a transformer stack, causing wasted computation on simple queries and insufficient flexibility for harder ones that need deeper reasoning. Adaptive-depth methods can improve efficiency, but prior approaches rely on costly inference-time search, architectural changes, or large-scale retraining, and in practice often degrade accuracy despite efficiency gains. We introduce Dr.LLM, Dynamic routing of Layers for LLMs, a retrofittable framework that equips pretrained models with lightweight per-layer routers deciding to skip, execute, or repeat a block. Routers are trained with explicit supervision: using Monte Carlo Tree Search (MCTS), we derive high-quality layer configurations that preserve or improve accuracy under a compute budget. Our design, windowed pooling for stable routing, focal loss with class balancing, and bottleneck MLP routers, ensures robustness under class imbalance and long sequences. On ARC (logic) and DART (math), Dr.LLM improves accuracy by up to +3.4%p while saving 5 layers per example on average. Routers generalize to out-of-domain tasks (MMLU, GSM8k, AIME, TruthfulQA, SQuADv2, GPQA, PIQA, AGIEval) with only 0.85% accuracy drop while retaining efficiency, and outperform prior routing methods by up to +7.7%p. Overall, Dr.LLM shows that explicitly supervised routers retrofit frozen LLMs for budget-aware, accuracy-driven inference without altering base weights.