TriRoute：用于联合自适应注意力、专家和 KV 缓存分配的统一学习路由

arXiv:2607.06601v1 公告类型：新 摘要：条件计算可以将语言模型质量与每个令牌的推理成本解耦，但领先的技术单独作用于单轴：专家混合 (MoE) 稀疏 FFN，深度混合 (MoD) 跳过整个转换器块，KV 缓存量化压缩注意力内存。我们认为这三个决策（注意力分辨率、专家选择和缓存位宽）是强耦合的，应该联合做出：一个足够稀有、足以保证充分关注的令牌可能也需要高精度缓存，无论哪个专家处理它。我们引入了 TriRoute，一个跨所有三个轴共享的单个轻量级控制器，它针对每一层的每个 token 发出协调的策略：（i）注意力模式（跳过/本地/完整），（ii）一组稀疏的 FFN 专家（具有恢复 MoD 的空专家），以及（iii）KV 缓存位宽。控制器在拉格朗日预算约束下通过异构松弛（Gumbel-Softmax 进行分类决策的直通估计和专家的负载平衡 top-k 门控）进行端到端训练，将平均计算和内存成本转变为可控旋钮。我们在朴素联合训练中识别出跨轴路由崩溃级联，其​​中一个轴上的崩溃传播到其他轴，并通过每轴归一化和耦合感知平衡损失来解决它。在计算最佳令牌计数下从 160M 到 1.3B 参数的纯解码器模型上，TriRoute Pareto 在匹配的推理 FLOP 和内存上主导了最佳独立 MoD+MoE+KV 量化组合，同时更好地保留了纯粹困惑度优化所侵蚀的稀有实体、代码和算术的尾部情况稳健性。事后分析揭示了可解释的结构：控制器将充分的注意力和高精度缓存分配给句子开头的位置、罕见的子词和命名实体，同时廉价地路由功能词。