Sparsity-Controllable Dynamic Top-p MoE for Large Foundation Model Pre-training

TL;DR

This paper introduces DTop-p MoE, a dynamic routing method for sparse Mixture-of-Experts models that controls sparsity through a PI controller, improving efficiency and adaptability in large model pre-training.

Contribution

It proposes a novel dynamic Top-p routing mechanism with a PI controller for adaptive sparsity control and layer-wise routing normalization, advancing MoE scalability and efficiency.

Findings

DTop-p outperforms Top-k and fixed-threshold Top-p baselines.

It maintains precise expert activation control across tokens and layers.

Demonstrates strong scaling with model size and dataset complexity.

Abstract

Sparse Mixture-of-Experts (MoE) architectures effectively scale model capacity by activating only a subset of experts for each input token. However, the standard Top-k routing strategy imposes a uniform sparsity pattern that ignores the varying difficulty of tokens. While Top-p routing offers a flexible alternative, existing implementations typically rely on a fixed global probability threshold, which results in uncontrolled computational costs and sensitivity to hyperparameter selection. In this paper, we propose DTop-p MoE, a sparsity-controllable dynamic Top-p routing mechanism. To resolve the challenge of optimizing a non-differentiable threshold, we utilize a Proportional-Integral (PI) Controller that dynamically adjusts the probability threshold to align the running activated-expert sparsity with a specified target. Furthermore, we introduce a dynamic routing normalization mechanism that adapts layer-wise routing logits, allowing different layers to learn distinct expert-selection patterns while utilizing a global probability threshold. Extensive experiments on Large Language Models and Diffusion Transformers demonstrate that DTop-p consistently outperforms both Top-k and fixed-threshold Top-p baselines. Our analysis confirms that DTop-p maintains precise control over the number of activated experts while adaptively allocating resources across different tokens and layers. Furthermore, DTop-p exhibits strong scaling properties with respect to expert granularity, expert capacity, model size, and dataset size, offering a robust framework for large-scale MoE pre-training.