Distributionally Robust Mixture-of-Experts Training

Distributionally Robust Mixture-of-Experts Training. By Xin Teng and two others, from New York University and NYU Shanghai. Published on the fifth of October. Its code repository is worked out: about ninety-five source files, with a readme and a licence.

Mixture-of-Experts transformers grow capacity by activating only a few experts for each token. That sparsity hides a reliability problem. Routing is never perfect, and a model whose load is nicely balanced can still send a token to an expert that was never trained well enough for that kind of input. Balanced traffic says nothing about whether the expert on the receiving end is ready.

DRMoET, short for Distributionally Robust MoE Training, is a drop-in objective that shifts the target. It treats the experts in each layer as endogenous robustness groups, and optimizes the routing outcomes that carry high loss, instead of merely equalizing how much traffic each expert gets.

Each layer keeps a distribution over its experts. That distribution is updated by an entropy-regularized softmax rule applied to expert losses that are weighted by activation and smoothed with an exponential moving average. The effect is to strengthen plausible routing paths outside the top choices, while the standard MoE computation stays unchanged.

Experiments follow the FLAME-MoE recipe at two scales: seven hundred forty-six million total parameters, and ten point three billion total. At both, DRMoET raises downstream averages over standard FLAME-MoE and over auxiliary-loss-free balancing. At the larger scale, trained on sixty-seven billion tokens, the seven-task average rises from zero point six six two five to zero point six seven six seven, while the auxiliary-loss-free baseline reaches zero point six four three one.

Mechanistic analyses show lower variance in expert loss with the mean loss almost unchanged, four point three percent lower excess loss when tokens are deliberately misrouted to mid-k experts, and better specialization of experts by domain.

The argument is about what a sparse model should be balancing. Spreading tokens evenly is one goal. Making sure that the expert a token lands on can handle it, even when routing slips, is another, and the authors position that routing robustness as a practical objective for scaling sparse MoE reliably.