Dr.LLM: Dynamic Layer Routing in LLMs

AI summary

Overview

• Research area: Natural Language Processing — adaptive computation and inference efficiency for large language models.
• Technical level: Intermediate. Readers should be comfortable with transformer blocks, hidden states, and basic supervised learning concepts; the paper explains its MCTS and loss formulations in detail.
• Scope: The paper introduces Dr.LLM, a retrofittable framework that adds small per-layer routers to a frozen LLM so each transformer block is skipped, executed once, or repeated, with routers trained on offline search-derived supervision.

What This Paper Is About

Standard LLMs push every token through every transformer layer, which wastes computation on easy inputs and offers no extra depth for hard ones. Prior adaptive-depth methods either sacrifice accuracy for speed, demand architectural changes and large-scale retraining, or require expensive search at inference time. Dr.LLM's goal is to give a frozen pretrained model learnable, per-layer skip/execute/repeat decisions that increase accuracy while reducing the average number of layers executed, without altering the base weights.

Key Contributions

1. Supervised dynamic routing for frozen LLMs. Per-layer routers choose skip, execute, or repeat for their block, trained end-to-end on only 4,000 execution paths discovered offline, while the base model stays frozen.
2. Length-aware MCTS for path supervision. An offline Monte Carlo Tree Search finds layer edits under a compute budget and retains only accuracy-preserving or accuracy-improving paths, producing a compact supervision dataset with no base-model modification.
3. Lightweight router design. Windowed mean pooling stabilizes decisions on long contexts, focal loss with effective-number class weights handles the heavily imbalanced skip/execute/repeat labels, and a two-layer bottleneck MLP keeps overhead negligible.
4. Accuracy plus efficiency without retraining or inference search. Across six models, Dr.LLM improves accuracy on ARC and DART while saving layers, and it outperforms prior routing methods on shared benchmarks.

Main Findings

• In-domain gains are consistent, never negative. On ARC (logic) and DART (math) across six LLaMA-3.2 and Qwen-2.5 backbones, accuracy improved in every case, with mean gains of +2.25 percentage points and roughly 5.0 fewer layers per example. ARC improved modestly (+0.9 to +2.5%p) while DART improved more (+1.4 to +4.0%p), consistent with math needing deeper iterative refinement.
• Largest single result. LLaMA-3B-Base rose from 11.8% to 15.8% on DART (+4.0%p) while saving 4.12 layers per query; LLaMA-8B-Instruct gained +2.8%p on DART while saving 11.0 layers per query.
• Wall-clock savings. End-to-end generation for 1,000-token sequences dropped 15.3% (29.21s vs 34.49s baseline), with router overhead of only 0.27s per query, under 1% of total latency. Because routers run once per input sequence, they remain compatible with KV caching.
• Strong out-of-domain generalization. Across MMLU, GSM8k, AIME24, TruthfulQA, SQuADv2, GPQA Diamond, PIQA, and AGIEval, accuracy dropped only 0.85%p on average across four instruct models while efficiency was retained. On GSM8k all four models gained an average 0.40%p; on AIME24 accuracy was unchanged. Some cases improved, e.g., +2.5%p on GPQA Diamond with LLaMA-3B.
• Beats prior routing methods. Against LayerSkip, ShortGPT, MindSkip, and FlexiDepth on LLaMA3-8B, Dr.LLM reached 67.4% average accuracy versus 59.7% for the closest competitor, a 7.7%p gap, despite the comparison benchmarks being in-domain for the prior methods and out-of-domain for Dr.LLM. FlexiDepth also required training on 326k examples; Dr.LLM used 4k.
• Routing follows a structured phase pattern. Early layers are almost always executed once, middle layers are frequently skipped, and late layers are often repeated, especially on DART. The late-layer repetition effect is stronger in the 8B model, suggesting larger models exploit additional depth for complex reasoning.
• Routers learn from internal state, not question type. A control router trained on first-layer embeddings (for all layers) performed far worse, dropping 8.6%p on DART versus Dr.LLM and falling below the vanilla model by 5.2%p. This indicates routers map internal representations to routing decisions rather than memorizing shallow input patterns.
• Batching survives dynamic depth. Per-layer microbatch synchronization costs little because repeats occur in only 2.3% of layers and routing decisions agree 80–89% of the time across sequences. Throughput improved 8.65% on average across batch sizes, including 26.89% at batch size 1.
• Ablations favor compact, windowed routers. A bottleneck dimension of 128 gave the best accuracy (+3.4%p); two linear layers outperformed deeper routers; more pooling windows consistently improved both accuracy and layer savings; focal loss beat weighted cross-entropy by +1.1%p on ARC and +1.8%p on DART; zero-initialized biases were more stable than frequency-based initialization.
• Training is cheap. Routers add 11M trainable parameters for 3B models (0.14% of base weights) and 16.8M for 8B models (0.56%). Training completes within four hours on a single A100 40GB GPU using 20% of VRAM, at an average per-layer routing accuracy of 96.8%.
• Search cost is offline only. MCTS required roughly 961k forward passes to collect 4k supervision examples from 24,330 visited paths, but no search occurs at inference. About 30% of retained edited paths beat the default path in accuracy.

Methodology in Plain English

The researchers start with an ordinary pretrained LLM and leave its weights completely untouched. Alongside each transformer block they attach a very small neural network, the router, whose only job is to look at the model's internal state and answer one question: should this block be skipped, run once, or run twice?

To teach the routers, they first need examples of good routing choices. They run a search algorithm, Monte Carlo Tree Search, offline for each training question. The search explores different combinations of skipping and repeating blocks, evaluates each candidate path on the actual task, and keeps only paths that match or beat the default full-depth path while using fewer layers. A length penalty in the search score steers it toward shorter configurations, and the search stops early once it finds a correct answer with fewer layers. Constraints keep the search tractable: at most two consecutive skips, at most one repeat per block, and total path length capped at twice the original depth.

These winning paths become training labels for the routers, one label per layer (0 for skip, 1 for execute, 2 for repeat). Since "execute" dominates the labels, the routers are trained with focal loss and class-balancing weights so the rare skip and repeat decisions are not ignored. During training the ground-truth path is used to drive execution, so each router is graded independently rather than depending on what earlier routers chose. At inference, routers simply take the highest-probability action greedily, with no search. Router decisions are made from mean-pooled summaries of hidden states over several windows rather than per-token signals, which keeps the decision stable on long sequences and allows the router to run only once per input, preserving cache compatibility.

Why This Matters

• Research impact. Dr.LLM reframes adaptive depth as a supervision problem rather than an inference-time search problem. By showing that search-derived targets can be distilled into tiny routers, it separates the expensive discovery step from the cheap deployment step, and it demonstrates that adaptive computation can raise accuracy rather than trading it away. The structured early/middle/late routing pattern also offers a concrete empirical signal about which phases of a transformer stack perform which kind of work.
• Real-world applications:
  • Serving LLM APIs where per-query cost and latency matter, since routers shave wall-clock time and layers without changing the served model.
  • On-device or edge inference, where a frozen base model plus a small adapter avoids the memory and retraining budget of full fine-tuning.
  • Reasoning-heavy assistants in math and logic tutoring, where the router's tendency to repeat late layers helps multi-step problems.
  • Budget-aware deployment, where an operator can tune a single control parameter to trade accuracy against compute continuously.
• Industry relevance. The retrofit property is the main practical appeal: organizations with existing deployed models can add routers without retraining or re-quantizing the base weights, and the training cost fits on a single 40GB GPU in a few hours. The measured throughput gains and demonstrated batch-processing strategy address the usual objection that per-layer routing destroys GPU utilization, which is a prerequisite for production adoption.

Future Directions

1. Scaling the router training set and model range. The current routers are trained on 4k examples from two reasoning datasets and evaluated up to 8B parameters; whether the same recipe scales to much larger models and broader data mixtures is untested.
2. Improving router accuracy on rare decisions. Per-layer routing accuracy is 96.8% but macro F1 is only 61%, with most errors conservatively defaulting to execute. Better modeling of skip and repeat could unlock further layer savings.
3. Joint or hierarchical routing across tokens and layers. The paper notes that token-level methods like Mixture-of-Depth address local redundancy while layer-level control reallocates global compute; combining the two is an open direction.
4. Building continuous control into deployment schedules. The paper sketches a scalar control parameter that interpolates between skip, router, and execute policies. Turning that into a principled, calibrated mechanism for serving systems under variable load remains future work.

Target Audience

This paper is most useful to machine learning researchers working on inference efficiency, adaptive computation, or LLM compression, and to practitioners who deploy LLMs and care about latency, throughput, and serving cost. It also suits readers interested in reinforcement-search-style data generation for training small auxiliary modules, and those studying how transformer depth is used across layers. Readers without a background in transformer internals will find the core idea accessible, but the MCTS formulation, focal loss weighting, and routing analysis assume intermediate familiarity with modern NLP methods.

Authors’ abstract

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. Code is available at https://github.com/parameterlab/dr-llm.