Distributionally Robust Mixture-of-Experts Training
DRMoET improves imperfect-routing robustness, raising a 10.3B MoE model's seven-task average from 0.6625 to 0.6767.
Distributionally Robust MoE Training (DRMoET) treats layer-wise experts as robustness groups and optimizes high-loss routing outcomes instead of only equalizing token traffic. It updates a per-layer expert distribution with an entropy-regularized softmax over EMA-smoothed, activation-weighted expert losses while preserving standard MoE computation. Under the FLAME-MoE recipe at 10.3B total parameters and 67B tokens, the seven-task average rose from 0.6625 to 0.6767, compared with 0.6431 for auxiliary-loss-free balancing. Analyses report lower expert-loss variance, 4.3% lower excess loss under forced mid-k misrouting, and stronger domain-expert specialization.
- DRMoET optimizes high-loss expert routing rather than only balancing traffic.
- At 10.3B parameters, seven-task average rises from 0.6625 to 0.6767.
- Auxiliary-loss-free balancing scores 0.6431 on the same average.
- Forced mid-k misrouting excess loss falls by 4.3 percent.
From the original article
1 min read at huggingface.co
Mixture-of-Experts (MoE) transformers scale capacity by activating only a few experts per token, but this sparsity creates a hidden reliability problem: when routing is imperfect, load-balanced models may send tokens to experts that are insufficiently trained for the assigned inputs.
By Xin Teng, Muxiao Li, Hongyi Wen · huggingface.co
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.