新上线今天0 投票
TriRoute：统一学习路由，联合自适应分配注意力、专家和KV缓存
大型语言模型（LLM）的推理成本一直是部署中的核心挑战。业界已探索多种条件计算技术来解耦模型质量与每token推理成本，但现有方法大多各自为政：混合专家（MoE）稀疏化前馈网络，混合深度（MoD）跳过整个Transformer块，KV缓存量化压缩注意力内存。然而，这些决策实际上高度耦合——一个需要全注意力的罕见token，很可能也需要高精度缓存，无论由哪个专家处理。
近期一篇论文提出 TriRoute，一个轻量级统一控制器，为每个token的每一层联合输出三项决策：注意力模式（跳过/局部/全局）、稀疏专家集合（含空专家以恢复MoD功能）以及KV缓存位宽。该控制器通过异构松弛技术（Gumbel-Softmax结合直通估计处理分类决策，负载均衡top-k门控处理专家选择）进行端到端训练，并引入拉格朗日预算约束，将平均计算和内存成本转化为可控旋钮。
研究团队发现，朴素联合训练会导致跨轴路由崩溃级联——某一轴的崩塌会传播至其他轴。为此，他们提出 每轴归一化 和 耦合感知平衡损失 加以解决。在160M至1.3B参数的解码器模型上，TriRoute在匹配推理FLOPs和内存的条件下，帕累托优于 独立MoD+MoE+KV量化组合的最佳效果，同时在罕见实体、代码和算术等尾部案例上展现出更强的鲁棒性。
后验分析揭示了可解释的结构：控制器倾向于为句子起始位置、罕见子词和命名实体分配全注意力与高精度缓存，而对功能词则采用廉价路由。这一工作为LLM推理优化提供了新视角：联合决策比孤立优化更有效，有望推动更高效、更智能的模型部署方案。