Dirichlet-Prior Shaping: Guiding Expert Specialization in Upcycled MoEs

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.

Modern AI systems improve as they grow larger, but training them from scratch is costly. A practical alternative is to “upcycle” existing models into Mixture-of-Experts (MoE) systems, where different parts handle different inputs. However, since these experts start identically, they can fail to specialize, limiting the benefits of the architecture. We introduce Dirichlet-Prior Shaping Loss (DPSL), a method that guides how the model assigns inputs to experts. Instead of leaving this entirely unconstrained, DPSL nudges the model toward structured decisions, encouraging different experts to focus on different types of inputs, such as images or text, without manual design. Our approach leads to better expert specialization and improved performance on vision-language tasks, helping build more efficient and capable AI systems at lower cost.