Dr.LLM: Dynamic Layer Routing in LLMs

TL;DR

Dr.LLM tackles the inefficiency of fixed-depth LLMs by introducing dynamic, per-layer routing in a frozen model with $L$ transformer blocks. Routers decide to skip, execute, or repeat a layer and are trained with explicit supervision derived from offline length-aware Monte Carlo Tree Search (MCTS) paths, eliminating inference-time search. The approach uses windowed mean pooling, a compact two-layer router, and focal loss to handle class imbalance, achieving compute-efficient inference with no changes to base weights. Empirically, Dr.LLM improves accuracy on ARC and DART, saves 3–11 layers per example, and generalizes to out-of-domain benchmarks with only about a 0.85 percentage-point drop, outperforming prior routing methods by up to 7.7 percentage points.

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.

Figures (10)

• Figure 1: Dr.LLM improves accuracy while reducing computation. Number of layers used per example vs. accuracy on ARC and DART, averaged on six models.
• Figure 2: Our layer routing based on hidden states. Dr.LLM augments a frozen decoder-only LLM with per-layer routers that decide to skip, execute, or repeat a block once. Routers read windowed summaries of hidden states and are trained from MCTS-derived targets (Sec. \ref{['sec:training-data']}). For clarity, the diagram also highlights the router internals and the flow of hidden states across layers.
• Figure 3: Length-aware MCTS used to collect the supervised training dataset of per-layer routing configurations (skip/execute/repeat). For each input, MCTS explores modified layer paths and retains accuracy-preserving or improving ones under a compute budget.
• Figure 4: Analysis of routing decisions per layer, dataset, and model.(a) Layer frequency of LLaMa 3B and 8B base (B) and instruct (I) 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-Challenge (A-2), and DART levels 1–5 (from D-1 to D-5).
• Figure 5: Ablation study. We apply Dr.LLM on LLaMa3.2-3B and control: (a) the effect of bottleneck dimension, (b) the effect of number of linear layers, and (c) the effect of number of windows.