Dr.LLM: Dynamische Laagroutering in LLM's

Dr.LLM: Dynamic Layer Routing in LLMs

October 14, 2025

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

cs.AI

Samenvatting

Grote Taalmodellen (LLMs) verwerken elk token door alle lagen van een transformatorstack, wat leidt tot verspilde berekeningen bij eenvoudige queries en onvoldoende flexibiliteit voor moeilijkere queries die dieper redeneren vereisen. Adaptieve-dieptemethoden kunnen de efficiëntie verbeteren, maar eerdere benaderingen vertrouwen op kostbare zoekacties tijdens inferentie, architectuurwijzigingen of grootschalige hertraining, en leiden in de praktijk vaak tot een verminderde nauwkeurigheid ondanks efficiëntiewinsten. Wij introduceren Dr.LLM, Dynamische routing van Lagen voor LLMs, een aanpasbaar framework dat voorgetrainde modellen uitrust met lichtgewicht per-laag routers die beslissen om een blok over te slaan, uit te voeren of te herhalen. Routers worden getraind met expliciete supervisie: met behulp van Monte Carlo Tree Search (MCTS) leiden we hoogwaardige laagconfiguraties af die de nauwkeurigheid behouden of verbeteren binnen een rekenbudget. Ons ontwerp, windowed pooling voor stabiele routing, focal loss met klassebalancering, en bottleneck MLP-routers, zorgt voor robuustheid bij klasseonbalans en lange sequenties. Op ARC (logica) en DART (wiskunde) verbetert Dr.LLM de nauwkeurigheid met tot +3,4%p terwijl het gemiddeld 5 lagen per voorbeeld bespaart. Routers generaliseren naar taken buiten het domein (MMLU, GSM8k, AIME, TruthfulQA, SQuADv2, GPQA, PIQA, AGIEval) met slechts 0,85% nauwkeurigheidsverlies terwijl de efficiëntie behouden blijft, en overtreffen eerdere routingmethoden met tot +7,7%p. Over het algemeen toont Dr.LLM aan dat expliciet gesuperviseerde routers bevroren LLMs aanpassen voor budgetbewuste, nauwkeurigheidsgedreven inferentie zonder de basisgewichten te wijzigen.

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.