专家-路由器耦合损失（ERC）优化MoE模型训练，字节团队提出控制专家专业化新方法

混合专家 (Mixture-of-Experts， MoE) 模型通过多个小型专家网络和一个路由器实现稀疏计算，每个输入token仅激活少数专家。这种设计在保持计算效率的同时，实现了参数规模的大幅扩展，使得训练万亿参数级别的模型成为可能。理想情况下，路由器应当精确掌握每个专家的实际能力，以实现高效的token路由。然而，传统的MoE模型缺乏明确的约束机制来保证这一点。由于路由器无法直接访问专家参数，它只能通过试错来学习路由策略，这常常导致token被错误路由，其梯度会反过来干扰专家的专业化过程。

为了解决这一核心问题，字节跳动Seed团队联合中国人民大学高瓴人工智能学院提出了一种轻量级的辅助损失——专家-路由器耦合损失 (Expert-Router Coupling loss， ERC)。该损失能以极小的计算开销，将路由器的决策逻辑与专家的实际能力紧密耦合起来。ERC损失的计算效率很高，仅需处理 n² 个激活值（n 为专家数量），形成一个与批次大小无关的固定计算成本。相比之下，以往的一些耦合方法计算成本通常随token数量线性增长。研究团队在参数量分别为3B和15B的MoE大语言模型上进行了预训练，验证了ERC损失的有效性。此外，ERC损失还能在训练过程中灵活控制并定量追踪专家的专业化水平，为深入理解MoE模型的工作机制提供了宝贵的洞见。

• 论文标题：Coupling Experts and Routers in Mixture-of-Experts via an Auxiliary Loss
• 论文链接：https://arxiv.org/pdf/2512.23447

方法

研究团队提出的专家-路由器耦合损失 (ERC)，旨在以极小的额外开销将路由器与专家紧密耦合。该损失将路由器参数矩阵 R ∈ ℝⁿ×ᵈ 解释为聚类中心，其中每一行 R[i] 作为被路由至第 i 个专家的token集合 ℭ_i 的中心。如图1所示，ERC损失包含三个关键步骤：

(1) 生成代理令牌
对路由器嵌入 R[i] 添加有界随机噪声 δ_i（δ_i∼U(1−ε_i,1+ε_i)^d）得到 R̃[i]，作为 ℭ_i 中token的代理表示。噪声 δ_i 的范围被限制在相邻聚类中心之间最小距离的一半，以确保该噪声能够模拟 ℭ_i 内部的输入变化，同时避免跨越聚类边界。

ε_i ≤ ‖R[i]−R[j]‖_2 / (2‖R[i]‖_2)

(2) 计算专家激活矩阵
受先前工作的启发，研究团队使用中间激活范数作为衡量专家能力与输入token匹配程度的指标。每个代理token R̃[i] 由所有 n 个专家的 W_g 参数进行处理，共产生 n² 个中间激活值。计算所有专家的中间激活范数，产生一个矩阵 M ∈ ℝⁿ×ⁿ。其中，M[i,j] 表示专家 j 在给定输入 R̃[i] 时的激活范数。

M[i,j]=‖R̃[i]⋅W_g^j‖

(3) 施加损失约束
对于所有 i≠j，当非对角元素 M[i,j] 或 M[j,i] 超过 αM[i,i] 时，ERC损失就会施加惩罚。其中，α 是一个标量超参数：

ℒ_ERC = 1/n² Σ_(i=1 to n) Σ_(j≠i to n) ( max(M[i,j] - αM[i,i], 0) + max(M[j,i] - αM[i,i], 0) )

研究团队通过最小化上述ERC损失，实现了专家与路由器之间的紧密耦合，具体体现在两个机制上：

• 专家专业化：代理token R̃[i] 在第 i 个专家中引发的激活强度高于在所有其他专家中引发的强度。这表明第 i 个专家已被优化为最匹配其对应token集合 ℭ_i 的特征。
• 精确的token路由：第 i 个专家对指定向量 R̃[i] 的激活响应强于对任意其他 R̃[j]（j≠i）的响应。这说明 R[i] 与第 i 个专家的能力高度对齐，从而确保路由器将最需要该专家的token分配给它。

专家-路由器耦合损失在 PyTorch 中的伪代码实现如下：

import torch
import torch.nn as nn
import PseudoExpertClass

class MoE(nn.Module):

    def __init__(self, args):
        super().__init__()
        self.experts = PseudoExpertClass(args)
        self.R = torch.nn.Parameter(torch.empty(args.n, args.d))
        self.alpha = args.alpha

    def erc_loss(self, M):
        row_diff = (M - self.alpha * torch.diag(M).unsqueeze(1))
        row_diff_clamped = torch.clamp(row_diff, min=0.0)

        col_diff = (M - self.alpha * torch.diag(M).unsqueeze(0))
        col_diff_clamped = torch.clamp(col_diff, min=0.0)

        mask = torch.ones_like(A) - torch.eye(A.size(0), device=A.device)
        total_diff = (row_diff_clamped + col_diff_clamped) * mask
        return total_diff.mean()

    def get_noisy_router(self, R):
        with torch.no_grad():
            norm_R = torch.norm(R, dim=1)
            distances = torch.cdist(R, R, p=2)
            distances.fill_diagonal_(float('inf'))
            min_dist, _ = torch.min(distances, dim=1)
            eps = min_dist / 2 / norm_R

            low = (1 - eps).unsqueeze(1)
            high = (1 + eps).unsqueeze(1)
            noise = torch.rand_like(R)
            return (low + noise * (high - low)) * R

    def forward(self, x):
        erc_loss = 0.0
        if self.training:
            R = self.get_noisy_router(self.R)
            M = torch.norm(torch.einsum('jDd,id->ijD', self.experts.Wg, R), dim=-1)
            erc_loss = self.erc_loss(M)

        logits = x.view(-1, x.shape[-1]) @ self.R.T
        scores = logits.softmax(dim=-1)
        expert_weights, expert_indices = torch.topk(scores, dim=-1)

        return self.experts(x, expert_weights, expert_indices), erc_loss

评估

研究团队将引入ERC损失的MoE模型与原始MoE模型、以及注意力专家 (Attention over Experts， AoE) 基线进行了比较。所有模型均基于OLMo架构从头开始训练，参数量为3B。模型包含12层，d=1536、D=768。每个Transformer层包含16个注意力头与 n=64 个专家，每个输入token激活 K=8 个专家。所有模型均在开源数据集dolmap-v1.5-sample的500B个token上进行训练。对于ERC损失，若无特别说明，默认使用 α = 1，损失权重固定为1。

研究团队在ARC-Challenge、CommonsenseQA、COPA、BoolQ、HellaSwag、OpenbookQA、SciQ、Social IQa、WinoGrande和MMLU等多个任务上对模型进行了评估。

图3(a)展示了所有任务的平均准确率。结果表明，引入ERC损失的MoE模型取得了稳定的性能提升，显著优于原始MoE模型，并有效缩小了AoE与MoE之间的性能差距。具体任务的详细结果如图8所示。

在效率方面，使用与不使用ERC损失的MoE模型在训练吞吐量和内存开销上几乎相同。相比之下，AoE模型需要增加1.6倍的训练时长和1.3倍的内存开销，其训练效率问题限制了进一步扩展。

专家-路由器耦合损失与常用的负载均衡损失具有兼容性。如图3(b)所示，结合ERC损失的MoE与原始MoE在负载均衡损失方面的差异约为 10⁻⁵ 量级。考虑到整体负载均衡损失的数值保持在 10⁻² 左右，这一差异可以忽略不计。

研究团队进一步将专家数量 n 增加至256（保持 K=8）并将模型深度加倍，该配置下模型总参数量达到15B。如表1所示，ERC损失带来的性能优势在多个更具挑战性的公开基准测试中依然持续存在。这一致的性能提升表明，该方法即使在大规模场景下也能有效缓解专家与路由器之间的解耦问题。在整个大规模训练过程中，未观察到任何损失尖峰或异常梯度。

专家参数的可视化如图4所示。原始MoE中的专家参数特征未能形成有意义的簇结构；相比之下，引入ERC损失的MoE显示出清晰的聚类结构，彰显了专家模块的专业化能力。

ERC损失不仅能够促进专业化，还可作为探索专业化特性的强大工具。研究团队通过以下两个特性来展示这种能力。

特性1：通过 α 实现可控的专业化程度调节
在ERC损失中，超参数 α 控制专家与路由器之间的耦合强度。当 α = 0 时，ERC损失促使 R[i] 与其他专家的参数正交，从而实现最大程度的专业化。而当 α → 1 时，损失允许所有专家对 R[i] 的响应差异更小，从而降低专业化。值得注意的是，α = 1 仅会最大程度弱化ERC损失的约束，其保留的专业化程度仍强于原始MoE模型中自发形成的专业化。

特性2：ϵ 为专业化提供了量化指标
噪声幅度 ϵ 与 α 高度相关，并能反映训练过程中专家专业化程度的变化。这种相关性是因为随着 α 增大，专家被允许变得更加同质化；专家之间同质性的增强会削弱路由器中各簇中心之间的分离；而簇中心间距的减小将导致 ϵ 变小。因此，ϵ 可作为跟踪专家专业化程度的量化指标。

这项研究为优化 Transformer 架构中的MoE组件提供了一种新颖且高效的思路。通过在训练中引入轻量级的专家-路由器耦合损失，不仅能稳定提升模型在下游任务上的性能，还能实现对专家专业化水平的精细控制和定量分析，这对于理解和设计更高效的大型稀疏模型具有重要意义。更多关于深度学习与模型训练的讨论，欢迎访问云栈社区交流分享。