Dr.LLM: Dynamic Layer Routing in LLMs

📄 arXiv: 2510.12773v1 📥 PDF

作者: Ahmed Heakl, Martin Gubri, Salman Khan, Sangdoo Yun, Seong Joon Oh

分类: cs.CL, cs.AI, cs.LG

发布日期: 2025-10-14

备注: 17 pages, Under submission

💡 一句话要点

Dr.LLM：通过动态层路由提升大语言模型推理效率与精度

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 动态路由 自适应深度 推理加速 显式监督

📋 核心要点

  1. 现有LLM对所有token进行全层处理，导致简单查询计算浪费，复杂查询灵活性不足。
  2. Dr.LLM通过轻量级路由器动态决定每层是否执行，以提升效率和精度。
  3. 实验表明，Dr.LLM在提升精度的同时节省计算，并能泛化到领域外任务。

📝 摘要（中文）

本文提出Dr.LLM，一种用于大语言模型（LLMs）的动态层路由框架。该框架通过轻量级的逐层路由器，决定跳过、执行或重复Transformer堆栈中的每个块，从而避免了对所有token都进行全层处理的低效性，并为需要深度推理的复杂查询提供更大的灵活性。路由器通过显式监督进行训练：使用蒙特卡洛树搜索（MCTS）导出高质量的层配置，在计算预算下保持或提高精度。该设计采用窗口池化以实现稳定的路由，使用焦点损失与类别平衡，以及瓶颈MLP路由器，确保在类别不平衡和长序列下的鲁棒性。在ARC（逻辑）和DART（数学）数据集上，Dr.LLM在平均节省5层计算的情况下，将准确率提高了高达+3.4%p。路由器推广到领域外任务（MMLU、GSM8k、AIME等）时，仅损失0.85%的准确率，同时保持效率，并且优于先前的路由方法高达+7.7%p。Dr.LLM表明，显式监督的路由器可以对冻结的LLM进行改造，以实现预算感知、精度驱动的推理，而无需更改基础权重。

🔬 方法详解

问题定义：现有的大语言模型在推理时，无论输入的复杂程度如何，都会对每个token都进行Transformer堆栈的所有层的计算。这导致对于简单的输入，存在大量的冗余计算；而对于复杂的需要深度推理的输入，模型的层数可能不足以充分表达。现有的自适应深度方法通常需要昂贵的推理时搜索、架构更改或大规模的重新训练，并且在提高效率的同时，往往会降低精度。

核心思路：Dr.LLM的核心思路是为预训练的LLM配备轻量级的逐层路由器，这些路由器可以动态地决定每个Transformer块是应该被跳过、执行还是重复。通过这种方式，模型可以根据输入token的复杂程度，自适应地调整计算深度，从而在保证精度的前提下，提高推理效率。这种设计避免了对模型架构进行大规模修改，也无需重新训练整个模型。

技术框架：Dr.LLM的整体框架包括以下几个主要模块：1) 预训练的LLM：作为基础模型，提供token embedding和Transformer层。2) 逐层路由器：每个Transformer层都配备一个路由器，用于决定该层是否执行。3) 训练策略：使用蒙特卡洛树搜索（MCTS）生成高质量的层配置，并使用显式监督的方式训练路由器。4) 推理过程：根据路由器的决策，动态地调整模型的计算深度。

关键创新：Dr.LLM的关键创新在于其显式监督的路由器训练方法。与以往依赖隐式学习或强化学习的方法不同，Dr.LLM使用MCTS生成高质量的层配置作为训练目标，从而能够更有效地训练路由器。此外，Dr.LLM的设计无需修改基础LLM的权重，使其能够轻松地应用于各种预训练模型。

关键设计：Dr.LLM的关键设计包括：1) 窗口池化：用于稳定路由决策，避免相邻层之间的路由决策过于频繁地切换。2) 焦点损失与类别平衡：用于解决类别不平衡问题，确保路由器能够有效地处理各种类型的输入。3) 瓶颈MLP路由器：使用具有瓶颈结构的MLP作为路由器，以减少参数量和计算开销。路由器接收来自当前层输入和前一层的路由决策作为输入，输出一个概率分布，指示该层应该被跳过、执行或重复。

🖼️ 关键图片

img_0

📊 实验亮点

Dr.LLM在ARC（逻辑）和DART（数学）数据集上，在平均节省5层计算的情况下，将准确率提高了高达+3.4%p。在领域外任务（MMLU、GSM8k、AIME等）上，仅损失0.85%的准确率，同时保持效率，并且优于先前的路由方法高达+7.7%p。这些结果表明，Dr.LLM能够在提高效率的同时，保持甚至提高模型的精度，并且具有良好的泛化能力。

🎯 应用场景

Dr.LLM具有广泛的应用前景，可以应用于各种需要高效推理的大语言模型应用场景，例如智能客服、机器翻译、文本摘要、代码生成等。通过动态调整计算深度，Dr.LLM可以在资源受限的设备上运行大型语言模型，并提高推理速度，从而降低部署成本。此外，Dr.LLM还可以用于提高模型的鲁棒性和泛化能力，使其能够更好地适应各种不同的任务和领域。

📄 摘要（原文）

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.