Dr.LLM : Routage dynamique des couches dans les modèles de langage de grande taille

Dr.LLM: Dynamic Layer Routing in LLMs

October 14, 2025

papers.authors: Ahmed Heakl, Martin Gubri, Salman Khan, Sangdoo Yun, Seong Joon Oh

cs.AI

papers.abstract

Les modèles de langage de grande taille (LLMs) traitent chaque jeton à travers toutes les couches d'une pile de transformateurs, entraînant un gaspillage de calcul pour les requêtes simples et une flexibilité insuffisante pour les requêtes plus complexes nécessitant un raisonnement approfondi. Les méthodes à profondeur adaptative peuvent améliorer l'efficacité, mais les approches précédentes reposent sur une recherche coûteuse au moment de l'inférence, des modifications architecturales ou un réentraînement à grande échelle, et en pratique, elles dégradent souvent la précision malgré les gains d'efficacité. Nous présentons Dr.LLM, Dynamic routing of Layers for LLMs, un cadre rétrofitable qui équipe les modèles préentraînés de routeurs légers par couche décidant de sauter, d'exécuter ou de répéter un bloc. Les routeurs sont entraînés avec une supervision explicite : en utilisant la recherche arborescente de Monte Carlo (MCTS), nous dérivons des configurations de couches de haute qualité qui préservent ou améliorent la précision sous un budget de calcul. Notre conception, incluant le regroupement par fenêtres pour un routage stable, la perte focale avec équilibrage des classes et les routeurs MLP à goulot d'étranglement, garantit la robustesse face au déséquilibre des classes et aux longues séquences. Sur ARC (logique) et DART (mathématiques), Dr.LLM améliore la précision jusqu'à +3,4% tout en économisant en moyenne 5 couches par exemple. Les routeurs généralisent à des tâches hors domaine (MMLU, GSM8k, AIME, TruthfulQA, SQuADv2, GPQA, PIQA, AGIEval) avec seulement une baisse de précision de 0,85% tout en conservant l'efficacité, et surpassent les méthodes de routage précédentes jusqu'à +7,7%. Globalement, Dr.LLM montre que des routeurs explicitement supervisés rétrofitent des LLMs figés pour une inférence consciente du budget et axée sur la précision sans altérer les poids de base.

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.