ICML2026

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

Can Jin, Hongwu Peng, Mingcan Xiang, Qixin Zhang, Xiangchi Yuan, Amit Hasan, Ohi Dibua, Yifan Gong, Yan Kang, Dimitris Metaxas

3 citations

Abstract

Sparse Mixture-of-Experts architectures are essential for scaling model capacity efficiently, yet the standard Top-kk routing imposes a rigid sparsity pattern that ignores the intrinsic variance in token difficulty and layer-specific computational needs. Top-pp 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-pp implementations with fixed global probability thresholds provide only marginal gains over Top-kk, suffer from hyperparameter sensitivity, and result in uncontrolled computational costs. In this paper, we propose DTop-pp, a sparsity-controllable dynamic routing mechanism that learns the Top-pp 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-pp consistently outperforms both Top-kk and fixed Top-pp baselines while matching the average FLOPs of Top-kk MoE. Our analysis confirms that DTop-pp 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.