Towards Principled Design of Mixture-of-Experts Language Models under Memory and Inference Constraints
Seng Pei Liew, Kenta Shinzato, Yuyang Dong
TL;DR
The paper tackles how to design Mixture-of-Experts language models under memory and inference constraints beyond the conventional focus on total and active parameters. It introduces a five-variable MoE framework and demonstrates through ablations and scaling-law analyses that $N_{total}$ and expert sparsity $s:=n_{exp}/n_{topk}$ primarily determine performance, with an $n_{exp}$ penalty arising when budgets fix $N_{total}$ and $s$. It derives a disambiguated scaling law $L \propto {N_{total}}^{-0.052} {n_{exp}}^{0.023} {n_{topk}}^{-0.018}$ and advocates maximizing $N_{total}$ while minimizing $s$ and $n_{exp}$ within constraints, supported by an iterative optimization routine. The findings offer a rigorous, implementable guide for memory- and compute-constrained MoE design and are validated by comparing configurations under equivalent budgets.
Abstract
Modern Mixture-of-Experts (MoE) language models are designed based on total parameters (memory footprint) and active parameters (inference cost). However, we find these two factors alone are insufficient to describe an optimal architecture. Through a systematic study, we demonstrate that MoE performance is primarily determined by total parameters ($N_{total}$) and expert sparsity ($s:=n_{exp}/n_{topk}$). Moreover, $n_{exp}$ and $n_{topk}$ do not "cancel out" within the sparsity ratio; instead, a larger total number of experts slightly penalizes performance by forcing a reduction in core model dimensions (depth and width) to meet memory constraints. This motivates a simple principle for MoE design which maximizes $N_{total}$ while minimizing $s$ (maximizing $n_{topk}$) and $n_{exp}$ under the given constraints. Our findings provide a robust framework for resolving architectural ambiguity and guiding MoE design.