让大模型推理更省力:通过记忆路由历史提升效率
Do Dynamic Routers Need Memory? HeRo: History-Aware Routing for Efficient LLM Inference

- 引入路由记忆机制,结合历史决策与当前状态做路由选择
- 在Llama 3.1-8B上跳过26.87%参数仍保持100.24%性能
- 适合需要高效推理的部署场景,尤其擅长多步推理任务
动态层路由通过为不同输入跳过冗余层来降低大语言模型的推理成本。现有方法将每层路由决策视为仅依赖当前隐藏状态的局部操作,忽略了路由在深度上的序列性与路径依赖性:早期决策影响后续路由器看到的表示,且所有决策共同优化层使用目标。本文提出历史感知路由(HeRo),引入路由器记忆机制,在模型深度上显式维护路由状态。该记忆通过线性注意力构建,逐步聚合前序路由得分及其引发的残差更新,形成紧凑的历史表示。在每一路由层,路由器同时基于累积状态和当前隐藏表示选择执行分支。该方法应用于逐标记前馈网络路由,仅需训练轻量级路由器与适配器,无需修改预训练参数。在Llama 3.1-8B、Llama 2-7B和Llama 2-13B上,HeRo在十种基线中表现最优。在Llama 3.1-8B上,其跳过26.87%参数时仍达密集模型100.24%性能;在更紧计算预算下跳过38.82%参数,性能保留97.01%。消融实验表明,移除路由历史会显著降低性能,尤其在多步推理与代码生成任务中,验证了显式路由记忆能实现更准确、自适应的动态路由。
原文摘要 · Abstract (English)
Dynamic layer routing reduces the inference cost of Large Language Models (LLMs) by learning to skip layers for individual tokens. Existing methods, however, treat each routing decision as a local operation conditioned solely on the current hidden state which is a formulation that overlooks the sequential, path-dependent nature of routing across depth: earlier decisions shape the representations seen by downstream routers, and the layer-usage objective couples all decisions jointly. We propose History-Aware Routing (HeRo), a dynamic routing framework that resolves this mismatch by introducing a router memory mechanism to maintain an explicit routing state across model depth. The memory is constructed via linear attention, incrementally aggregating preceding routing scores and their induced residual updates into a compact history representation. At each routed layer, the router conditions jointly on this accumulated state and the current hidden representation to select the executed branch. Instantiated for token-wise FFN routing, HeRo trains only lightweight routers and adapters on a frozen backbone, requiring no modification to pretrained parameters. Across Llama 3.1-8B, Llama 2-7B, and Llama 2-13B, HeRo consistently achieves the highest aggregate performance retention among ten baselines. On Llama 3.1-8B, it bypasses 26.87% of model parameters while achieving 100.24% of dense model performance across seven benchmarks, and retains 97.01% while bypassing 38.82% of model parameters under a tighter computation budget. Ablation studies confirm that removing routing history consistently degrades performance, most notably on multistep reasoning and code generation, validating that explicit routing memory enables more accurate and adaptive dynamic routing than solely conditioning on hidden state.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。