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

# 面向联合自适应注意力、专家和KV缓存分配的统一学习路由
来源：https://arxiv.org/html/2607.06601
###### 摘要

条件计算有望将语言模型的质量与其每个词元的推理成本解耦，但主流技术各自作用于*单一*轴，且孤立发展：混合专家（MoE）稀疏化前馈网络，混合深度（MoD）跳过整个Transformer块，而KV缓存量化压缩注意力内存。我们认为这三个决策——*注意力分辨率该多高*、*选择哪些专家*以及*令牌应获得多少位缓存*——是紧密耦合的，应该联合做出。我们提出TriRoute，一个轻量级的统一控制器，对每一层的每个词元，它在三个轴上协同输出一个协调策略：(i) 注意力模式（跳过/局部/全局），(ii) 一组稀疏的FFN专家（包括一个可恢复MoD的空专家），以及 (iii) KV缓存位宽，它决定了该词元如何被后续查询记住。该控制器通过语言建模目标端到端训练，采用异构松弛方案——对分类的注意力和比特决策使用Gumbel-Softmax结合直通估计，对专家使用负载均衡的top-k门控——并在单个拉格朗日预算约束下进行，该约束将平均计算和内存成本变为一个可控制的旋钮。我们识别并缓解了影响朴素联合训练的*跨轴路由坍缩级联*问题，采用逐轴归一化和耦合感知平衡损失。在从160M到1.3B参数的仅解码器模型上，以计算最优的令牌数量训练，TriRoute绘制出的帕累托前沿在匹配的推理FLOPs和内存下，优于独立调优的MoD + MoE + KV量化组合的最佳性能，同时更好地保留了尾部案例的鲁棒性（稀有实体、代码、算术），这些是纯困惑度优化所削弱的。事后分析表明，控制器学习了可解释的结构：它将对句首词元、稀有子词和命名实体投入注意力和缓存位，同时对功能词和可预测的延续进行廉价路由。我们的结果表明，一个共享的学习控制器——而非三个手动调优的机制——是将固定推理预算花在最需要的地方的原则性方法。

## 1 引言

部署大语言模型（LLM）的主要成本发生在推理阶段，每次一个词元，而且成本惊人地*均匀*：一个密集Transformer在句法上无足轻重的功能词上花费的浮点运算和分配的键/值（KV）缓存内存，与一个锚定段落含义的稀有命名实体完全相同。这种均匀性在统计上是浪费的。一系列关于*条件计算*的研究试图打破这种局面，目前有三种技术主导实践，每种作用于Transformer块的不同部分：

- •**混合专家（MoE）** 稀疏化前馈网络（FFN），将每个词元路由到一小部分专家，从而将参数数量与活跃FLOPs解耦 [42 (https://arxiv.org/html/2607.06601#bib.bib2), 31 (https://arxiv.org/html/2607.06601#bib.bib4), 20 (https://arxiv.org/html/2607.06601#bib.bib3)]。
- •**混合深度（MoD）** 使*深度*适应不同词元，学习一个逐块门控，让词元绕过整个注意力+FFN子层 [38 (https://arxiv.org/html/2607.06601#bib.bib6)]，呼应了早期的自适应深度和提前退出工作 [24 (https://arxiv.org/html/2607.06601#bib.bib22), 15 (https://arxiv.org/html/2607.06601#bib.bib23), 18 (https://arxiv.org/html/2607.06601#bib.bib24), 41 (https://arxiv.org/html/2607.06601#bib.bib25)]。
- •**KV缓存量化** 压缩主导长上下文服务的注意力内存，以2–4位存储键和值，并精心选择分组尺度 [33 (https://arxiv.org/html/2607.06601#bib.bib14), 26 (https://arxiv.org/html/2607.06601#bib.bib15)]。

这些机制几乎总是被*独立地*研究和调优。然而，它们所做的决策并非独立。考虑一个稀有实体词元，比如*“...signed by Nakamura on Tuesday”*中的姓氏。MoD可能正确判断其FFN变换是可预测的，可以跳过；然而，正因为该词元稀有且信息丰富，它很可能需要*完整*的注意力分辨率来与其共指对象绑定，并且其键/值应该以*高*精度存储，以便后续查询能忠实地检索到它。像*“the”*这样的功能词在所有三个方面则正好相反。因此，“正确”的计算量并非每个词元的单一标量（如MoD隐含假设的那样），而是一个跨越异构资源的*向量*耦合选择，并且这些选择之间的相关性既不是恒定的，也不是先验明显的。

#### 本文。

我们提议用*一个单一的学习控制器*来做出这三个决策，并与模型联合训练。我们将由此产生的架构称为TriRoute，因为一个共享的路由器引导三条路径穿过Transformer块。对于每一层ℓ的每个词元t，控制器输出一个协调策略：

1. 1.一个**注意力模式** a ∈ {skip, local-w, full}，控制词元关注序列的范围；
2. 2.一个**稀疏专家选择** e，在E个FFN专家中选择，并指定一个**空专家**，作为MoD风格FFN跳过的特例；
3. 3.一个**KV位宽** b ∈ {2, 4, 8, 16}，决定了词元自身的键/值写入缓存的精度，从而决定了*未来*词元能多忠实地关注它。

隔离机制：MoE (FFN稀疏性) | MoD (深度跳过) | KV量化 (缓存位) —— β₁, β₂, β₃ (三个预算，分别调优)
TriRoute (我们的)：控制器 g_φ | 注意力 (跳过/局部/全局) | 专家 (top-k/空) | KV位 (2/4/8/16) | C⋆ (一个全局预算)

统一：词元 → "the" "signed" "by" "Nakamura" "on" "Tuesday" ■attn ■expert ■bits

图1：从三个隔离机制到一个控制器。MoE、MoD和KV量化各自作用于单一轴，具有手动调优的预算（左）。TriRoute（右）用一个共享的逐词元控制器g_φ替换它们，该控制器在一个全局预算下，在注意力分辨率、专家选择和KV精度上输出耦合策略。词元的绘制面积与它们接收的计算量+内存成比例：稀有实体*“Nakamura”*保留完整注意力和8位缓存，但跳过FFN，而*“the”*在三个轴上都很廉价——这是隔离机制无法共同表达的模式。

联合做出这些决策实际上比孤立做出更难，原因有三，这构成了我们的技术贡献。

(1) **决策空间是异构的。** 专家路由是在一个大的离散集上进行选择，有成熟的负载均衡工具包 [20 (https://arxiv.org/html/2607.06601#bib.bib3), 56 (https://arxiv.org/html/2607.06601#bib.bib5), 55 (https://arxiv.org/html/2607.06601#bib.bib29)]；注意力模式和位宽是小的序数类别，其选项在敏感性上差异巨大（跳过注意力比将缓存从8位降到4位危害大得多）。通过所有这些进行反向传播需要结合 Gumbel-Softmax / 直通估计器 [27 (https://arxiv.org/html/2607.06601#bib.bib7), 35 (https://arxiv.org/html/2607.06601#bib.bib8), 4 (https://arxiv.org/html/2607.06601#bib.bib9)] 与 softmax 门控，并使其梯度保持在可比较的尺度上。

(2) **路由坍缩可能跨轴级联。** 稀疏路由器容易坍缩——将所有词元发送到一个选项——MoE文献通过辅助负载均衡和z-loss来控制 [42 (https://arxiv.org/html/2607.06601#bib.bib2), 20 (https://arxiv.org/html/2607.06601#bib.bib3), 56 (https://arxiv.org/html/2607.06601#bib.bib5)]。有了三个耦合的路由器，一个轴上的坍缩会*诱导*另一个轴上的坍缩：如果注意力路由器学会激进地跳过，FFN路由器看到的输入分布是退化的、低方差的，并停止区分词元，这进而使得比特路由器的任务不适定。我们凭经验描述了这种*跨轴坍缩级联*，并引入了一个耦合感知的平衡目标来阻止它。

(3) **预算必须是一个单一的可控旋钮。** 实践者关心成本-质量权衡曲线上的一个点，而不是三个独立的稀疏性超参数。我们将FLOPs和内存折叠成一个可微的成本模型，并通过一个拉格朗日约束强制执行平均目标预算C⋆，其对偶变量在线自适应，使得一个单一标量就能扫遍整个帕累托前沿。

#### 贡献。

- • 我们将*统一的自适应计算*表述为一个跨三个耦合轴——注意力分辨率、FFN专家和KV缓存精度——的逐词元、逐层路由问题，并将其实例化为TriRoute，据我们所知，这是第一个用一个控制器联合学习所有三个轴的架构（第3节）。
- • 我们给出了一个异构松弛和平衡配方：逐轴温度退火的直通Gumbel估计器，一个耦合感知的负载均衡损失，用于防止跨轴坍缩级联，以及一个在线拉格朗日预算控制器，暴露一个单一的成本旋钮（第3.4、3.6和3.7节）。
- • 我们研究了统一所开辟的设计空间：*共享 vs. 独立*路由器（表示共享有助于，直到高稀疏度下的干扰），以及路由*粒度*（词元级、逐头、逐层组），发现逐头注意力路由结合词元级FFN/位路由是一个最佳点（第5.2节）。
- • 在160M–1.3B参数的仅解码器模型上，以计算最优的词元预算训练 [25 (https://arxiv.org/html/2607.06601#bib.bib10)]，TriRoute在匹配的推理FLOPs和内存下帕累托主导了独立调优的MoD、MoE和KV量化的最佳组合，并且——关键的是——更好地保留了*尾部案例*在稀有实体、代码和算术上的准确性，而平均困惑度最优的自适应模型倾向于在这些方面退化（第5节）。
- • 我们展示了学习到的策略是可解释的：路由模式沿着语言轴（句子边界、稀有子词、句法功能）聚类，提供了关于*哪里*节省了计算量的机理说明（第6节）。

我们发布了一个参考PyTorch控制器和成本模型实现（附录C），以使统一的公式易于采用。

## 2 相关工作

TriRoute位于四个研究方向的交汇处。我们回顾每一个，并关键指出，所有先前的工作都优化*单一*轴；TriRoute的贡献在于跨轴的耦合控制器。

#### 稀疏混合专家（FFN轴）。

通过学习门控进行条件计算可追溯到 Shazeer 等人 [42 (https://arxiv.org/html/2607.06601#bib.bib2)]，他们引入了具有噪声 top-k softmax 和负载均衡损失的稀疏门控 MoE 层。GShard [31 (https://arxiv.org/html/2607.06601#bib.bib4)] 和 Switch Transformer [20 (https://arxiv.org/html/2607.06601#bib.bib3)] 将其扩展到万亿参数，并将路由简化为 top-1，暴露了*坍缩*和*不平衡*这对孪生故障模式，这激发了辅助损失。ST-MoE [56 (https://arxiv.org/html/2607.06601#bib.bib5)] 加入了路由器 z-loss 以增强稳定性。替代词元选择路由的方法包括专家选择路由 [55 (https://arxiv.org/html/2607.06601#bib.bib29)]、作为最优分配问题的 BASE 层 [32 (https://arxiv.org/html/2607.06601#bib.bib27)] 以及固定哈希 [39 (https://arxiv.org/html/2607.06601#bib.bib28)]。系统工程 [37 (https://arxiv.org/html/2607.06601#bib.bib31), 22 (https://arxiv.org/html/2607.06601#bib.bib32)] 和最近的开源模型 [28 (https://arxiv.org/html/2607.06601#bib.bib55), 12 (https://arxiv.org/html/2607.06601#bib.bib56)] 使 MoE 成为在不按比例增加 FLOPs 的情况下扩展容量的默认方式。所有这些都只路由 FFN；注意力和缓存是密集的。TriRoute将MoE纳入其专家轴，并通过*空专家*涵盖其深度跳过的特例。

#### 自适应深度与提前退出（块轴）。

自适应计算时间 [24 (https://arxiv.org/html/2607.06601#bib.bib22)] 为RNN学习了一个暂停单元；Universal Transformers [15 (https://arxiv.org/html/2607.06601#bib.bib23)] 将其应用于Transformer。深度自适应Transformer [18 (https://arxiv.org/html/2607.06601#bib.bib24)] 和自信自适应语言建模 (CALM) [41 (https://arxiv.org/html/2607.06601#bib.bib25)] 在词元预测置信度高时提前退出，而LayerDrop [19 (https://arxiv.org/html/2607.06601#bib.bib26)] 随机丢弃层以实现可修剪深度。混合深度 (MoD) [38 (https://arxiv.org/html/2607.06601#bib.bib6)] 与我们的深度行为最接近：一个逐块 top-k 路由器选择哪些词元被块处理，哪些通过残差流绕过，提供了一个适合于训练的静态计算图。MoD 以*整个块*（注意力和FFN一起）为粒度进行路由，并且不涉及缓存精度。TriRoute的不同之处在于 (i) 分解块，使得注意力和FFN可以*分别*决定，以及 (ii) 将深度与缓存分配耦合——一个词元可以跳过其FFN，同时保留完整注意力和高精度缓存，这是MoD无法表达的。CoLT5 [2 (https://arxiv.org/html/2607.06601#bib.bib61)] 类似地将词元路由到长输入上的轻量/重型注意力和FFN分支，但使用两个独立的路由器，并且没有缓存精度决策。

#### 高效与自适应注意力。

大量工作使用固定稀疏模式 [7 (https://arxiv.org/html/2607.06601#bib.bib33), 3 (https://arxiv.org/html/2607.06601#bib.bib34), 50 (https://arxiv.org/html/2607.06601#bib.bib36)]、低秩/哈希近似 [30 (https://arxiv.org/html/2607.06601#bib.bib35)] 或 IO 感知精确核 [13 (https://arxiv.org/html/2607.06601#bib.bib37)] 来降低注意力成本。多查询和分组查询注意力 [43 (https://arxiv.org/html/2607.06601#bib.bib13), 1 (https://arxiv.org/html/2607.06601#bib.bib12)] 通过共享头来缩小KV缓存。*学习的、词元自适应*的注意力比较少见：自适应注意力跨度 [46 (https://arxiv.org/html/2607.06601#bib.bib52)] 学习了逐头跨度，而混合注意力头 [53 (https://arxiv.org/html/2607.06601#bib.bib53)] 和 SwitchHead [11 (https://arxiv.org/html/2607.06601#bib.bib54)] 在*注意力内部*应用 MoE 风格路由（为每个词元选择头/投影）。这些只自适应注意力，保持FFN和缓存不变。我们的注意力轴是互补的：它是一个粗略的分辨率门控（跳过/局部/全局），与专家和比特决策共同优化，而不是单独调优。

#### KV缓存压缩与量化（内存轴）。

对于长上下文服务，KV缓存（而非参数）主导内存。两类方法解决这个问题：*驱逐*，即丢弃低效用词元——H2O保留“重击者” [54 (https://arxiv.org/html/2607.06601#bib.bib16)]，StreamingLLM保留注意力汇聚点和最近词元 [49 (https://arxiv.org/html/2607.06601#bib.bib17)]——以及*量化*，即以低位宽存储键和值。KIVI [33 (https://arxiv.org/html/2607.06601#bib.bib14)] 展示了无需调优的2位非对称每通道/每...