用注意力权重来引导 MoE 路由能让 OLMoE 在 GSM8K 上提升 3.37 个百分点

MoE 模型的路由机制一直有个痛点：Router 选专家全靠当前 Token 的隐藏状态，能拿到的上下文信息太少。arXiv 2609.20974v1 提出的 Attention-Aware Routing（AAR）换了个思路，直接把滑动窗口内的注意力权重（包含时间特征和频谱特征）喂给路由。这种做法把路由和注意力机制耦合在一起，在不动 Base Transformer 权重、只训练路由参数的情况下，能有效缓解长文本生成发散的问题，且对数学推理有明显加持。

为什么要把注意力权重交给路由

传统的 MoE 路由是基于隐藏状态（Hidden State）做决策的，但隐藏状态和上下文的全局状态是耦合在一起的。AAR 的核心逻辑是把注意力权重作为一种“上下文状态摘要”，将其与隐藏状态解耦，作为额外的特征输入给 Router。

这种设计的实际效果体现在两个方面：
第一，它能抑制那些没意义的冗长错误生成。实测发现，应用 AAR 后，模型给出错误答案时的长度会缩短，而正确答案的长度保持不变，这意味着模型在意识到走错路时能更快“收敛”，而不是在那儿瞎编一段长文。
第二，它揭示了路由与注意力之间存在一种“耦合电路”。当第 l 层的路由发生变化时，这种影响会通过残差流传递，从而放大第 l+1 层的注意力汇聚（Attention Sinks），在不需要更新注意力机制本身的情况下，就改变了注意力的分布。

AAR 在 OLMoE 上的实测数据

研究团队在 OLMoE 模型上做了验证，实验设置非常干净：完全冻结 Base Transformer，只训练路由参数，以此隔离变量。

• 性能提升： 在 GSM8K 数学推理数据集上，AAR 相比于仅进行路由微调（Routing-only SFT）的基线，得分提升了 3.37 个百分点。
• 深度敏感性： 这是一个很关键的坑。AAR 不是在所有层全量铺开就行的，它对层深非常敏感。如果不管三七二十一在所有层都用，反而会损害事实检索（Factual Retrieval）的能力；但如果把它放在网络的深层，数学推理的增益依然存在。

这意味着模型在不同深度处理信息的方式不同：浅层更多负责检索，深层更多负责推理。这种“检索-推理”的张力决定了 AAR 必须进行层选择性地应用。

如何理解 AAR 的实现逻辑

如果要复现或者研究这个机制，重点在于 AAR 提取特征的方式。它不是简单地把注意力权重丢进去，而是通过一个滑动窗口提取两种特征：

• 时间特征（Temporal Features）： 捕捉 Token 序列在时间维度上的注意力分布演变。
• 频谱特征（Spectral Features）： 通过分析注意力权重的频率分布，提取出与隐藏状态解耦的上下文摘要。

具体到操作流程，可以概括为： 1. 维护一个注意力权重的滑动窗口。 2. 提取该窗口内的时间与频谱特征。 3. 将这些特征与当前的 Token Hidden State 拼接。 4. 将拼接后的向量输入给路由权重矩阵，决定专家选择。

实际应用中的权衡

如果你打算在自己的 MoE 项目中尝试类似的耦合路由，需要注意以下几点：

成本与开销： 由于 AAR 只需要训练路由参数，其训练开销远低于全量 SFT。但推理时需要维护注意力权重的窗口并计算频谱特征，这会带来一定的额外计算量。

适用场景：

• 如果你的任务偏向数学推理、逻辑推演，把 AAR 放在深层大概率有效。
• 如果你的任务是纯粹的知识问答（Fact Retrieval），谨慎在浅层使用，否则可能会出现检索精度下降的情况。

风险点： 这种方法本质上是在利用注意力机制的副作用来引导路由。虽然在 OLMoE 上验证了可行性，但在不同规模或不同预训练目标的 MoE 模型上，这种“耦合电路”的强度可能不同，不保证在所有模型上都能获得 3.37 pp 这种量级的提升。