一句话定义:负载均衡损失(Load Balancing Loss)是训练混合专家模型(MoE, Mixture of Experts)时,额外加进总损失里的一个惩罚项,目的是逼着路由器(Router,也叫 Gating Network)把 token 尽量均匀地分给各个专家,而不是总丢给同几个"红人"。
先看它要解决什么问题
MoE 的做法是:把 Transformer 里那个又大又笨的前馈网络(FFN)拆成 N 个并行的"专家",每个 token 进来时,由一个很小的路由器打分,只挑分数最高的 k 个专家来处理。这样模型总参数量可以做到很大,但每个 token 实际只算一小部分,推理成本大幅下降。
问题出在训练上。路由器一开始是随机的,但总有某个专家"运气好",早期多接了几个 token,梯度多、学得快,于是下一轮路由器更愿意选它——正反馈一旦转起来,token 会迅速向少数专家集中,剩下的专家常年接不到活、拿不到梯度,参数几乎不更新。业内管这叫专家坍缩(Expert Collapse):显存被 N 个专家占着,实际干活的只有两三个。
打个比方
一个客服团队有 8 个人,前台负责派单。前台发现老王回复又快又好,就把单子全给老王。老王越练越熟,其他人越来越生疏,前台更不敢给他们派单。最后老王一请假,整个团队直接瘫痪。
负载均衡损失就是给前台的 KPI 加一条:你派出去的单子,各人接单量必须大致均衡,偏得越厉害扣分越多。于是前台被迫在"谁最合适"和"大家都有饭吃"之间找平衡。
它长什么样
通常的做法是:统计一个批次里 token 分给各专家的实际比例 f_i,以及路由器给各专家的平均概率 p_i,让 N × Σ(f_i × p_i) 尽量接近 1。越接近 1 说明分配越均匀。这个值乘以一个权重系数,和主任务的语言建模损失加在一起,一起反传。
和相邻概念的区别
| 概念 | 作用 | 性质 |
|---|---|---|
| Top-k 路由 | 决定"token 分给谁" | 主机制 |
| 负载均衡损失 | 决定"分得均不均" | 软引导,附加损失 |
| 容量因子与丢 token(Capacity Factor / AI 词典:Token">Token Dropping) | 限定每个专家最多接多少 token,超了就丢 | 硬约束 |
| L2、Dropout | 防过拟合 | 普通正则化 |
一句话:路由负责分活,负载均衡损失负责盯着别分偏,容量因子则是给每个人设个接单上限。
对从业者和普通人的意义
做训练的人:MoE 训不起来、loss 抖、某些专家梯度长期为零,第一件事就是打印路由分布看看是不是塌了。但那个权重系数不能乱调——太小压不住坍缩,太大又会强行把 token 塞给不合适的专家,损害模型效果,属于典型的"需要手感"的超参。
普通职场人只需要记住一件事:MoE 是今天大模型能做到"参数巨大但调用便宜"的关键结构之一,而让它真正训得起来的那颗螺丝钉,就是这条看起来不起眼的辅助损失。具体实现各家不同,以官方文档和论文页面为准。
