Towards a Statistical Understanding of Mixture-of-Experts

Siyuan He; Bokai Yang; Jie Hu; Ziwen Gao; Yuhong Yang

stat.MLSep 3, 2026

Towards a Statistical Understanding of Mixture-of-Experts

Authors: Siyuan He, Bokai Yang, Jie Hu, Ziwen Gao, Yuhong Yang

Abstract

Mixture-of-experts (MoE) architectures increase model capacity by combining a collection of expert predictors through input-dependent routing, while often activating only a small subset of experts for each input. Despite their growing importance in modern large-scale models, the statistical roles of their design choices, especially routing, sparse activation, and shared experts, remain only partially understood, as existing theory has largely focused on parametric or correctly specified MoE models. In this paper, we view MoE as a form of localized aggregation and show how this localization reshapes the approximation-estimation-computation tradeoff. We derive oracle risk bounds for learning dense and sparse routing with evolving experts, separating approximation, expert-learning, and router-estimation errors, and characterize how sparse Top-K routing can retain the benefits of localized aggregation while controlling per-input computation. We also interpret gating through the geometry of input space, relating routing performance to regions of local expert advantage, and show how shared experts, as adopted in architectures such as DeepSeekMoE, can extract common predictive structure so that routed experts focus on residual local variation. Together, these results provide a unified statistical framework for understanding MoE through input-dependent expert aggregation, in which expert specialization and computational tradeoffs are governed by local predictive structure.

Explore similar work

May 6, 2026cs.LG

Expert Routing for Communication-Efficient MoE via Finite Expert Banks

Resource-efficient machine learning increasingly uses sparse Mixture-of-Experts (MoE) architectures, where the gate acts as both a learning component and a routing interface controlling computation, communication, and accuracy. Motivated by finite-rate interpretations of MoE gating, we treat the gate as a stochastic channel and use

I(X;T)

to quantify the routing information available to the selected expert. To make the associated information quantities tractable beyond synthetic examples, we develop a finite-bank MNIST construction using pretrained CNN experts and a discrete, data-dependent selection rule. Since the selected model belongs to a finite candidate set, the algorithmic mutual information

I(S;W)

admits a closed-form discrete-entropy estimator from the empirical posterior

q(W|S)

. Sweeping a data-dependence parameter

α

, we observe that

\widehat I(S;W)

monotonically tracks the generalization gap, while the Xu-Raginsky bound exhibits the expected looseness. We also compare with a uniform union-bound baseline and introduce an empirical estimator of

I(X;T)

together with a Blahut-Arimoto procedure for tracing an accuracy-rate curve over the expert bank. The proposed framework provides a practical tool for analyzing resource-aware MoE inference systems and for interpreting

I(X;T)

and

D(R_g)

as design proxies for efficient expert routing.

Mohammad Reza Deylam Salehi, Ali Khalesi

Mixture-Of-Expert Inference Mixture-of-Experts Models Experts Expert Routing Channel Modeling Mnist+1

Ask PDF

Jun 1, 2026cs.LG

ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

Mixture-of-Experts (MoE) models scale by activating only a small subset of experts per token. However, training such models remains challenging because top-

k

routing is discrete and non-differentiable, requiring gradient estimators for expert selection whose design remains a central open problem. We introduce ProbMoE, a probabilistic routing framework that models expert selection as a distribution over cardinality-constrained expert subsets and formulates routing as probabilistic inference in this discrete subset space. We first propose ProbMoE Exact-

k

routing, which samples

k

-expert subsets in the forward pass, and the backward pass uses gradients through each expert's exact marginal probability as a tractable surrogate for the true gradient. ProbMoE naturally generalizes to a dynamic-

k

routing setting, where both training and inference constrain the routing cardinality to the same predefined range, allowing adaptive expert allocation per token. Across benchmarks and model backbones, ProbMoE Exact-

k

achieves strong performance compared to competitive baselines, with improved expert utilization and routing diversity; ProbMoE Dynamic-

k

achieves comparable performance with fewer activated experts.

Heng Zhao, Zilei Shao, Guy Van den Broeck +1

Probabilistic Inference Cardinality

Ask PDF

May 31, 2026cs.AI

DAG-MoE: From Simple Mixture to Structural Aggregation in Mixture-of-Experts

Mixture-of-Experts (MoE) models have become a leading approach for decoupling parameter count from computational cost in large language models, yet effectively scaling MoE performance remains a challenge. Prior work shows that fine-grained experts enlarge the space of expert combinations and improve flexibility, but they also impose substantial routing overhead, creating a new scalability bottleneck. In this paper, we explore a complementary axis for scaling -- how expert outputs are aggregated. We theoretically show that replacing the standard weighted-summation aggregation with structural aggregation expands the expert-combination space without altering the experts or router, and enables possible multi-step reasoning within a single MoE layer. To this end, we propose DAG-MoE, a sparse MoE framework that employs a lightweight module to automatically learn the optimal aggregation structure among the selected experts. Extensive experiments under standard language modeling settings show that DAG-MoE consistently improves performance in both pretraining and fine-tuning, surpassing traditional MoE baselines.

Jiarui Feng, Hanqing Zeng, Karish Grover +11

Mixture-of-Experts Models

Ask PDF