Dr.LLM: Dynamisches Schicht-Routing in LLMs

papers.abstract

Große Sprachmodelle (LLMs) verarbeiten jedes Token durch alle Schichten eines Transformer-Stacks, was zu verschwendeter Rechenleistung bei einfachen Anfragen und unzureichender Flexibilität für komplexere Aufgaben, die tiefere Schlussfolgerungen erfordern, führt. Adaptive Tiefenmethoden können die Effizienz verbessern, aber bisherige Ansätze beruhen auf kostspieliger Inferenzzeit-Suche, architektonischen Änderungen oder groß angelegtem Neu-Training und führen in der Praxis oft zu einer Verschlechterung der Genauigkeit trotz Effizienzgewinnen. Wir stellen Dr.LLM, Dynamic Routing of Layers for LLMs, vor, ein nachrüstbares Framework, das vortrainierte Modelle mit leichtgewichtigen pro-Schicht-Routern ausstattet, die entscheiden, ob ein Block übersprungen, ausgeführt oder wiederholt wird. Die Router werden mit expliziter Supervision trainiert: Mithilfe von Monte-Carlo-Baumsuche (MCTS) leiten wir hochwertige Schichtkonfigurationen ab, die die Genauigkeit unter einem Rechenbudget bewahren oder verbessern. Unser Design, das Fenster-Pooling für stabiles Routing, Focal Loss mit Klassenausgleich und Bottleneck-MLP-Router umfasst, gewährleistet Robustheit bei Klassenungleichgewicht und langen Sequenzen. Bei ARC (Logik) und DART (Mathematik) verbessert Dr.LLM die Genauigkeit um bis zu +3,4 %P, während durchschnittlich 5 Schichten pro Beispiel eingespart werden. Die Router generalisieren auf domänenübergreifende Aufgaben (MMLU, GSM8k, AIME, TruthfulQA, SQuADv2, GPQA, PIQA, AGIEval) mit nur 0,85 % Genauigkeitsverlust bei Beibehaltung der Effizienz und übertreffen bisherige Routing-Methoden um bis zu +7,7 %P. Insgesamt zeigt Dr.LLM, dass explizit überwachte Router eingefrorene LLMs für budgetbewusste, genauigkeitsgetriebene Inferenz nachrüsten können, ohne die Basisgewichte zu verändern.

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.

Dr.LLM: Dynamisches Schicht-Routing in LLMs

Dr.LLM: Dynamic Layer Routing in LLMs

papers.abstract

Support