论文

动态路由需要历史吗?HeRo的历史感知高效LLM推理

Do Dynamic Routers Need Memory? HeRo: History-Aware Routing for Efficient LLM Inference

模型推理推理加速

摘要

动态层路由通过学习跳过各个标记的层来降低大型语言模型 (LLM) 的推理成本。然而,现有方法将每个路由决策视为仅以当前隐藏状态为条件的本地操作,这种表述忽略了跨深度路由的顺序、路径依赖性质:早期决策塑造了下游路由器所看到的表示,而层使用目标将所有决策联合起来。我们提出了历史感知路由(HeRo),这是一种动态路由框架,通过引入路由器内存机制来维持跨模型深度的显式路由状态,从而解决了这种不匹配问题。记忆是通过线性注意力构建的,增量地将先前的路由分数及其引起的残差更新聚合成紧凑的历史表示。在每个路由层,路由器联合根据此累积状态和当前隐藏表示来选择执行的分支。 HeRo 针对Token智能 FFN 路由进行实例化,仅在冻结主干上训练轻量级路由器和适配器,无需修改预训练参数。在 Llama 3.1-8B、Llama 2-7B 和 Llama 2-13B 中,HeRo 在十个基准中始终保持最高的总体性能保留率。在 Llama 3.1-8B 上,它绕过了 26.87% 的模型参数,同时在七个基准中实现了 100.24% 的密集模型性能,并在更严格的计算预算下保留了 97.01%,同时绕过了 38.82% 的模型参数。消融研究证实,删除路由历史会持续降低性能,尤其是在多步推理和代码生成方面,这验证了显式路由内存比单独调节隐藏状态能够实现更准确和自适应的动态路由。