DTop-p MoE: Sparsity-Controlled Dynamic Top-p MoE for Foundation Model Pre-training

Sparse Mixture-of-Experts architectures are essential for scaling model capacity efficiently, yet the standard Top-$k$ routing imposes a rigid sparsity pattern that ignores the intrinsic variance in token difficulty and layer-specific computational needs. Top-$p$ routing is more adaptive because it selects experts until their cumulative routing probability reaches a threshold, allowing confident tokens to use fewer experts and ambiguous tokens to recruit more. However, we demonstrate that existing naive Top-$p$ implementations with fixed global probability thresholds provide only marginal gains over Top-$k$, suffer from hyperparameter sensitivity, and result in uncontrolled computational costs. In this paper, we propose **DTop-$p$**, a sparsity-controllable dynamic routing mechanism that learns the Top-$p$ probability threshold with a Proportional-Integral controller and uses dynamic routing normalization to support layer-wise expert selection under a global sparsity constraint. Extensive experiments on Large Language Models and Diffusion Transformers demonstrate that **DTop-$p$** consistently outperforms both Top-$k$ and fixed Top-$p$ baselines while matching the average FLOPs of Top-$k$ MoE. Our analysis confirms that **DTop-$p$** exhibits strong scaling properties across expert granularity, total expert capacity, model size, and dataset size, offering a robust and efficient MoE framework for foundation model pre-training.

Modern AI systems like large language models are enormous, and running every part of them for every word they process is wasteful and expensive. A popular fix is "Mixture-of-Experts": instead of using the whole network, the model routes each piece of text to a small handful of specialized sub-networks, or "experts." The standard approach forces every input to use the same fixed number of experts. But that's like assigning the same number of staff to every task regardless of difficulty—simple inputs get more help than they need, while hard ones get too little. A natural alternative lets the model decide how many experts to use per input, but existing versions of this idea are unpredictable: the computing cost swings wildly during training, which makes them hard to budget for and prone to crashing. We borrowed a tool from classical engineering—the same kind of feedback controller that keeps a thermostat steady—to automatically tune how many experts get used, holding the average cost to a fixed target while still letting individual inputs draw on more or fewer experts as needed. Across language and image tasks, our method beats the standard approach at the same computing cost. This makes large AI models more efficient and easier to scale under practical training budgets.