Dr.LLM: Dynamic Layer Routing in LLMs
Authors: Ahmed Heakl
Martin Gubri
Salman Khan
Sangdoo Yun
Seong Joon Oh
Venue: ICLR 2026 Poster (2026)
Keywords: Large Language Models
Efficient Inference
Adaptive Computation
Test-time Optimization
Monte Carlo Tree Search
Dynamic Layer Routing
Primary area: foundation or frontier models, including LLMs
TL;DR: Dynamic Layer Routing in LLMs

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.

Published as a conference paper at ICLR 2026
DR.LLM: DYNAMIC LAYER ROUTING IN LLMs
Ahmed Heakl 22 Martin Gubri', Salman Khan', Sangdoo Yun", Seong Joon Oh! 0h..4,5 Parameter Lab, 'MBZUAI, "NAVER AI Lab, University of Tübingen, Tübingen AI Center https:/ htw.:/abb.coompprrmtteeellaaaa/a 11m
ABSTRACT
Large Language Models (LLMs) process every token through all layers of a trans- former stack, causing wasted computation on simple queries and insufficient flexi- bility for harder ones that need deeper reasoning. Adaptive-depth methods can im- prove efficiency, but prior approaches rely on costly inference-time search, archi- tectural changes, or large-scale retraining, and in practice often degrade accuracy despite efficiency gains. We introduce Dr.LLM, Dynamic routing of Layers for LLMs, 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), high-quality layer configurations that preserve or improve accuracy under a com- pute budget. Our design, windowed pooling for stable routing, focal loss with class balancing. and bottleneck MIP routers, ensures robustness under class im- balance and long sequences. On ARC (logic) and DART (math), Dr.LLM im- proves accuracy by to +3,4%p while saving 5 layers per example on average. Routers generalize to oft-of-domain tasks (MMLU, GSM8k, AIME, TruthfulQA, SQuADv2, GPQA, PIQA, AGIEval) with only 0.85% accuracy drop while retain- ing efficiency, and outperform prior routing methods by up +7.7%p. Overall, Dr.LLM shows that explicitly supervised routers retrofit frozen LLMs for budget- aware, accuracy-driven inference without altering base weights.
derive
INTRODUCTION
Large language models (LLMs) typically process every token through fixed stack of transformer layers, regardless of the input's difficulty. This static-depth regime results in wasted computation for easy prompts and insufficient flexibility for challenging reasoning tasks. To address this, prior work has investigated adaptive depth mechanisms at test time, including early-exit strategies (Elhoushi et al.] 2024), layer pruning (Men et al. 2024), recurrent or looped blocks (Bae et al.
Accuracy Dr.LLM
41%
Vanilla LLM
40%
Layer-wise routing
2025) >
Dynamically skip repeat layers
dynamic routing methods (He et al.) 2024 Luo et al. 2025
mixture-of-depth approaches (Raposo et al. 2024). mixture- of-experts architectures (Shazeer et al. 2017), and search- based routing frameworks (Li et al.] 2025). Despite their promise, these methods typically suffer from one more lim- itations: (i) they trade accuracy for speed, (ii) they require ar- chitectural modifications and retraining on substantial amounts of data, or (iii) they rely on costly inference-tine -ieeeeee-imeeeee search that is difficult to deploy at scale.
# #Layers used
Figure 1: Dr.LLM improves ac- curacy while redueing computa- tion. Number of layers used per example vs. accuracy on ARC and DART, averaged on six models.
We propose Dr.LLM (Dynamic Routing of Layers for LMs),
a retrofittable framework that equips a frozen, pretrained L.M with lightweight, per-layer routers that decide whether to skip, execute, repeat their layer. Routers operate in windowed, mean- pooled hidden states and are trained with explicit supervision obtained from an offline Monte Carlo Tree Search (MCTS). For each input, MCTS discovers execution paths, that is, which layers to skip or repeat, such that they preserve or improve task accuracy under a constrained compute budget.
"Corresponding authors

---

Published as a conference paper at ICLR 2026
Training the routers with supervised learning on 4k optimized paths is lightweight, since only the few router parameters are updated while the LLM remains frozen. The trained router removes the need for any search at inference time, enabling compute-efftcient inference that increases accuracy, without modifying the base model weights.
Empirically, Dr.LLM improves the accuracy of reasoning- tasks while reducing the average number of executed layers (Fig. [ where "number layers" denotes the sum of active layers across all forward passes, layer applications in total. Also, Dr.LLM cuts generation time by 15.3% on Ik-token runs, with router overhead under 1.. On ARC (logic) and DART (math). accuracy improves in all cases, with mean gains of +2.25 percentage points (%p) and 5.0 fewer layers per example across six models. Routers generalize out-of-domain (e.g. MMLU, GSM8K, AIME, TruthfulQA, SQuADv2, GPQA, AGIEval, PIQA) with only 0.85%p average accuracy drop while retaining efficiency, indicating that learned routing policies transfer beyond the supervised domains. Lastly, Dr.LLM outperforms all prior SoTA routing methods by up to +7.7% % p accuracy. Our contributions are as follows:
32-layer model generating long sequences accumulates several hundred
Supervised dynamic routing for frozen LLMs. We introduce per-layer routers that decide to either skip, execute, or repeat their layer. We train the routers end-to-end on only 4k execution paths optimized for accuracy, discovered offline.
Effective path supervision via MCTS. We present a length-aware MCTS to find layer edits (skips/repeats) under budget and to retain only eecuracy-peeervvng or improving paths, gener- ating a compact supervision dataset without modifying the base weights.
Lightweight router training. We propose windowed mean-pooling for stable decisions on long contexts and use focal loss with dlass-nbbllnncing weights, combined with lightweight two- layer linear model, which together handle class imbalance and keep the computation overhead negligible.
Accuracy increase and compute efficiency. On ARC and DART across six models, accuracy improves in all cases, with up to +4.0%p and 11.0 layers saved per example in that case, without architectural changes, retraining, or inference-time search.
Robust generalization. Routers transfer to out-of-domain benchmarks with only a 0.85%p aver- age drop, showing that policies learned during training remain useful beyond the training tasks.
RELATED WORK
Adaptive-depth methods span pruning, early exits, recurrence, and routing; Table condenses their trade-offs across accuracy, retrofitting, efficiency, practicality, and frozen-base compatibility.
Pruning and Early Exit. Classical model compression prunes redundant weights, heads, or layers post hoc (Sajjad jet al. 2023). Early-cxit networks extend this by attaching auxiliary classifiers at intermediate layers (Xin et al.; 2020 Zhou al. 2020 Teerapittayanon
Table I:
Comparison of dynamic routing methods for frozen LLMs. Accuracy does the method improve accu- racy over the baseline. Retrofit: can it be added to pretrained models with minimal effort. Cheap I: enables efficient in- ference without heavy overhead. Cheap T: enables efficient training with limited data. LLM $: base model remains changed. Symbols: /= strong support, X= not supported. Lr.L the only method satisfying all five criteria.
al. 2016). letting easy inputs terminate early. While effective, such classifiers need calibration, add overhead, and complicate deployment. LayerSkip (Elhoushi al. 2024) improves this by training with dropout and a shared exit loss, removing the need for multiple classifiers. Yet still requires finetuning or training from scratch for the LLM and cannot repeat layers. In contrast, our approach supervises skip/execute/repeat directly via MCTS, eliminating auxiliary exits and enabling repetition without retraining base weights.
Method
Accuracy 1 Retrofit Cheup I Cheap T LLM $
CoLa
/
Mixture of Depths Universal Transformer
X
X
X
/
X
LLM-Prunet
X
X
X
Mixture of Experts Mixture of Recursions
/
X
X
LayerSkip
X
X
MindSkip FlexiDepth
x
X
/
/
DrL.IM (Ours)
/
/
/
2

---

Published as a conference paper at ICLR 2026
Figure 2: Our layer routing based on hidden states. Dr.LLM augments a frozen decoder-only LLM with
What
is
x
2?
The
ans
wer
4
Router Internals
Router Decision
Layer-wise Routing
X N layers
per-layer routers that decide to skip, execute, repeat a block once.
Average
Transformer Block
Routers read windowed summaries of hidden states and are trained from MCTS- derived targets (Sec. 4
>>> Skip
V Execute
C Repeat
Linear év
GELU
Linear e
Router
For clarity, the diagram also highlights the router internals and the flow of hidden states across layers.
Nww windows
Avarage
WeST
Input Tokens
Output Tokens
Input takens
Recurrence and Looped Architectures. Another line work adapts depth by repeating compu- tation. Universal Transformers (Dehghani et al. 2018) learn halting policy per token, while looped transformers Yang et al. [2023] Giannou et al. [023] Geiping et al.) 2025) iteratively reapply blocks for refinement ("slow thinking") These models are flexible but require architectural redesign, full retraining, and incur higher inference cost. We also support targeted repetition, by attaching shallow controllers to frozen layers, avoiding structural changes or pretraining. Moreover, we allow skips to offset the layer increases from looping.
Dynamic Routing and Modular Inference. Routing- based methods let inputs select modules dy- namically. MoE architectures (Fedus et al.) 20222 Shazeer et al. 2017) expand capacity by routing tokens to experts, but demand large-scale retraining. CoLa (Li al., 2025) is closer our setting: treats pretrained layers as modules and searches, via MCTS, for input-specific "chains of lay- ers." However, CoLa requires costly search inference and, critically, access to gold labels during search to decide which path is "correct," making it impractical for deployment. We instead perform MCTS offline to generate supervision and then train routers that make decisions cheaply at infer- ence. Other adaptive-depth methods, such FlexiDepth (Luo et al. (025) and MindSkip (He et al.) 2024), retrofit routing to pretrained models but require extensive training (hundreds of thousands of examples) and often reduce accuracy save compute. By contrast, our routers are trained from only 4k MCTS-derivee examples and in experiments improve accuracy while lowering cost. Mixture-of Depth (MoD) (Raposo al. 2024) takes different angle, routing at the token level by sending only a subset tokens through deeper layers, but modifies the base weights. This intra-layer mecha- nism complements sequence-level skip/execute/reeeat routing: token-level signals identify local redundancy, while layer-level control reallocates global compute.
SUPERVISED TRAINING OF THE ROUTER
Let a pretrained decoder-only LLM with L transformer blocks be M = [B]...., sequence of length T, let (1) E Txd denote its initial hidden states. The classical forward pass applies each block once: (0) By (((()))) We instead seek a discrete per-layer policy
BB). For token
Ve (skip, execute, repeat).
where skip bypasses B3, execute applies it once, and repeat applies twice succession. The vector y = (V1...., induces a custom execution path, while we freeze the base weights.
3.1 ROUTER ARCHITECTURE
As shown in Figure each block B, paired with a lightweight MLP LLinear-GLLU-Lineall re : Rd 3 which outputs logits for {skip, execute, repeat). The router operates compact summary of the hidden states (f- from the previous layer. Routers executed once per input sequence at inference, adding negligible overhead (constant with the number of generated tokens), and remaining fully compatible with KV caching, unlike most layer routing methods.
3

---

Published as a conference paper at ICLR 2026
Training Datasets
Monte Carlo Search Algorithm
Output
ARC, DART
Selection
Expansion Accion sion
Evallaation
Update
Contigurat on highest
Sipasioo ofgeete3
Search for earh esample
UCB score
Compute reward
Pro pagaze reward
Layer
Action
Action Space
One node Wig for layer
Layer iin [1.N]
Skpp Exctute
Sep lyyr
Expandoon don
Repeat
mubioos maximum
Figure 3: Length-aware MCTS used collect the supervised training dataset of per-layer routing For each input, MCTS explores modified layer paths and retains securacy-preservvin or improving ones under compute budget.
configurations
To stabilize decisions on long contexts while keeping overhead negligible, we adopt windowed mean pooling: the first ] tokens are divided into contiguous windows [ with
my = 1 [S..] Eies H (t-1) the mean-pooled representation. Router votes are aggregated by aver-
aging logits:
W
W Σ re(m))) ==
Pf softmax(=))
DP arg max PL,e- cE(0.1,2)
We default to W-8 (clamped by T). Router weights are Xavier-uniform initialized (Glorot & Bengio 2010) and biases initialized to zero. Only {ne}/_ are trainable, while model parameters are frozen. At inference, the decision de governs block execution where SKIP passes H(9) (t-11), EXECUTE applies B, once, and REPEAT applies B3 twice in succession.
3.2 TRAINING REGIME
For each prompt-response pair cresstion-aswarr pair (g. a). the method proposed Sec. yields "ground truth" path, x², that we utilize to supervise the training process introduced here. These paths preserve or improves task reward under a compute budget. We convert x to per-layer labels Vi co *)) E (0, 1, producing tuples (g,y*, .a) meaning (question, y", answer).
Because execute e dominates, we apply focal loss (Lin et al. 2017) with effective-number weights. Let the global class counts be ne for (skip, exec, repeat) and E (0, (0,
0g
1 1-B B ne
Σ 1
1-3 3
= Σ
L
Oyt (1 Pt.p) log
We use 7=2, 8-0.999 by default] During training we apply teacher forcing for execution only, ie., we replace the router decision with the ground-truth label ÿr + to follow labeled path while supervising logits with L. This avoids making router, depend on router,-1 outputs, which would otherwise slow training and lower accuracy by 1.7%. At inference, decisions are greedy:
ÿe arg max Pe; search is used.
Routers add O(Lh)) parameters for hidden size and router width (we use h=128), and one small MLP per layer at inference. Windowed pooling is linear in T and inexpensive relative to transformer block. Skip reduces compute; repeat adds targeted compute when beneficial. We report accuracy, per-class fl for {skip, exec, repeat), }, and average executed layers.
TRAINING DATA GENERATION VIA MCTS
This section describes the search-based generation of the supervised training dataset of layer con- figurations for the router. We supervise routers using tuples (g. y" a), where g the input ques- tion/prompt, a the gold answer, and (skip, execule, repeat) the best per-layer routing targets found by discrete search over modified forward passes (Fig. 3 The search is offline and does not modify base weights.
Setting 7= recovers weighted cross-entropy.
4

---

Published as a conference paper at ICLR 2026
Algorithm Length-aware MCTS for question answer pair (q.
Require: Default path ло .I,....
... LJ, number of simulations Ns. constants (c, A, Prant) 1: Create root r with л (r) <- no: cache & <- 23; best path xt" f- 2
2: for n==== 1 to N2 do
3: Select: traverse from to a leaf using nodes UCB score; w.p. Prand pick random child
4: Expand: add one untried action to obtain child u (respecting path-length cap)
5: Evaluate: if (uu) & &, model constrained to (u)) set ((]( R(
6: Backpropagate: propagate to ancestors
7: if R( =1 and [T["]] < <|=| or xt then
8:
- π (
> update to the shortest correct path b enhance default path answer W-+
9: break if =f
10: end if
end for
12:
Convert n² into per-layer labels y* E { 1, (skip/executerreeeatt
13: return y"
4.1 EDITED EXECUTION PATHS AND ACTIONS
Let the base model have L blocks and default path Tu = [L..... L]. An edited path [e......4 Lx] preserves the original order of blocks but may omit certain layers (skip) or apply a given layer twice (repeat once). We allow skips of at most two consecutive layers, and we allow most a single repeat for any block, which controls the compute growth, i.c., the total edited path length is capped at | 2L.
4.2 LENGTH-AWARE MCTS
Each node stores triple of path visits v(Tr), and cumulative reward (Ir). During selection, we maximize a UCB (Upper Confidence Bound) score (inspired by [i et al. (2025) with an explicit length penalty to favor compact paths:
(
In
(T) L
UCB(T)=
((T)
exploitation
exploration
length penalty
where V is the parent's visit count. We use c=1.8, 333.0, and with probability Prand=0.1 pick a random child to encourage exploration. For each simulation, we expand one untried action, evaluate the edited path once, and backpropagate the task reward R E [0, 1] (no length penalty) through the root node. We run a fixed budget N,=50 simulations or stop early if (i) reach correctness and (ïi) find a strictly shorter correct path than the best-so-far. Path evaluations are mem- orized to avoid duplicates. We retain only accuracy- preserving/improinng paths (vs. the default To).
Table 2: Data generation statistics. Vis- ited is the total number of candidate paths explored, and Sampled is the subset of paths that improve or preserve accuracy.
Dataset
Original
Sampled
Visited
ATnferences
ARC-E
2.25k
2,090
82.6k
ARC-C
1.12k
600
1.119
44.2k
DAKT-1
117k
200
967
38 28.2k
DART-2
296k
400
2,242
88.6k
DART-3
364k
600
3,695
146.0k
237.6k.
We ran MCTS across ARC and DART, collecting 4k supervision examples. About 30% of these edited paths achieved higher accuracy than the default path while the rest preserved accuracy and reduce the num- ber of layers with exact statistic shown in Tab. 2 The average number of layers saved is 1.82. Although the search required 961.0k forward passes, it is performed entirely offline; at inference time, routing decisions made directly by the trained routers without any search.
DAKT-5
445k
1000
8,203
324.0k
Total
1.63M
4000
24,330
961.0k
Compared to Li et al. (2025), we found that reducing the repetition block size from 4 to 1 made the search substantially faster, while achieving the same accuracy gains and layer savings with only 50 simulations instead 200. We also found that lowering the length penalty from 5 to reduced the number of search samples by 14.8%p
5

---

Published as a conference paper at ICLR 2026
5 EXPERIMENTS
We evaluate Dr.LLM across both in-domain reasoning tasks and a diverse suite of out-of-domain (OOD) benchmarks test generalization under distribution shift. Our experimental setup is de- signed to answer three key questions: (i) Does supervised dynamic routing improve accuracy relative to static baselines? (ii) How much computational efficiency is gained in terms of average executed layers? (iii) Are the learned routing policies robust to new tasks and model families?
Models. We retrofit Dr.LLM onto six backbone models spanning two families: LLaMA-3.2 (Dubey et al. 2024) (3B Instruct, 3B Base, 8B Instruct, 8B Base) and Qwen-2.5 (Yang et al.) 2024) (3B Instruct, 7B Instruct). These models cover variety of sizes and both instruction-tun and base variants.
Training Data. Routers are supervised using 4K MCTS-derived tuples (Sec. 4p from ARC- Easy/Challenge (Clark et al.] (018)) and DART-Math (Tong al., 2024). We selected these datasets for three reasons: (1) they provide stratified difficulty levels (ARC-Easy vs. ARC-Challenge, DART- to DART-5), (2) they target logic and multi-step mathematical reasoning, where adaptive compu- tation is beneficial, and (3) they have train/test splits allowing us test in-domain distributions.
Training Setup. We train all routers on a single NVIDIA A100 40GB GPU. Given the small number of trainable parameters (1IM for 3B models, 0.14% of base weights; 16.8M for 8B models, 0.56%), training is efficient and completes within 4 hours while using only 20% of the GPU VRAM. We use Adam W (Loshchilov & Hutter 2017) with cosine schedule, learning rate x 10 10-3, weight decay 0.01, 500 warmup steps, and total of 25 epochs. The effective batch size is 16, and training is performed in 6 precision. Routers achieve macro F1-scon of 61%, reflecting class balance across the highly imbalanced actions (skip: pp 7.6%, repeat: 2.3%), while the actual per-layer ac- curacy reaches 96.8%. Most errors conservatively predict execute, slightly increasing compute but not harming accuracy, which explains why Dr.LLMM consistently improves task performance de- spite the moderate Fi-score. We tested different initialization schemes for router biases, (i) empiri- cal class frequencies and (ii) zero- cero-initialization, finding the latter more stable and yielding stronger downstream accuracy.
Wall-clock latency and router overhead We measure end-to-end generation time for sequences of 1,000 tokens. Dr.LLM achieves a 15.3% wall- tile reduction (29.21 vs. 34.49 s baseline), while the router adds only 0.27 s per query (<1% of total latency).
In-Domain Evaluation. We first evaluate routers on ARC and DART test splits. These tasks serve direct measure of whether routers recover the MCTS supervision signal and yield improvements under controlled conditions.
Out-of-Domallo Evaluation. To assess robustness, we evaluate the router-equipped models on a broad range of benchmarks: MMLU (Hendrycks et al. (020)) for factual knowledge, GSM8k (strict_matchh (Cobbe et al., (021) for grade-school math, TruthfulQA (meql) (Lin et al. 1021 for adversarial factuality, GPQA Diamond (Rein al., 2024) and AIME24 (MAA) 2024) for challeng- ing mathematical reasoning, AGlEval (Zhong et al. 2023) for exam-style reasoning, SQuADv2 (f1) (Rajpurkar et al. 2018) for reading comprehension, and PIQA (Bisk al., (020) for commonsense reasoning. All benchmarks are reported using acc_norm computed from log- likelihoooos in the 1m-eval-ha rness framework (Gao et al. 2024), except GSM8k, TruthfulQA, and SQuADv2 which follow their respective metries. Evaluations tion length 2048 tokens, and greedy decoding.
with default settings, maximum genera-
6 RESULTS & DISCUSSION
We evaluate Dr.LLM on in-domain tasks, test its robustness on out-of-domain benchmarks, and analyze routing patterns with ablations.
6.1
IN-DOMAIN PERFORMANCE ON ARC AND DART
Table summarizes in-domain results on ARC (logic) and DART (math), showing that routers consistently improve accuracy while reducing the average number of layers executed across all six

---

Published as a conference paper at ICLR 2026
Table 3: Routers consistently improve accuracy and reduce executed layers across all models. In-domain results on ARC (logic) and DART (math). Accuracy in %. (+x) indicates accuracy gains, (-x) indicates layer savings, and (+0.0) indicates no change.
ARC
DART
Total
Model
Accuracy
Num Layers
Accuracy
Num Layers
Accuracy
Num Layers
LLaMA-3B-Instrrett
73.5
103.9
35.2
422.0
46.1
331.1
+ Router
74.5 (+1.0)
99.50 (-4.25)
38.6 (+3.4)
413.3 (-8.66)
48,9 (+2.7)
323.7 (-7,40)
LLaMA--B-Instrrett
88.5
106.9
38.4
320.0
52.7
518.2
Router
89.4 (+0.9)
104.0 (-2.94)
41.2 (+2.8)
309.1 (-10.96)
54.7 (42.3)
509.6 (-8.66)
LLaMA-3B-
48.0
56.0
11.8
548.4
22.1
815.4
+ Router
49.0 (+1.0)
55.70 (-0.28)
15.8 (+4.0)
544.3 (-4.12)
25.3 (+3.2)
812.4 (-3.02)
LLaMA-8B-BaL
22.5
56.0
17.2
536.7
18.7
798.7
Router
23.5 (+1.0)
55.60 (-0.42)
20.2 (+3.0)
531.0 (-5.74)
21.1 (+2.4)
794.5 (-4.22)
Owen-3B-Instnne
53.0
115.7
30.2
536.3
36.7
832.4
Router
55.5 (+2.5)
115.5 (-0.23)
32.4 (+2.2)
531.7 (-4.55)
39.0 (+2.3)
828.9 (-3.31)
wven-7B-Instruce
94.5
112.0
45.4
277.8
59.4
460,9
Router
94.5 (+0.0)
111.8 (-0.20)
46.8 (+1.4)
273.1 (-4.67)
60.4 (+0.9)
457.5 (-3.39)
Table Generalization out-of-domain benchmarks. Accuracy in %. Router models maintain accuracy with 0.85%p average drop while preserving efficiency. All evaluated models are instruct, we use LLaMa-3.2 and Qwen2.5. TQA is TruthfulQA, and GPQA D is GPQA Diamond.
Model
MMLU
AIME24
TQA
GSM8K
SQuADv2
GPQA D
AGlEval
PIQA
Avg. A
LLaMA3B
60.5
3.3
31.32144
64.9at.2
32.6s1.1
27.2
35.72i51
75.6s1.06
+ Router
49.5an
3.341.64
30.42131
64.3a1as
30.6x142
29.8
33.8
71.92t.7
-0.94
LLaMA8B
67.92072
6.7a1.7
36.92145
73.2:130
29.1sass
34.32031
43.2.0.2
80.9a1.06
+ Router
66.810
6.711.4
36.6.1.40
74.9
28.61033
32.3u0
41.5.00i
79.2x1.07
-0.70
Qwen3B
65.32082
6.7a1.36
41.9a1.50
11.1
21.52009
33.32034
54.2a0.51
78.1.
Router
62. Sann
6.741.38
41.9a1.4
11.5at.29
20.
32.4
49.42)
78. 78.941.04
-1.05
Qwen7B
71.7.083
10.0
47
75.6
20.8
32.8.a
61.2.00s
79.7
+ Router
71 71.2408
10.0
47.9
75.741
20.240.3
32.8
S7.2.055
78.8.0.22
-0.70
models. On ARC. gains are modest (+0.9-2.5%p), reflecting that logic questions already require relatively shallow reasoning. contrast, DART exhibits larger improvements (+ (+1.4-4.0%p), where the router often assigns repeat to late layers. effectively allocating more computation to iterative refinmment needed for multi-step math problems. For example, LLaMA-3B. 3B-Base improves from 11.8% to 15.8% accuracy (+4.0%p) while saving 4.12 layers per query on average, and Qwen-3B- Instruct gains +2.2%p while cutting 4.6 layers per query. Notably, instruction-tunet models start with substantially higher accuracy than their base counterparts, yet still benefit from routing: e.g., LLAMA-8B--ME ---nIBrrIIttttt improves by +2.8%p on DART while saving 11.0 layers per query on average. Importantly. Dr.LLM never degrades accuracy and always saves inference compute with 3-11 fewer layers per query. These results demonstrate that Dr.LLM not only reduces computation but also improves accuracy, with the largest benefits tasks requiring deeper or repeated reasoning steps.
6.2 GENERALIZATION TO OUT-OF-DOMAIN BENCIMARKS
Table evaluates Dr.LLM-euuppped models on diverse suite of out-of-distribution benchmarks, from in-domain and oft-of-distributioo mathematical reasoning benchmarks (AIME24, GSM8K) to out-of-domains benchmarks specialised in knowledge (MMLU, AGIEval, GPQA Diamond), factu- ality (TruthfulQA), comprehension (SQuADv2), and commonsense (PIQA). Despite not trained to handle these types of questions, the routers maintain a good generalization with 0.85 % average ac- curacy drop across the eight benchmarks and four instruct models. The routers decision generalizes to other in-domain benchmarks: all four models gain an average 0.40%p accuracy on GSM8k and maintain the exact same accuracy AIME24, while reducing compute. In oft-of-domain bench- marks, the accuracy drop is limited to 1.20%p on average. Notably, in some cases routers even improve accuracy, such GPQA Diamond with LLaMA-3B (+25%%). In all cases, the router main-
7

---

Published as a conference paper at ICLR 2026
1.12
24-
1.05
100
1,05
LaVe Erequee
0.95
0.02
0,90
1,00
24
0.80
Liger COOD
5.50
0.85
0.75
cadyy
0/p4
t
mp
0.80
20
de
(a) Heatmaps (LLaMa 3B and 8B)
(b) Boxplot (LLaMa 3B)
(c) Boxplot (LLaMa 8B)
Figure 4: Analysis of routing decisions per layer, dataset, and model. (a) Layer frequency of LLaMa 3B and 8B base (B) and instruct (/) models across ARC and DART. (b,c) Layer frequency grouped by early, middle, and late layers. The x-axis corresponds to the dataset difficulty levels: ARC-Easy (A-1), ARC-Challeng (A-2), and DART levels 1-5 (from D-1 to D-5).
tains its efficiency by saving layers. These results indicate that router policies transfer beyond their domain, suggesting that the learned skip and repeat patterns capture general structural redundancies in transformer computation. Thus, Dr.LLM not only yields efficiency and accuracy improvements in-domain, but also preserves robustness when deployed to unseen, distribution-shiff benchmarks.
6.3 COMPARISON TO EXISTING METHODS
Most adaptive-depth approaches either sacrifice accuracy for efficiency impose costly architectural changes. For exam- ple. in Tab. 55) FlexiDepth (Luo et al. 2025) saves four layers on LLaMA-8B but suffers -6.1%p accuracy drop on GSM8k, while MindSkip (He et al. 2024) reduces compute yet loses 7.8% manEval. ShortGPT (Men et al., 2024)
Table 5: Comparison of Dr.LLM with existing meth- ods on reasoning and coding benchmarks. Results on LLaMa3-8B reported from FlexiDepth (Luo et al.) 2025) with layers saved. Although these benchmarks are in-domain for prior methods and out-of-domain for Dr.LLM, ours still achieves the highest accuracy.
Hu-
Method
GSM8K
MMLU
HellaSwag
HumanEval
Avg.
LayerSkip
0.4
65.9
63.6
0.0
32.5
also improves efficiency but underper- forms on reasoning, reaching only 53.6%
ShonGPT
53.6
9.2
MindSkip
37.8
66.4
69.8
18.9
48.2
on GSM8k compared to Dr.L.LM's 74.9%.
FlexiDepth
65.7
66.3
74.3
32.3
59.7
Even FlexiDepth, the method closest in ac-
Dr.LMM
74.9
66.8
79.3
48.6
67.4
curacy to Dr.LLM, requires training: is trained on Tulu-v2 (Ivison et al. 2023) with 326k examples, incurring substantial compute. By con- trast, Dr.LLM achieves higher accuracy with far lower overhead, trained on only 4k MCTS-derived examples using single GPU, despite the fact that these benchmarks in-domain for prior routing methods but out-of domain for Dr.LLM.
6.4 ANALYSIS OF LAYER ROUTING PATTERNS
We analyze router decisions across layers, models, and datasets to identify which layers can be skipped and which improve accuracy when repeated. Fig. visualizes the learned routing policies for LLaMA 3B and 8B models. The heatmaps (Fig. 4al show structured patterns rather than ran- dom skipping: early layers are consistently executed close to once, middle layers are frequently down-weighted, and late layers are often repeated, especially on reasoning-intensive DART tasks. The boxplots (Figs. (cb]4c) confirm this trend: early layers exhibit the lowest variance in execution frequency (stable usage), middle layers show wider skip distributions, and later layers are biased toward repeated execution, indicating their role in iterative refinement. This effect is stronger in the 8B model, where late-layer repetition dominates, suggesting that larger models rely more heavily additional depth for complex reasoning. Together, these results indicate that Dr.LLM learns routing behaviors aligned with transformer computation phases: maintaining stability in early input process-
8

---

Published as a conference paper at ICLR 2026
ARC
DART
-ap Total
Layer Savings
Number burb of Layers Saved
0.3
0.3
0.16
Accuracy In
0.2
0.1
256 (a) Bottieneck Dimenston
1024
Number Linear Layers
(c) Number Windows
Figure 5: Ablation study. We apply Dr.LLM on LLaMa3.2-31 and control: (a) the effect of bottle- neck dimension, (b) the effect of number of linear layers, and (c) the effect of number of windows.
ing, economizing in middle layers, and reinvesting compute in later blocks where deeper reasoning most beneficial.
6.5 ABLATION STUDIES
Router internals. We ablate the router components to understand their effect accuracy and efficiency (Fig.5) Varying the bottleneck dimension (Fig. (5) shows that smaller hidden sizes (64- 128) strike the best balance: bottleneck of 128 yields the highest accuracy gains (+3,4%p), while larger dimensions reduce both accuracy and layer savings, likely due to overfitting. Next, tuning the number of linear layers (Fig. 5)) indicates that both accuracy and compute gains best when the router is composed of two linear layers. Deeper routers fail to improve routing, confirming that compact routers more suitable. The number of pooling windows (Fig. (5)) strongly influences accuracy gains: more windows consistently increase both the accuracy and the number of layers saved. Averaging the hidden states of all input tokens is a signal that is too coarse for the router to learn. Finally, the focal loss better accounts for the class imbalance of the router labels than the weighted cross-entropy loss (+1.1%p. on ARC and + +1.8%p. on DART). These trends highlight that Dr.LLM benefits from (i) compact router architecture, and (ii) windowed contexts to learn fine-grained hidden state features.
What do the routers learn? Since the semantics of the question is predictive of LLM accuracy in this question (Ulmer et al.) 2024), ask if the routers truly Iearn from the internal state of the model or from the types of input. Routers could learn ques- tion patterns (e.g. skip the seventh layer for math questions). Table reppots the ARC and DART ac- curacies of a router trained on input embeddings (for all layers), rather than on the hidden states of the pre- vious layer. This new router performs considerably worse than Dr.LLM (-8.6%p on DART), and even worse than the vanilla model without layer routing (-5.2%)p. Therefore, Dr.LLM learns to dynamically map the internal model states to the decision to skip or repeat layers, instead of relying on shallow static signals from the inputs.
Table 6: Dr.LLM routes layers from their
state, not from the type of question. Benchmark accuracy of routers trained on the hidden states of the previous-layer or of the first layer. In %.
Router Features
ARC
DART
Prev. layer Hi-1 (Dr.LMM)
74.5
38.6
First layer H2 (emboddings)
70.9
30.0
No routing (vanilla model)
73.5
35.2
6.6 BATCH PROCESSING EFFICIENCY
A natural concern with per-layer routing is whether skip/repeat actions disrupt batch process- ing. We address this by synchronizing per-layer: sequences are grouped by their routing decision (skip/execute/repeatt)) at each layer and processed together as a microbatch. When repeat decision occurs, all sequences must join the group to complete both passes through the layer before pro- ceeding. However, the synchronization overhead is low because repeat actions occur in only 2.3% of layers. The high routing consensus (80-89% agreement) ensures that most layers process the majority of sequences in single unified batch, maintaining high GPU utilization.
Table reports throughput measurements (tokens/second)) on ARC-E, ARC-C, and DART 1-5 datasets using LLaMA-3B-Insture
Dr.LLM achieves an average improvement of 8.65% across all batch sizes. Importantly, Dr.LLM always increases throughput, showing that synchronization overhead is not problem for large batch sizes.

---

Published as a conference paper at ICLR 2026
Table 7: Dr.LLM consistently improves throughput across all batch sizes. Throughput mea-
surements (tokens/second) on ARC-E, ARC-C, and DART 1-5 using 1--MA-----nstrucce Despite per-layer synchronization requirements, Dr.LLM achieves 8.65% average improvement.
Batch Size
Original Throughput
Dr.LLM Throughput
Improvement
25.05
31.78
26.89%
2
37.97
41.87
10.29%
70.74
76.12
7.60%
8
107.74
115.56
7.25%
16
117.86
122.89
4.26%
32
125,43
130.83
4.31%
64
103.18
108.05
4.72%
128
109.12
113.38
3,90%
Average
87.14
92.56
8.65%
Decision Consensus. The key insight explaining why synchronization overhead remains minimal that sequences naturally exhibit high routing consensus. We measured inter-prediction similarity across all six models in our study, finding that routing decisions agree 80-89% of the time across dif- ferent sequences (Table R). This high similarity means that at most layers, sequences make identical skip/execute/repeat decisions, allowing our microbatching strategy group them efficiently.
Table 8: Routing decision consensus across models. Inter-predie inter-prediction similarity measures the percentage of identical skiplexecute/repeat
121.4
Throughpur (totemeree))
decisions made by routers across different quences.
100
90.4
Model
Similarity (%)
80-
LLaMA-3B-Instree
80.47
LLaMA-8B-
86.20
Inter-Prediction Similarity (1%)
LLaMA-3B 33--Bsee
80.86
LLaMA-88 LaMMA-B--Base
83.49
Figure Batching efficiency increases with rout- ing consensus. Higher inter-prediction similarity (agrecment in skip/cxcccteerreeatt decisions across sequences) translates to improved throughput.
Qwen-30 3B Instructtt
89,41
8B-Instruct
85.26
This consensus directly drives batching efficiency. Batches with 95% consensus achieve 100.28 tokens/sec, while perfect consensus reaches 121.58 tokens/sec (Table (), demonstrating that higher routing agreement naturally improves parallel processing. Despite per-layer synchronization, high routing consensus and infrequent repeats enable Dr.LLM to achieve both dynamic depth allocation and parallel efficiency. Computational savings from skipped and optimally routed layers outweigh synchronization overhead, yielding consistent throughput improvements across all batch sizes.
CONCLUSION
We introduced Dr.LLM, retrofittable framework that equips frozen LLMs with lightweight routers for skip/execute/repeat decisions. Supervised on high-quality paths from length-aware MCTS, Dr.LL.M removes inference-time search and architectural changes while improving both efficiency and accuracy. On ARC and DART, it yields up +3.4%p accuracy with 3-11 layers saved per query, outperforms prior routing methods by up to +7.7%p, and generalizes to out-of-domain benchmarks with only 0.85%p drop. Routing analysis reveals structured patterns, early layers preserved, mid- dle pruned, late reused, showing that adaptive compute allocation is both learnable and aligned with transformer computation phases. Overall, Dr.LLM demonstrates that explicit supervised routing reconciles efficiency, accuracy, and robustness without retraining, providing a practical step toward budget-aware reasoning and scalable adaptive inference.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
This work was supported by the NAVER corporation.
The authors also grateful to Comelius Emde for his careful proofreading.
REFERENCES
Sangmin Bae, Yujin Kim, Reza Bayat, Sungnyun Kim, Jiyoun Ha, Tal Schuster, Adam Fisch, Hrayr Harutyunyan, Ziwei Ji, Aaron Courville, et al. centure-of-recursiinn Learning dynamic recur- sive depths for adaptive token-level computation. arXiv preprint arXiv:2507.1022 2025.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical com- monsense natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try are, the ai2 reasoning challenge. arXiv preprint arXiv: 0803.0545, 2018.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv: 10. 14168, 2021.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Fukasz Kaiser. Universal transformers. arXiv preprint wrXiv:1807.03817 2018.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, al. The llama 3 herd of models. arXiv e-prints, pp. XrXiv-2407, 2024.
Mostafa Elhoushi, Akshat Shrivastava, Diana Liskovich, Basil Hosmer, Bram Wasti, Liangzhen Lai, Anas Mahmoud, Bilge Acun, Saurabh Agarwal, Ahmed Roman, al. Layerskip: Enabling early exit inference and self-speculative decoding. arXiv preprint arXiv:2404. 16710, 2024.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, (3110)):139, 2022.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Fos- ter, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, h, Haonan Li, Kyle McDonell, Niklas Muen- nighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. The language model evaluation harness, 07 2024. (tLhttps::/ zww.do.nooooooooooooo /t//tttsss//eecooooorg///////// wtt.:/dd..rgg/ecood//////////
Jonas Geiping, Sean McLeish, Neel Jain, John Kirchenbauer, Siddharth Singh, Brian Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, and Tom Goldstein. Scaling up test-time compute with latent reasoning: A recurrent depth approach. arXiv preprint arXiv:2502.05177 2025.
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason Lee, and Dimitris Papailiopoulos. Looped transformers programmable computers. In International Conference on Machine Learning, pp. 11398-11442. PMLR, 2023.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249-256. JMLR Workshop and Conference Proceedings, 2010.
Shwai He, Tao Ge, Guoheng Sun, Bowei Tian, Xiaoyang Wang, and Dong Yu. Router-tuning: A simple and effective approach for enabling dynamic-depth in transformers. arXiv preprint arXiv:2410.1 112400113888 3184, 2024.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.088 03300, 2020.
11

---

Published as a conference paper at ICLR 2026
Hamish Ivison, Yizhong Wang, Valentina Pyatkin, Nathan Lambert, Matthew Peters, Pradeep Dasigi, Joel Jang, David Wadden, Noah A Smith, Iz Beltagy, et al. Camels in changing cli- mate: Enhancing Im adaptation with tulu 2. arXiv preprint arXhv:2311.10702 2023.
Ziyue Li, Yang Li, and Tianyi Zhou. Skip layer or loop it? test-time depth adaptation of pretrained Ilms. arXiv preprint arXiv:2507.0799 2025.
Stephanie Lin, Jacob Hilton, and Owain Evans. Truthfulqa: Measuring how models mimic human falsehoods. arXiv preprint arXiv:2 2021.
Tsung- Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollár. Focal loss for dense object detection. Proceedings of the IEEE international conference on computer vision, pp- 2980-2988, 2017.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017.
Xuan Luo, Weizhi Wang, and Xifeng Yan. Adaptive layer-skipping in pre-trained llms. arXiv preprint arXiv:2503.23700 2025.
MAA.
American Invitational Mathematics URL.
invitational
mathematics
examination
aime.
In Febru-
American
Examination
AIME 2024.
ary
2024.
https:/ paaps://maa.ouu htt.a/aaa...aaaaaaaaaaaaaaaae htt/:mmahhhhmppae cnmmpttiiioss..
american- iotnn-iivitatiiiinnn www.raattoooa-mmtttnaaaaese mathematicsseeee weteemmtiics-exemnatiiiio. -xw.mmataiinnnnnnee aime
Xin Men, Mingyu Xu, Qingyu Zhang, Bingning Wang, Hongyu Lin, Yaojie Lu, Xianpei Han, and Weipeng Chen. Shortgpt: Layers in large language models are more redundant than you expect, 2024. URL https://arxiv. org/abs/2403.03355
2(3):4, 2024.
Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don't 't know: Unanswerable questions for squad. arXiv preprint arXiv: 03822, 2018.
David Raposo, Sam Ritter, Adam Santoro, Greg Wayne, Theophane Weber, Matt Botvinick, Aäron van den Oord, and Razvan Pascanu. Mixture-of-depths: Dynamically allocating compute in transformer- based language models. arXiv preprint arXiv:2404.02255 2024.
David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Di- rani, Julian Michael, and Samuel R Bowman. Gpqa: A graduate-level google-proof q&a bench- mark. In First Conference on Language Modeling, 2024.
Hassan Sajjad, Fahim Dalvi, Nadir Durrani, and Preslav Nakov. On the effect of dropping layers of pre-trained transformer models. Computer Speech & Language, 77:101429, 2023.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le. Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gater mixture-of-expprtt layer. arXiv preprint arXiv: 100.v:/700..0.....0 2017.
Surat Tecrapittayanon, Bradley McDanel, and Hsiang-Tsuun Kung. Branchynet: Fast inference via early exiting from deep neural networks. In 2016 23rd international conference on pattern recognition (ICPR), pp. 2464-244444 2469. IEEF, 2016.
Yuxuan Tong, Xiwen Zhang, Rui Wang, Ruidong Wu, and Junxian He. Dart-math: Difficulty-aware rejection tuning for mathematical mooblem-solving Advances in Neural Information Processing Systems, 37:7821-7844 2024.
Dennis Ulmer, Martin Gubri, Hwaran Lee, Sangdoo Yun, and Seong Oh. Calibrating large lan- guage models using their generations only. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 15440-15459, 2024.
Xin, Raphael Tang, Jaejun Lee, Yaoliang Yu, and Jimmy Lin. Deebert: Dynamic early exiting for accelerating bert inference. arXiv preprint arXiv:2004. 12993, 2020.
An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei. et al. Qwen2. technical report. arXiv e-prints, pp- arXiv-2412, 2024.
12

---

Published as a conference paper at ICLR 2026
Liu Yang, Kangwook Lee, Robert Nowak, and Dimitris Papailiopoulos. Looped transformers better at learning learning algorithms. arXiv preprint arXiv22311.2222, 2023.
Wanjun Zhong, Ruixiang Cui, Yiduo Guo, Yaobo Liang, Shuai Lu, Yanlin Wang, Amin Saied, Weizhu Chen, and Nan Duan. Agieval: A human-centric benchmark for evaluating foundation models. arXiv preprint arXiv:2304. 06364, 2023.
Wangchunshu Zhou, Canwen Xu, Tao Ge, Julian McAuley, Ke Xu, and Furu Wei. Bert loses pa- tience: Fast and robust inference with early exit. Advances Neural Information Processing Systems, 33:18330-18341 2020.
13

---

Published as a conference paper at ICLR 2026
APPENDIX
A
AUTHOR CONTRIBUTIONS
All authors contributed to writing and editing the paper.
Ahmed Heakl proposed the initial idea and motivation for the work, drafted the experimental set- tings, implemented and ran all experiments, collected the data, analyzed results, prepared visual- izations, reviewed related work, wrote the first draft, edited the paper, and published the research artefacts.
Martin Gubri provided daily supervision, helped refine the experimental settings, consolidated the paper's narrative, proposed the out-of-distribution evaluation (Section 6.2) and the "What do the routers learn?" experiment (Table 6), offered technical support, contributed to the initial draft, and reviewed the final draft.
Ahmed Heakl and Martin Gubri jointly created the diagrams (Figures 1-3).
Salman Khan provided feedback during the ideation phase and reviewed the draft.
Seong Joon Oh and Sangdoo Yun suggested Figures 1 and 2, validated the experimental settings, and contributed writing and editing the paper.
Martin Gubri, Seong Joon Oh, and Sangdoo Yun provided weekly supervision and supported the project through organizational and funding contributions.
B SCORING, REWARD, AND ANSWER CHECKING
Given input y and a candidate path म, we run generation with the model constrained to and obtain a textual response à. We then map à to a scalar reward R(ā, a):
ARC (multi-choice). Extract letter A-D via strict regex match (accepting optional "Answer:") The reward is for correct letter, otherwise.
DART (math). Extract the boxed expression (inject if needed for base models), then compute grade_answer(e a) [0, 1] using a symbolic equivalence checker and robust string comparator.
52.000
700
51.995
51.950
Acrurcyy
37.0
51.922
36.8
400
Accuracy
Computation
51.900
16.8
300
51.975
Baseline
51855
Total Accursty
100
51925
Average Layers
35.57
02
622
6:4 Routing Protebilty
3
0.8
(a) Accuracy average layers used under control interpolation.
Routing probability skiplexecutefrepeat.
distribution
acTOSS
Figure 7: Fine-grained control in LLaMA-8B. (a) Accuracy as a function of interpolated routing decisions, compared to baseline (red) and ours (green). (b) Histogram of routing probabilities. Shifts from execute - skip correlate with higher accuracy, while repeat allocations increase computation.
14

---

Published as a conference paper at ICLR 2026
Weighted
0.08
0.350
Skip Accura
0.345
0.335
0.02
0.325
1000
2000
3000
4000
1000
2000
3000
4000
0.005
Exerute Accuracy
0.980
Accuracy
0.975
0.003
0.970
0.002
0.965
0.001
0.960
1000
2000
3000
4000
1000
2000
3000
4000
Training Steps
Training Steps
Figure 8: Effect loss choice under class imbalance. Macro Fl across training for weighted CE, focal, and plain CE. While all losses perform similarly on the majority execute class, only focal loss improves skip accuracy and yields non-trivial repeat accuracy, highlighting its necessity for minority classes.
C
CONTROL OF ROUTER DECISIONS
Beyond analyzing learned routing policies, we study whether router decisions can continuously controlled to balance accuracy and efficiency. Figure reports results for LLaMA-8B.
We introduce a scalar control parameter E [-1, 1] that interpolates router probabilities with fixed skip, execute, repeat distributions:
(1 t) Takip + Trouter:
PE -055). t = p+1
0.5
Tpp)
(1 t) Trouter t Texec:
p E (-0.5, 0.5), t t= P+0.5
(1 t) Texse + Frepects
p E (0.5, t 2-0.5
0.5
Here Tronter are the learned router probabilities, and Takip: Texce: Tropeat one-hot distributions over the three actions.
This formulation allows to smoothly traverse the spectrum from aggressive skipping to repeated execution, without retraining the router. Figure 7a shows that modest interpolation (p == -0.5) re- duces average layers while slightly increasing accuracy, suggesting that routers tend to over-execute by default. The distributional shifts in Figure cooorobbacoro this: reallocating mass from execute toward sk i correlates with accuracy gains, while reallocating toward repeat primarily increases computation with diminishing benefit.
In sum, router behavior not only learnable but also tunable post-training, enabling fine-grained control over the accuracy-efficiency trade-off through single scalar knob.
D FOCAL VS. CROSS-ENTROPE UNDER CLASS IMBALANCE
Router supervision is highly imbalanced: Nakip = 4,399, Rexecute 120,956, Prepest 1,457. Plain cross-entropy minimizes error by predicting the dominant execute class, yielding trivial accuracy s skip/repeat. Weighted CE partly compensates, but still collapses on repeat. Focal loss (Lin et al. 2017) reweights classes and Nown-modulates easy majority examples, forcing learning
actions. As shown in Fig. 8 all losses perform similarly execute, but focal substantially improves skip accuracy and Is the only setup where non-trivial repeat accuracy is learned. Thus, focal loss is essential to mitigate imbalance and enable useful skip/repeat routing.
15

---

Published as a conference paper at ICLR 2026
windows
windows
windows
0.4
0.5
Skip Accuraty
0.3
Macro F1
0.4
0.3
250
500
750
1000
1250
1500
1750
2000
250
500
750
1000
1250
1500
1750
2000
1.0
Execute Accuracy
Repeat Accuracy
30.20
0.15
0.7
0.10
0.6
0.5
250
500
750
1000
1250
1500
1750
2000
250
500
750
1000
1250
1500
1750
2000
Training Steps
Training Steps
Figure 9: Effecc window size on router training. Larger pooling windows consistently improve minority-class accurac. Gains saturate beyond 16 windows, suggesting diminishing returns.
E TRAINING ON MORE WINDOWS
Windowed mean pooling stabilizes router decisions by aggregating hidden states over larger con- texts. Figure 9 shows that increasing the number of windows yields consistent improvements for minority actions. Skip accuracy rises from 0.32 (1 window) to 0.42 (32 windows), and repeat ac- curacy nearly doubles from 0.12 0.23. Execute accuracy stays unchanged at 0.98, confirming that the majority class is unaffected. Macro-F1 improves from 0.42 to 0.53, with most of the gain realized between 8 and 16 windows, indicating that more granular context summaries significantly help routers capture rare actions without harming the dominant class.
F MCTS TRAINING DATA ANALYSIS
Labels distribution. To better understand the supervision signal provided to the routers, we analyze the distribution of sk p/exeeute/reepa actions across datasets and model families (Fig. [0). Across all models and datasets, the vast majority of labels are execute, typically ex- ceeding 90%, confirming the extreme class imbalance (Nexecate Nakipp Rezepeat) and motivating focal loss with rebalancing during training (Sec. D). Skip ratios vary across datasets: ARC-Easy and ARC-Challenge exhibit noticeably higher skip counts than DART, suggesting that logical reasoning tasks permit redundancy while mathematical reasoning tasks require more thorough computation. Repeats
overall (1-3% of labels) but occur consistently across all datasets, with higher fre- quency in more challenging DART levels, indicating that repetition is a targeted mechanism for dif- ficult problems rather than a generic operation. Model family and scale also influence distributions: LLaMA- models exhibit more balanced skip/cxecute ratios compared to their instruction-tune tuned counterparts, which strongly favor execution, while instruction tuned variants slightly increase re- peat counts. Larger 8B models reduce skips further, reflecting greater reliance on their depth, though still allocating some repeats when beneficial. Overall, the MCTS-derived labels capture structured, interpretable routing signals under heavy imbalance, requiring routers to learn policies where most layers execute but the rare skip and repeat actions play a disproportionate role efficiency and accuracy.
Decisions per layer. Figure [ reveals structured routing patterns that align with transformer com- putation phases. Across all model families, early layers (embedding and low-level processing) are almost always executed, indicating their necessity for stable representations. Middle layers show the highest variation, with frequent skips reflecting redundancy in feature composition. Late lay- display higher repeat frequencies, particularly for the more difficult DART tasks, suggesting
16

---

Published as a conference paper at ICLR 2026
40000
skip
skip
35000
49000
repeat
repeat
Total Courts
30000
Total Counts
25000
20000
15000
10000
5000
0
ABG/C
ARC-E
DARF1
DARF33
DARF-3
DART-4
DANT-S
ARC
ARC
DWVT11
DAHF22
DART-3
DART4
DART-3
(a) LLaMA 3B
(b) LLaMA 8B
60000
skip
skip
execule
execute
Tatal Counts
50000
repeat
40000
40000
30000
30000
Total
20000
20000
10000
10000
ARC
ARC-E
DANF1
DANT-3
DAKT-3
DART-
DAKT-S
ARC-C
ARC-B
DART-1
DAKF-2
DANT-3
DART
DAHT-S
LLaMA-Base 3B
(d) LLaMA-Bass 8B
35000
25000
execute
execute
repeat
repeat
Total Counts
25000
lotal Counts
20000
15000
25000
10000
10000
5000
0
ABC-C
ARCE
DART1
DARF22
DART-3
4
DART-S
ARC
ARC.E
DART.1
DART.2
DART-3
DART4
DART-5
Qwen 3B
Qwen 7B
Figure 10: Label distribution across models. Distribution of skip/execute/repeat actions across datasets for different planners: (a) LLaMA-3B, (b) LLaMA-8B, (c) LaLaMA--as---- (d) LLaMA- Base-8B, (e) Qwen-3B, (f) Qwen-7B.
that deeper refinement is allocated where multi-step reasoning is required. Instruction tuned mod- els exhibit more aggressive skipping than base models, supporting the view that fine-tuning creates functionally specialized layers that routers can prune more confidently. These trends confirm that Dr.LLM learns consistent, interpretable depth allocation policies across both model scale and family.
17

---

Published as a conference paper at ICLR 2026
(a) -lamma3.2-3B-Insstete
(b) Llama-3.1
24
28
20
24
20
Layer
16
12
1.100
(c) Llama-3 3.----3.2---Ba--
(d) Mlama-3.1-B--D..
1.075
24
28
1.050
20
24
1.025
20
Layer
16
1.000
12
0.975
Mean layer
0.950
Owen2.5-3B-Instten
Qwen2.5-7B-Instmue
0.925
32
24
0.900
28
20
20
16
16
12
12
ARC-E
ARC-C
DART-I
DART-2
DART
DART4
DARTS
arc-e
ARC-C
dart-1
DART-2
DART
DART
dART:
Figure 11: Per-layer routing frequency across datasets and models. Heatmaps show the mean usage per layer skip, = execute, 2 = repeat) for six backbones: (a) LLLMA--322-BB-Inssmn
(e) Qwen2.5-3B- Instruct, and (f) Siven2.5-BB-Insteeee The x-axis corresponds to benchmark subsets (ARC-E, ARC- C, DARTI-S Early layers consistently executed. middle layers are frequently skipped, and late layers are occasionally repeated, especially on more complex DART levels.
(b)
(c) LLaMA-3.. --.--------B------ 3.2--B-Base, (d) -L-MA-3.1-8B-BBa..
18


Reviews:
[{"review_id": "HU0jdYM1HJ", "reviewer": "Reviewer_p46J", "summary": "The paper proposes a method for dynamically skipping, executing or repeating layers in transformer architecture per example in order to improve the model quality while preventing wasted compute on simpler queries. It can be applied to any LLM that is already trained without altering the weights, only adding per-layer routers. The routers are trained using an offline tree-based search method and use hidden layers of the previous layers as inputs. The search algorithm is monte carlo based and includes an explicit length penalty to incentivize shorter paths to favor executing less number of layers. The method is evaluated on Llama and Qwen models and has shown to improve accuracy around 3% while saving 5 layers per example on average.", "questions": "- Figure 1 shows the network has 600+ layers. -> why there are so many layers? No state of the art LLM has that many layers.\n\n- The router accuracy during training is 61%, which seems low. What’s the implication of this, have you tried improving the router accuracy?\n\n- “all four models gain 0.40%p accuracy on GSM8k” -> This claim on page 7 does not match the results on Table 4. Llama 3B accuracy drops on GSM8k according to table 4. \n\n- Why are there no baseline llama numbers in Table 5? It’d be good to add the baseline in this table as some of the tasks in this table were not shown before in the paper. The table shows Dr.LLM is better than the state of the art methods however it does not show how it’s better than the pure LLM in HumanEval, Hellaswag, MMLU tasks.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The method is applicable to already trained LLMs with only minimal router training\n- It is a comprehensive method which combines layer skipping and looping which was previously explored separately in most former work\n- Research artifacts including the code and data is shared publicly", "weaknesses": "- There is no latency measurements shown in the paper.  While the method reduces 5 layers per example on average, the routers add additional overhead per layer and overall latency might increase because of that.\n- The method is not applicable to batching as the routing is done per example level, when the batch size is greater than 1, accuracy might suffer. \n- Out of domain generalization looks weak. While the method improves accuracy for in-domain tasks (the tasks the routers are trained for), for out of domain tasks accuracy drops.\n- I have several questions about the evaluation, listed in questions section below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5TeBs4VaZ8", "reviewer": "Reviewer_wfaJ", "summary": "This paper proposes Dr.LLM, a retrofittable framework that introduces per-layer dynamic routers for frozen pretrained LLMs. Each router decides whether to skip, execute, or repeat a transformer block, allowing adaptive compute allocation per input. The routers are trained with explicit supervision derived from offline Monte Carlo Tree Search (MCTS), which identifies optimal layer configurations under compute constraints. The approach achieves accuracy improvements on reasoning-heavy tasks (ARC, DART) while reducing the number of executed layers and generalizes well to out-of-domain benchmarks.", "questions": "If all concerns shown in weaknesses are resolved, I would raise my score.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Novel supervised routing framework that avoids inference-time search and large-scale retraining, yet improves both efficiency and accuracy.\n2. Strong empirical results showing consistent gains on logic and math benchmarks, plus solid cross-domain generalization with minimal accuracy drop.", "weaknesses": "1. Lack of wall-clock evaluation – The reported “layers saved per query” metric does not directly reflect real-world speedup. Since each layer involves router computation overhead, the paper should include wall-time results to substantiate efficiency claims.\n2. Limited training exploration – The routers are trained separately from the LLM. Joint training (LLM + routers) could yield better synergy and higher final performance, especially since inference is a long-term process while training is a one-time cost.\n3. Unclear batch inference feasibility – It’s uncertain whether Dr.LLM supports efficient batch processing, as per-token dynamic routing can lead to divergent compute paths and reduced parallelism. Discussion or experiments on batch inference scalability would strengthen the work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PjPg2uFqlW", "reviewer": "Reviewer_nSUU", "summary": "This paper introduces Dr.LLM, a new method to make Large Language Models (LLMs) more efficient and sometimes more accurate. It works by adding small \"routers\" to an existing LLM that dynamically decide whether to use, skip, or repeat each layer for a given problem.", "questions": "Same as above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The framework is retrofittable, meaning it can be applied to existing, pre-trained models without costly full-scale retraining. Keeping the base LLM weights frozen makes the approach highly practical and accessible.\n\n2. Using MCTS to generate an \"oracle\" dataset of optimal execution paths is a very clever way to supervise the routers. This avoids complex reinforcement learning and eliminates the need for a costly search process during actual inference.\n\n3. The claim that the trained routers generalize well to out-of-domain tasks suggests they are learning a robust, transferable policy for allocating computational resources rather than just overfitting to the training data.", "weaknesses": "1. The paper emphasizes that router training is lightweight, but the initial offline MCTS process to find the optimal paths could be extremely computationally expensive. The true cost of preparing the training data is not fully addressed.\n\n2. The search space for MCTS (all possible combinations of skipping/repeating layers) grows exponentially with model depth. It's unclear how well this search process scales to extremely deep models (e.g., 100+ layers).\n\n3. The actions (skip, execute, repeat) are discrete and simple. This might not be nuanced enough for problems where only a part of a layer (e.g., a few attention heads) is needed, potentially leaving further efficiency gains on the table.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FDuszW8a4H", "reviewer": "Reviewer_kHKr", "summary": "This paper introduces Dr.LLM, a dynamic layer routing method for large language models (LLMs) that operates at the sequence level. The motivation behind this approach arises from the limitations of traditional static-depth LLMs, which apply a fixed number of layers regardless of input complexity. Such models tend to waste computation on simple prompts and lack the flexibility needed to handle more challenging reasoning tasks. Dr.LLM addresses this by introducing a lightweight routing module at each transformer layer. Based on the sequence of embeddings from the previous layer, the router dynamically decides whether to skip the current layer, execute it once, or execute it twice. These routing decisions are supervised using a Monte Carlo Tree Search (MCTS)-based algorithm. The paper presents experimental results on both in-domain tasks, where the input distribution matches the data used to train the router, and out-of-domain tasks, which involve distribution shifts. Dr.LLM is also evaluated against prior routing methods.", "questions": "* A more precise definition of the per-example number-of-layers-used metric would be helpful. I assume this refers to summing the number of active layers across each forward pass during generation, but providing a clear and explicit definition would improve clarity.\n* “We allow skips of at most two consecutive layers…” What motivates the need for this specific limitation?\n* What does the \"Original\" column in Table 2 represent?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The paper is clearly written, and the authors present their method in a precise and comprehensible way.\n* The use of MCTS to generate routing paths is unique and answers the requirements of a routing method presented in the introduction section.\n* The method outperforms other routing strategies.", "weaknesses": "* The computational cost associated with generating the MCTS-based supervision is not sufficiently reported. While this supervision is computed offline, its upfront cost does not appear to be negligible. It would be valuable to include a comparison of this cost relative to the resources required to train the router itself. Additionally, since performance gains have been demonstrated primarily on in-domain data, expanding the router’s training data to cover a broader range of tasks seems essential for making Dr.LLM applicable to general-purpose LLM scenarios and therefore understanding the computational effort required to generate such supervision is important.\n* KV cache compatibility — The authors state at the beginning of Section 3.1 that the method is KV cache compatible, but it's unclear how this is achieved. During generation, the router can select a different path at each forward pass, potentially resulting in different inputs to each transformer block. This variability seems to conflict with standard KV cache usage, which assumes consistent execution paths. To clarify, consider the following example: suppose we are generating a sequence and have reached the point of generating the 100th token. At this step, we pass $H^{\\ell-1}$ through the router and obtain $y_{\\ell}=\\text{skip}$. However, in the next forward pass (i.e., when generating the 101st token), the router outputs $y_{\\ell} = \\text{exec}$. In this case, how is the KV cache for layer $\\ell$ defined to correctly support both scenarios? The mechanism for maintaining cache consistency under dynamic routing decisions needs further clarification.\n* The efficiency analysis in the experimental section appears incomplete. Is the overhead introduced by the routing modules truly negligible compared to the reported efficiency gains (0.9–2.5%) in Table 3?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)