注意力感知路由：MoE中路由与注意力的耦合

摘要

在混合专家（MoE）语言模型中，路由器通常仅基于token的隐藏状态来选择并加权专家，所利用的上下文信息有限。我们提出注意力感知路由（AAR），用从注意力权重滑动窗口中提取的时序与频谱特征来增强路由器，这些特征代表模型上下文状态的摘要，并与隐藏状态解耦。我们在基础Transformer完全冻结的前提下仅训练路由参数，将路由隔离为唯一变量。在OLMoE上，AAR相比仅微调路由的SFT基线将GSM8K提升3.37个百分点。除性能提升外，我们证明路由与注意力构成耦合回路：第l层的路由变化经残差流传播，会在第l+1层放大注意力汇聚（attention sink），在未直接更新注意力机制的情况下重塑注意力。此外，AAR减少了发散的长生成：错误答案变短，而正确答案长度不变。最后，AAR对深度高度敏感：不加区分地应用于所有层可能损害事实检索，而在网络较深处引入时数学推理增益依然保持。这种敏感性揭示了跨深度的检索—推理张力，并使层选择性AAR成为探测不同层注意力所携带路由相关信息的受控探针。

注意力感知路由：MoE中路由与注意力的耦合：论文配图 — 图 1：注意力感知路由。对于当前 token i，对前 M 个 token 的注意力权重在各个头之间取平均并组装成滑动窗口 w，再变换为增广窗口 w̃ = φ(w)。注意力路由器 A_l 将 w̃ 映射为专家 logits g^attn，并与标准路由器的 logits g^std = W_r h 合并和门控，从而在专家 E_1, …, E_N 上产生最终的路由决策。