MoE负载均衡改按反馈控制设计：PID (I)在数据分布切换前后都把最忙专家压得更低
核心概要
混合专家（MoE）训练中的数据分布切换会打乱路由器已收敛的专家负载，而以往均衡方法只在平稳数据上评估。这项工作把负载均衡写成反馈控制：每个专家负载偏离均分份额的量是控制误差，加到路由分数上的每专家偏置是控制输入，辅助损失和ALF-LB分别对应比例项与符号积分项。只保留积分项的PID (I)在OLMoE骨干的预训练中途切换数据混合、以及OLMoE-1B-7B的持续预训练两种设置下，切换前后MaxVio都低于对比方法，困惑度相当。
深度剖析
MoE负载均衡被重新表述为一个反馈控制回路，均衡器不再只是附加在训练目标上的惩罚项，而是根据每个专家负载偏离均分份额的误差来调节路由选择。 此前两类主流做法——辅助均衡损失和ALF-LB的无损失偏置更新——主要在平稳训练分布下开发和评估，此时每专家负载缓慢漂移，均衡机制可以稳定下来并保持。 在OLMoE骨干上，预训练从零开始、每层16至64个专家并在训练中途切换数据混合，PID (I)在切换前保持均衡、切换后恢复得更紧，MaxVio在每种设置下都最低，困惑度与辅助损失和ALF-LB相当。
控制输入是加到路由分数上的每专家偏置，它只改变专家选择，不改变门控权重，也不接收梯度。 辅助损失通过同一次反向传播把均衡压力混入路由器权重的梯度，而偏置式方法把这份压力移出梯度，只作用于选择时刻。 把路由器看作带一步延迟的静态单调映射后，只有积分项既能消除持续失衡、又能把逐步负载波动平均掉；比例项和微分项会把这种波动传递或放大进偏置。在OLMoE-1B-7B的持续预训练中，PID (I)冷启动和热启动的postMV分别为0.90和0.87，低于ALF-LB的1.01和1.03。
PID (I)不修改模型权重和训练目标，可以直接插入已有的预训练模型。 持续预训练从已发布的OLMoE-1B-7B检查点出发，先在退火语料上训练到均衡器进入稳态，再切换到代码语料，用来检验均衡器作为即插即用组件的行为。 冷启动检查点的路由器从未见过这些均衡器，其postMV仍接近热启动设置；两种设置下PID (I)的postMV都最低，新领域上的留出困惑度四种方法几乎相同。
启示与展望
演示范围限于OLMoE骨干：预训练从零开始、每层16至64个专家、每token激活八分之一专家，并在训练中途做一次瞬时领域切换；持续预训练从OLMoE-1B-7B-0125检查点出发，在代码语料上做冷启动或热启动。所研究的分布变化是两领域之间的一次硬切换，最大模型为6.9B总参数、1.3B激活参数。
作者指出，均衡器相对ALF-LB的优势随专家数增加而扩大，但三个规模不足以把这一趋势定为结论，生产级专家数下是否延续仍开放。研究只覆盖一次硬切换，未覆盖渐进漂移和两个以上领域间的反复切换；结果来自单一骨干，向专家选择路由或带容量限制的路由器迁移留待后续。持续预训练中偏置式均衡器在相同训练预算内遗忘更多，退火语料上的留出困惑度高于Vanilla和辅助损失。