Dirichlet-Prior Shaping: Guiding Expert Specialization in Upcycled MoEs

Leyla Mirvakhabova, Babak Ehteshami Bejnordi, Gaurav Kumar, Hanxue Liang, Wanru Zhao, Paul N. Whatmough
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:89134-89155, 2026.

Abstract

Upcycling pre-trained dense models into sparse Mixture-of-Experts (MoEs) efficiently increases model capacity but often suffers from poor expert specialization due to naive weight replication. We introduce Dirichlet-Prior Shaping Loss (DPSL), a novel router regularization technique that directly shapes routing probability distributions by matching expert assignments to a target Dirichlet prior resulting in enhanced expert specialization. DPSL enables encoding of inductive biases such as encouraging experts to focus on specific modalities or tasks, without requiring manual intervention. DPSL is a general tool applicable to any module that outputs categorical probability distributions, extending its utility beyond MoE training. Experiments on upcycled MoE vision-language models show that DPSL consistently outperforms upcycling strategies and regularization techniques across standard vision-language benchmarks, addressing the critical issue of poor specialization and fostering higher-performing models.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-mirvakhabova26a, title = {{D}irichlet-Prior Shaping: Guiding Expert Specialization in Upcycled {M}o{E}s}, author = {Mirvakhabova, Leyla and Ehteshami Bejnordi, Babak and Kumar, Gaurav and Liang, Hanxue and Zhao, Wanru and Whatmough, Paul N.}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {89134--89155}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/mirvakhabova26a/mirvakhabova26a.pdf}, url = {https://proceedings.mlr.press/v306/mirvakhabova26a.html}, abstract = {Upcycling pre-trained dense models into sparse Mixture-of-Experts (MoEs) efficiently increases model capacity but often suffers from poor expert specialization due to naive weight replication. We introduce Dirichlet-Prior Shaping Loss (DPSL), a novel router regularization technique that directly shapes routing probability distributions by matching expert assignments to a target Dirichlet prior resulting in enhanced expert specialization. DPSL enables encoding of inductive biases such as encouraging experts to focus on specific modalities or tasks, without requiring manual intervention. DPSL is a general tool applicable to any module that outputs categorical probability distributions, extending its utility beyond MoE training. Experiments on upcycled MoE vision-language models show that DPSL consistently outperforms upcycling strategies and regularization techniques across standard vision-language benchmarks, addressing the critical issue of poor specialization and fostering higher-performing models.} }
Endnote
%0 Conference Paper %T Dirichlet-Prior Shaping: Guiding Expert Specialization in Upcycled MoEs %A Leyla Mirvakhabova %A Babak Ehteshami Bejnordi %A Gaurav Kumar %A Hanxue Liang %A Wanru Zhao %A Paul N. Whatmough %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-mirvakhabova26a %I PMLR %P 89134--89155 %U https://proceedings.mlr.press/v306/mirvakhabova26a.html %V 306 %X Upcycling pre-trained dense models into sparse Mixture-of-Experts (MoEs) efficiently increases model capacity but often suffers from poor expert specialization due to naive weight replication. We introduce Dirichlet-Prior Shaping Loss (DPSL), a novel router regularization technique that directly shapes routing probability distributions by matching expert assignments to a target Dirichlet prior resulting in enhanced expert specialization. DPSL enables encoding of inductive biases such as encouraging experts to focus on specific modalities or tasks, without requiring manual intervention. DPSL is a general tool applicable to any module that outputs categorical probability distributions, extending its utility beyond MoE training. Experiments on upcycled MoE vision-language models show that DPSL consistently outperforms upcycling strategies and regularization techniques across standard vision-language benchmarks, addressing the critical issue of poor specialization and fostering higher-performing models.
APA
Mirvakhabova, L., Ehteshami Bejnordi, B., Kumar, G., Liang, H., Zhao, W. & Whatmough, P.N.. (2026). Dirichlet-Prior Shaping: Guiding Expert Specialization in Upcycled MoEs. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:89134-89155 Available from https://proceedings.mlr.press/v306/mirvakhabova26a.html.

Related Material