Dr.LLM: Instradamento Dinamico degli Strati nei Modelli Linguistici di Grande Dimensione

Dr.LLM: Dynamic Layer Routing in LLMs

October 14, 2025

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

cs.AI

Abstract

I modelli linguistici di grandi dimensioni (LLM) elaborano ogni token attraverso tutti gli strati di una pila di trasformatori, causando un calcolo sprecato per query semplici e una flessibilità insufficiente per quelle più complesse che richiedono un ragionamento più profondo. I metodi a profondità adattativa possono migliorare l'efficienza, ma gli approcci precedenti si basano su ricerche costose durante l'inferenza, modifiche architetturali o riaddestramento su larga scala, e nella pratica spesso riducono l'accuratezza nonostante i guadagni in efficienza. Introduciamo Dr.LLM, Dynamic routing of Layers for LLMs, un framework retrocompatibile che equipaggia modelli preaddestrati con router leggeri per ogni strato che decidono di saltare, eseguire o ripetere un blocco. I router vengono addestrati con supervisione esplicita: utilizzando il Monte Carlo Tree Search (MCTS), deriviamo configurazioni di strati di alta qualità che preservano o migliorano l'accuratezza sotto un budget di calcolo. Il nostro design, che include il pooling a finestra per un routing stabile, la focal loss con bilanciamento delle classi e router MLP a collo di bottiglia, garantisce robustezza in caso di squilibrio delle classi e sequenze lunghe. Su ARC (logica) e DART (matematica), Dr.LLM migliora l'accuratezza fino a +3,4% risparmiando in media 5 strati per esempio. I router si generalizzano a task fuori dominio (MMLU, GSM8k, AIME, TruthfulQA, SQuADv2, GPQA, PIQA, AGIEval) con solo un calo di accuratezza dello 0,85% mantenendo l'efficienza, e superano i metodi di routing precedenti fino a +7,7%. In generale, Dr.LLM dimostra che router supervisionati esplicitamente possono adattare LLM congelati per inferenze guidate dall'accuratezza e consapevoli del budget senza alterare i pesi di 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.