Research map / Large language models / Reasoning and chain of thought
Dr.LLM: Dynamic Layer Routing in LLMs
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Dr.LLM retrofits frozen LLMs with lightweight per-layer routers that choose skip, execute, or repeat for each block, trained on 4k MCTS-derived layer paths, improving accuracy on ARC and DART while saving layers.
Problem
LLMs push every token through all layers regardless of difficulty, wasting compute on easy inputs and under-allocating depth on hard ones; prior adaptive-depth methods trade accuracy for speed, need retraining, or require costly inference-time search.
Approach
Offline length-aware MCTS finds accuracy-preserving or improving skip/repeat layer paths per question; lightweight per-layer MLP routers over windowed mean-pooled hidden states are trained with focal loss on these labels, and applied greedily at inference on frozen base weights.
Key results
- On DART, LLaMA-3B-Base improves accuracy from 11.8% to 15.8% (+4.0 points) while saving about 4.1 layers per example on average.
- Out-of-domain benchmarks show an average drop of 0.85 points across four instruct models, with routers preserving layer savings.
- Dr.LLM achieves 67.4 average versus 59.7 for FlexiDepth on GSM8k, MMLU, HellaSwag, and HumanEval for LLaMA-8B (in-domain for baselines).
Surprising findings
- Routers trained on first-layer embeddings perform worse than the vanilla model without routing (DART 30.0 vs 35.2), indicating routing relies on internal states rather than question type.
- Router macro F1 is only about 61% yet accuracy still improves, because most errors conservatively predict execute.
Limitations
- Out-of-domain drop of 0.85 points on average, with some benchmarks (e.g., MMLU with Qwen3B, -2.5) losing accuracy.
- Routing labels are highly imbalanced; repeat and skip classes are rare (2.3% and 7.6%), and macro F1 is only 61%.
- MCTS supervision requires offline search costing about 961k forward passes, and per-layer synchronization complicates batching.
Experimental setup
Datasets: ARC-Easy ARC-Challenge DART-Math
Benchmarks: ARC DART MMLU AIME24 TruthfulQA GSM8k GPQA Diamond AGIEval SQuADv2 PIQA HumanEval HellaSwag
Metrics: accuracy average executed layers macro F1 wall-clock time throughput
Baselines: FlexiDepth MindSkip ShortGPT LayerSkip CoLa Vanilla model without routing
Builds on: CoLa LayerSkip Mixture of Depths Universal Transformer focal loss
Closest papers by meaning
- Universal Model Routing for Efficient LLM Inference (ICLR 2026 · similarity 0.91)
- BEST-Route: Adaptive LLM Routing with Test-Time Optimal Compute (ICML 2025 · similarity 0.91)
- RouteLLM: Learning to Route LLMs from Preference Data (ICLR 2025 · similarity 0.90)
- Confidence-Guided Stepwise Model Routing for Cost-Efficient Reasoning (AAAI 2026 · similarity 0.90)
- R2-Router: A New Paradigm for LLM Routing with Reasoning (ICML 2026 · similarity 0.90)
- $\textit{Read-ME}$: Refactorizing LLMs as Router-Decoupled Mixture of Experts with System Co-Design (NeurIPS 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.
