💬 NLP

Dr.LLM: Dynamic Layer Routing in LLMs

Dr. LLM is a retrofittable framework that equips pretrained LLMs with lightweight, MCTS-supervised routers to dynamically skip, execute, or repeat transformer layers, achieving significant compute savings while improving or maintaining accuracy across diverse tasks without altering the base model weights.

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

Published 2026-05-20

📖 4 min read☕ Coffee break read

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant, overworked librarian (the Large Language Model, or LLM) who answers your questions. Right now, no matter how simple your question is ("What is 2+2?") or how complex it is ("Solve this advanced physics problem"), this librarian is forced to walk through every single room in their massive library, checking every bookshelf, before giving you an answer.

This is inefficient. For simple questions, they waste time walking through rooms they don't need. For hard questions, they might not have enough time to dig deep enough because they are stuck on a rigid schedule.

Dr.LLM is like giving this librarian a smart, lightweight traffic controller for each room in the library.

How It Works: The "Traffic Controller" Analogy

Instead of the librarian blindly walking through every room, Dr.LLM installs a tiny, fast decision-maker (a "router") at the entrance of every room. When the librarian arrives at a specific room, the router looks at the current situation and makes one of three quick decisions:

Skip: "This room isn't needed for this question. Let's walk right past it to save time."
Execute: "We need to do our work here. Go in, read the books, and come out."
Repeat: "This is a tricky part. We need to go in, do the work, and then go in again to double-check or refine the answer."

How They Trained the Traffic Controllers

The tricky part is teaching these routers what to do without hiring a new team of experts to rewrite the whole library (which would be expensive and slow).

The authors used a method called MCTS (Monte Carlo Tree Search). Think of this as a super-smart simulation team that ran thousands of "what-if" scenarios offline. They asked: "If we skip this room, does the answer get worse? If we repeat this room, does the answer get better?"

They found the perfect "paths" through the library for different types of questions. They then used these perfect paths to train the tiny routers. Once trained, the routers know exactly when to skip, go, or repeat, without needing to run simulations while answering your questions.

The Results: Faster and Smarter

The paper claims that this system does two amazing things at once:

It saves energy (Compute): On average, the librarian skips about 5 rooms per question. This makes the system faster and cheaper to run.
It gets smarter (Accuracy): Surprisingly, by skipping the boring parts and repeating the hard parts, the librarian actually gives better answers than before.
- On logic puzzles (ARC), accuracy went up by about 1%.
- On math problems (DART), accuracy jumped by up to 3.4%.

Does It Work on New Questions?

The authors tested these routers on questions they had never seen before (like general knowledge, reading comprehension, or different types of math). Even though the routers were trained on specific logic and math puzzles, they generalized well. The accuracy only dropped by a tiny fraction (0.85%), proving that the routers learned a general "style" of thinking rather than just memorizing specific answers.

The "Secret Sauce"

The paper highlights a few key features that make this work:

No Rebuilding: You don't have to tear down the library or rebuild the librarian. You just add these tiny routers on top of the existing system.
Windowed Pooling: The routers don't look at every single word one by one (which would be slow). Instead, they look at "chunks" or windows of the conversation to make stable decisions.
Smart Training: They used a special math trick (Focal Loss) to make sure the routers didn't just always say "Execute" (which is the safe, lazy choice). They forced the routers to learn when to actually skip or repeat.

In a Nutshell

Dr.LLM is like giving a rigid, one-size-fits-all robot a pair of smart glasses. Now, the robot can see which tasks are easy and which are hard, allowing it to skip the easy steps and double-check the hard ones. The result is a system that is faster, cheaper, and surprisingly more accurate than the original, without needing to be completely rebuilt.

Technical Summary: Dr.LLM – Dynamic Layer Routing in LLMs

Problem Statement

Large Language Models (LLMs) currently operate under a static-depth regime, processing every token through a fixed stack of transformer layers regardless of input complexity. This approach leads to two primary inefficiencies:

Wasted Computation: Simple queries expend unnecessary compute by traversing all layers.
Insufficient Flexibility: Complex reasoning tasks often require deeper processing or iterative refinement, which static architectures cannot provide without retraining or architectural overhaul.

Existing adaptive-depth methods (e.g., early exits, layer pruning, mixture-of-depths) face significant limitations: they often trade accuracy for speed, require extensive retraining or architectural modifications, or rely on costly inference-time search that hinders scalability. Furthermore, many prior approaches degrade model accuracy despite efficiency gains.

Methodology: Dr.LLM

The authors propose Dr.LLM (Dynamic Routing of Layers for LLMs), a retrofittable framework that equips frozen, pretrained LLMs with lightweight, per-layer routers. These routers decide whether to skip, execute, or repeat a specific transformer block for a given input.

1. Router Architecture

Structure: Each transformer block $B_\ell$ is paired with a lightweight Multi-Layer Perceptron (MLP) consisting of Linear-GELU-Linear layers.
Input: The router operates on a compact summary of the hidden states from the previous layer. To ensure stability on long contexts, the authors employ windowed mean pooling, dividing the sequence into contiguous windows and averaging the hidden states within each window.
Output: The router outputs logits for three classes: $\{skip, execute, repeat\}$ .
Overhead: Routers are executed once per input sequence. The overhead is negligible (constant relative to token count) and fully compatible with KV caching.

2. Supervised Training via MCTS

Unlike methods that rely on reinforcement learning or inference-time search, Dr.LLM trains routers using explicit supervision derived from an offline Monte Carlo Tree Search (MCTS).

Search Process: For each input, MCTS explores modified execution paths (skipping or repeating layers) under a compute budget. It retains only paths that preserve or improve task accuracy compared to the default full-depth execution.
Dataset: The search generates a compact dataset of approximately 4,000 optimized execution paths (tuples of question, optimal layer actions, and answer).
Training Objective: The routers are trained end-to-end on this dataset using Focal Loss with class-balancing weights to handle the severe imbalance between the dominant "execute" class and the minority "skip" and "repeat" classes.
Freezing: Crucially, the base LLM weights remain frozen; only the router parameters are updated.

3. Inference

At inference time, the trained router makes greedy decisions ( $\hat{y}_\ell = \arg\max p_\ell$ ) for each layer without any search.

Skip: Bypasses the layer ( $H^{(\ell)} = H^{(\ell-1)}$ ).
Execute: Applies the layer once.
Repeat: Applies the layer twice in succession.
Batching: To handle synchronization in batched processing, sequences are grouped by their routing decisions at each layer. The authors note that high routing consensus (80–89% agreement across sequences) minimizes synchronization overhead.

Key Contributions

Supervised Dynamic Routing for Frozen LLMs: Introduction of per-layer routers that decide to skip, execute, or repeat, trained end-to-end on only 4k MCTS-derived paths without modifying base weights.
Effective Path Supervision via MCTS: A length-aware MCTS algorithm that discovers accuracy-preserving or improving layer configurations offline, generating a compact supervision dataset without inference-time search.
Lightweight Router Training: A design featuring windowed mean-pooling for long-context stability, focal loss for class imbalance, and a compact two-layer linear router, ensuring negligible computational overhead.
Accuracy and Efficiency Gains: Demonstrated improvements in reasoning accuracy alongside reduced layer execution, outperforming prior state-of-the-art routing methods.
Robust Generalization: Showed that routing policies learned on specific reasoning tasks transfer effectively to out-of-domain benchmarks with minimal accuracy degradation.

Experimental Results

In-Domain Performance (ARC and DART)

Evaluated on six models (LLaMA-3.2 and Qwen-2.5 variants), Dr.LLM consistently improved accuracy while reducing the average number of executed layers:

Accuracy: Improved by up to +3.4 percentage points (p.p.) on DART (math) and +1.0 p.p. on ARC (logic). For instance, LLaMA-3B-Base saw a jump from 11.8% to 15.8% on DART.
Efficiency: Saved an average of 5.0 layers per example across models, with specific cases saving up to 11.0 layers.
Latency: Achieved a 15.3% reduction in wall-clock generation time for 1,000-token runs, with router overhead contributing less than 1% to total latency.

Out-of-Domain Generalization

When applied to diverse benchmarks (MMLU, GSM8k, AIME, TruthfulQA, SQuADv2, GPQA, AGIEval, PIQA) without retraining:

The routers maintained efficiency while suffering only a 0.85 p.p. average accuracy drop across eight benchmarks.
In some cases (e.g., GPQA Diamond with LLaMA-3B), accuracy actually improved (+2.5 p.p.).
This suggests the learned policies capture general structural redundancies in transformer computation rather than overfitting to specific task patterns.

Comparison to State-of-the-Art

Dr.LLM outperformed prior routing methods (e.g., FlexiDepth, MindSkip, ShortGPT) by up to +7.7 p.p. in accuracy on reasoning benchmarks. Unlike FlexiDepth, which requires training on hundreds of thousands of examples and often sacrifices accuracy, Dr.LLM achieves higher accuracy with significantly lower training overhead (4k examples, single GPU) and no architectural changes.

Routing Pattern Analysis

Analysis of routing decisions revealed structured behaviors:

Early Layers: Consistently executed (stable usage).
Middle Layers: Frequently skipped (economizing compute).
Late Layers: Often repeated, particularly for complex reasoning tasks, indicating a role in iterative refinement.
Feature Learning: Routers trained on previous-layer hidden states significantly outperformed those trained on static input embeddings, confirming that Dr.LLM learns to map internal model states to routing decisions dynamically.

Significance and Claims

The paper claims that Dr.LLM demonstrates that explicitly supervised routers can retrofit frozen LLMs to achieve budget-aware, accuracy-driven inference. Key takeaways include:

No Retraining Required: The base model weights remain untouched, preserving the model's original capabilities.
Accuracy-Positive Efficiency: Unlike many adaptive methods that trade accuracy for speed, Dr.LLM improves both simultaneously on reasoning-heavy tasks.
Practical Deployment: By removing the need for inference-time search and architectural modifications, Dr.LLM offers a practical, scalable solution for adaptive inference.
Generalizability: The learned routing policies transfer robustly across model families and task domains, suggesting that adaptive compute allocation is a fundamental, learnable property of transformer architectures.

The authors conclude that Dr.LLM provides a practical step toward scalable adaptive inference, reconciling efficiency, accuracy, and robustness without the prohibitive costs associated with retraining or architectural redesign.