TDMoE: Tail-probability-based Dynamic-k MoE

Mixture-of-Experts (MoE) large language models typically rely on fixed top-k routing that activates the same number of experts for every token at every layer, which ignores token-wise and layer-wise heterogeneity and can waste computation on easy inputs while under-allocating capacity to harder ones. Existing dynamic routing approaches either depend on predefined heuristics with hand-tuned hyperparameters to switch routing behaviors based on input features, or introduce new routing modules that often require reinforcement learning because expert selection involves discrete sampling. To address these limitations, we propose Tail-probability-based Dynamic-k MoE (TDMoE), a dynamic routing framework that enables stable and efficient optimization through a fully differentiable end-to-end training strategy. TDMoE introduces a lightweight module KPpredictor to predict the number of experts to activate for each token at each layer, after which the router performs expert selection conditioned on the predicted expert count. During training, we replace discrete top-k sampling with a tail-probability-based soft routing scheme, and during inference we retain standard hard top-k routing for efficiency. We further incorporate auxiliary regularization objectives that control overall expert utilization and sharpen the expert-count distribution toward decisive hard routing. We evaluate TDMoE by fine-tuning on multiple reasoning benchmarks and comparing against a Dynamic-MoE baseline using the OpenCompass evaluation toolkit.