Dr.LLM：大语言模型中的动态层路由

论文(Dr.LLM: Dynamic Layer Routing in LLMs)介绍了一种巧妙的技术，在提高准确率的同时降低大语言模型的计算量。路由器在提高准确率的同时，每次查询可减少约3到11层的计算。论文的方法是可改装的，为冻结的大语言模型添加逐层路由器，决定跳过、执行或重复每个模块。路径通过对层编辑进行简短的蒙特卡洛树搜索进行离线监督，然后在线执行时无需搜索。在平均节省层数的同时，提高了逻辑和数学任务的准确率。

啥是Dr.LLM： 附加在每一层的微型MLP路由器读取窗口化的均值池化隐藏状态，并输出三种操作之一：跳过、执行一次或重复一次。基础权重保持冻结，并且与KV缓存保持兼容。

监督机制如何工作： 长度感知的MCTS探索在计算预算约束下跳过或重复层的编辑前向传播，只保留那些保持或改善金标准答案奖励的路径。然后使用焦点损失和类别重新平衡在大约4000条发现的路径上训练路由器。

主要结果： 在六个主干网络上的ARC和DART任务中，路由器在提高准确率的同时，每次查询减少约3到11层。 例如：LLaMA-3B-Base在DART上从11.8%提升到15.8%，平均节省4.1层。经过指令调优的8B模型在DART上也有提升，同时节省11层。 域外泛化能力也很强。

这个方法为什么有效： 早期层保持稳定，许多中间层被跳过，后期层有时会被重复，尤其是在更难的数学问题上，这将深度重新分配到迭代优化更有价值的地方。

冻结解码器LLM，为每个模块附加一个2层瓶颈MLP路由器。对隐藏状态使用窗口化均值池化。运行简短的离线MCTS对层编辑进行搜索，使用长度惩罚来收集改进或保持的路径。** 使用焦点损失训练路由器以处理严重的类别不平衡。在推理时，执行贪婪路由，无需搜索。可选择暴露单个标量来偏向更多跳过或更多重复，以实现更严格的成本控制。

本文转载自​​AI帝国​​，作者：无影寺