AI论文精选
⏶30

Dr.LLM：LLM中的动态层路由

发表
由 Ahmed HeaklAhmed Heakl 提交
作者: Ahmed HeaklAhmed Heakl, Martin GubriMartin Gubri, Salman Khan, Sangdoo Yun, Seong Joon Oh

摘要

大型语言模型 (LLM) 会通过 transformer 堆栈的所有层来处理每个 token，导致对简单查询的计算浪费，而对于需要更深层推理的更复杂查询则灵活性不足。 自适应深度方法可以提高效率，但之前的方法依赖于昂贵的推理时间搜索、架构修改或大规模重新训练，并且实际上尽管提高了效率，但通常会降低准确性。 我们引入 Dr.LLM，即 LLM 的动态层路由，这是一个可复用的框架，为预训练模型配备轻量级的每层路由器，用于决定跳过、执行或重复一个块。 路由器经过显式监督训练：使用蒙特卡洛树搜索 (MCTS)，我们在计算预算下推导出高质量的层配置，以保持或提高准确性。 我们的设计，包括用于稳定路由的窗口池化、具有类别平衡的焦点损失以及瓶颈 MLP 路由器，确保了在类别不平衡和长序列下的鲁棒性。 在 ARC (逻辑) 和 DART (数学) 上，Dr.LLM 的准确性提高了高达 +3.4%p，同时平均每例节省 5 层。 路由器在泛化到非领域外任务 (MMLU, GSM8k, AIME, TruthfulQA, SQuADv2, GPQA, PIQA, AGIEval) 时，准确性仅下降 0.85%，同时保持效率，并且比之前的路由方法高出高达 +7.7%p。 总的来说，Dr.LLM 表明，显式监督路由器可以对冻结的 LLM 进行复用，以实现预算感知、以准确性为驱动的推理，而无需改变基础权重。
查看 arXiv 页面查看 PDF
Dr.LLM：LLM中的动态层路由
Dr.LLM：LLM中的动态层路由
Dr.LLM：LLM中的动态层路由

评论

Ahmed HeaklAhmed Heakl
论文作者
论文提交者

大型语言模型（LLM）会通过 Transformer 堆栈的所有层处理每个 token，对于简单的查询会浪费计算资源，而对于需要更深层推理的复杂查询则缺乏灵活性。

Dr.LLM（LLM 的动态层路由）是一个可追溯框架，它为预训练模型添加了轻量级的每层路由器。
每个路由器决定是否跳过、执行或重复一个层，从而在无需重新训练或架构更改的情况下实现自适应深度。

路由器通过蒙特卡洛树搜索（MCTS）的显式监督进行训练，生成高质量的层配置，在计算预算内保持或提高准确性。
Dr.LLM 通过窗口池化、焦点损失和瓶颈 MLP 进行稳定，在类别不平衡和长序列下保持鲁棒性。

📈 结果
- 在 ARC（逻辑）和 DART（数学）上，Dr.LLM 将准确性提高了 +3.4%，同时为每个输入节省了约 5 层。
- 仅使用 0.85% 的准确率下降，路由器就能够泛化到 MMLU、GSM8k、AIME、TruthfulQA、SQuADv2、GPQA、PIQA 和 AGIEval。
- 与之前的路由方法（LayerSkip、FlexiDepth、MindSkip）相比，性能提高了 +7.7%。

> 💡 Dr.LLM 为冻结的 LLM 提供了预算感知、精度驱动的推理——无需修改基础权重。