Optimal Transport Aggregation for Distributed Mixture-of-Experts

Faïcel Chamroukhi; Nhat Thien Pham

cs.LGDec 15, 2023

Optimal Transport Aggregation for Distributed Mixture-of-Experts

Authors: Faïcel Chamroukhi, Nhat Thien Pham

Organizations: IRT SystemX, 91120 Palaiseau, France · Normandie Univ, UNICAEN, CNRS, LMNO, 14000 Caen, France

Abstract

This paper considers the setting in which local Mixture-of-experts (MoE) models are fitted independently on distributed datasets and must be combined into a single global MoE model. Weighted averaging the local mixtures, whose experts are not aligned across machines, often yields a large mixture with many experts instead of $K$ and does not provide an estimator of the true model parameters. We propose to reduce this large mixture to a $K$ -expert MoE by minimizing an expected transportation divergence, in which the transport plan depends on the covariates. A relaxation of this problem separates the experts from the gating network: the experts are computed by a majorization--minimization algorithm, which decreases the objective monotonically, and the gating network is then fitted to the gates implied by the optimal plan, with an explicit bound on the resulting loss. We show that the aggregated estimator is consistent whenever the local estimators are, without requiring the local models to have the same labeling across machines. The procedure only needs a single one-way communication round. Experiments on simulated and real data show accuracy close to centralized estimation at a fraction of its computation time. Our source code is available at https://github.com/nhat-thien/Distributed-Mixture-Of-Experts.

Figures & tables

Estimator	$\mathcal{T}_{c}(\cdot,f^{*})$	$\mathcal{T}_{c}(\cdot,\bar{f}^{W})$	Test Loglik	$\textbf{MSE}(\cdot,\boldsymbol{\theta}^{*})$	RPE	ARI	Time (s)
GLB	$12.34\pm 1.97$	$105.31\pm 4.88$	$-52182.0\pm 139.7$	$0.0004\pm 0.0001$	$0.3917\pm 0.0057$	$0.9290\pm 0.0024$	$4255.698\pm 283.415$
OT-RED + warm start (ours)	$13.47\pm 2.21$	$\mathbf{104.13\pm 4.84}$	$-55333.0\pm 117.2$	$\mathbf{0.0456\pm 0.0008}$	$\mathbf{0.4834\pm 0.0049}$	$0.9155\pm 0.0031$	$30.721\pm 24.412$
OT-RED + kmeans clustering (ours)	$13.47\pm 2.21$	$\mathbf{104.13\pm 4.84}$	$-55333.0\pm 117.2$	$\mathbf{0.0456\pm 0.0008}$	$\mathbf{0.4834\pm 0.0049}$	$0.9155\pm 0.0031$	$34.015\pm 26.737$
GM	$\mathbf{12.86\pm 2.11}$	$104.71\pm 4.90$	$\mathbf{-55332.5\pm 117.8}$	$\mathbf{0.0456\pm 0.0008}$	$\mathbf{0.4834\pm 0.0049}$	$\mathbf{0.9155\pm 0.0030}$	$\mathbf{28.626\pm 24.219}$
A-AVR, W-AVR, MED, FED	Fail due to structural dimension mismatch ( $K_{m}\neq K$ )

Table 1: Aggregation performance under uncoordinated local expert capacities (

M=16,K_{m}=3,L=48\to K=5,d=20,N=100k

, mean

\pm

std over 50 runs).

Estimator	$\mathcal{T}_{c}(\cdot,\bar{f}^{W})$	RMSE (yrs)	MAE (yrs)	RPE	Corr. ( $r$ )	Time (s)
GLB	$346.4$	$9.16$	$6.37$	$0.7128$	$0.5365$	$323.9$
OT-RED	$\mathbf{245.4}$	$\mathbf{9.27}$	$\mathbf{6.46}$	$\mathbf{0.7300}$	$0.5202$	$\mathbf{63.5}$
GM	$307.6$	$9.34$	$6.52$	$0.7412$	$0.5104$	$117.8$
A-AVR	$307.1$	$9.32$	$6.49$	$0.7381$	$\mathbf{0.5207}$	$117.4$
W-AVR	$311.5$	$9.46$	$6.57$	$0.7604$	$0.5092$	$117.4$
MED	$439.9$	$9.75$	$6.62$	$0.8066$	$0.4726$	$117.8$

Table 2: Performance on the YearPredictionMSD audio benchmark (

M=32

machines,

K=4

experts).

Panel of Figure 1 — Figure 1: One-dimensional example with $M=2$ machines and $K=2$ experts. Top: regression functions; bottom: gating functions. Left: the weighted average $\bar{f}^{W}$ has $L=MK=4$ experts and fails to recover the gates, which collapse to $1/2$ , due to label switching. Right: the reduction $\bar{f}^{R}$ (see Section 2.3 ) has $K=2$ experts, with gates recovered from the transport plan, and is robust to local label switching.

Figure 2: Transportation divergence against supporting sample size

S

Proportions	Estimator	$\mathcal{T}_{c}(\cdot,f^{*})$	$\mathcal{T}_{c}(\cdot,\bar{f}^{W})$	Test Loglik	$\textbf{MSE}(\cdot,\boldsymbol{\theta}^{*})$	RPE	ARI	Time (s)
Balanced $\pi_{k}=0.20$	GLB	$14.65\pm 1.90$	$226.41\pm 7.03$	$-52176.7\pm 152.9$	$0.0003\pm 0.0001$	$0.3885\pm 0.0076$	$0.9291\pm 0.0033$	$2612.083\pm 237.065$
	OT-RED (ours)	$\mathbf{15.04\pm 2.08}$	$\mathbf{226.10\pm 7.04}$	$-52207.3\pm 145.4$	$0.0012\pm 0.0003$	$0.3891\pm 0.0074$	$0.9289\pm 0.0031$	$14.608\pm 2.572$
	W-AVR	$234129.9\pm 5247.6$	$1.41\times 10^{6}$	$-1.47\times 10^{6}$	$12.6341\pm 0.8846$	$0.8060\pm 0.0139$	$0.1146\pm 0.0172$	$12.137\pm 1.988$
	A-AVR	$17.38\pm 2.69$	$228.41\pm 7.18$	$\mathbf{-52181.4\pm 155.5}$	$\mathbf{0.0005\pm 0.0001}$	$\mathbf{0.3885\pm 0.0077}$	$\mathbf{0.9291\pm 0.0033}$	$12.166\pm 2.014$
	GM	$15.17\pm 2.08$	$226.27\pm 7.05$	$-52206.8\pm 145.2$	$0.0012\pm 0.0003$	$0.3891\pm 0.0074$	$0.9289\pm 0.0031$	$12.883\pm 2.331$
	MED	$306.61\pm 69.89$	$459.04\pm 31.09$	$-95360.9\pm 15541.4$	$0.2748\pm 0.1260$	$1.3159\pm 0.3571$	$0.6939\pm 0.0860$	$14.839\pm 3.966$

Table 3: Performances of estimators under homogeneous expert capacities (

M=16,N=100k,d=20,K=5

, mean

\pm

std over 50 runs).

Support Condition	$S/N_{m}$	$\mathcal{T}_{c}(\cdot,f^{*})$	$\mathcal{T}_{c}(\cdot,\bar{f}^{W})$	Test Loglik	$\textbf{MSE}(\cdot,\boldsymbol{\theta}^{*})$	RPE	ARI	Agg. Time (s)
$S=30$	$0.006$	$91.35\pm 104.87$	$285.55\pm 87.10$	$-58103.7\pm 13249.7$	$0.0969\pm 0.2457$	$0.5020\pm 0.1527$	$0.9075\pm 0.0311$	$0.433\pm 0.368$
$S=50$	$0.01$	$33.09\pm 19.97$	$240.73\pm 14.66$	$-53335.9\pm 1420.9$	$0.0234\pm 0.0243$	$0.4163\pm 0.0385$	$0.9227\pm 0.0096$	$\mathbf{0.347\pm 0.393}$
$S=100$	$0.02$	$19.64\pm 3.39$	$230.26\pm 9.55$	$-52078.1\pm 385.3$	$0.0073\pm 0.0058$	$0.3809\pm 0.0128$	$0.9299\pm 0.0033$	$0.462\pm 0.553$
$S=250$	$0.05$	$16.48\pm 2.03$	$226.81\pm 9.11$	$-51706.6\pm 130.5$	$0.0023\pm 0.0012$	$0.3693\pm 0.0064$	$0.9319\pm 0.0028$	$0.882\pm 0.781$
$S=500$	$0.1$	$15.77\pm 1.99$	$226.31\pm 9.13$	$-51654.6\pm 126.1$	$0.0016\pm 0.0006$	$0.3675\pm 0.0059$	$0.9323\pm 0.0027$	$1.401\pm 1.025$
$S=1000$	$0.2$	$15.60\pm 1.95$	$226.11\pm 9.11$	$-51635.4\pm 122.0$	$0.0013\pm 0.0003$	$0.3671\pm 0.0060$	$0.9324\pm 0.0028$	$3.204\pm 3.146$

Table 4: Sensitivity of proposed OT-RED to support sample size

S=|\boldsymbol{X}_{S}|

and source of

\boldsymbol{X}_{S}

(

N=100k,M=16,N_{m}=5000,d=20,K=5

, mean

\pm

std over 50 runs).

Estimator	$\mathcal{T}_{c}(\cdot,\widehat{f}_{\mathrm{GLB}})$	$\mathcal{T}_{c}(\cdot,\bar{f}^{W})$	RMSE	MAE	RPE	Corr.	Time (s)
GLB	$0.00$	$151.67$	$17.77$	$9.46$	$0.0464$	$0.9766$	$29.6$
OT-RED (ours)	$\mathbf{75.60}$	$\mathbf{112.12}$	$17.76$	$9.50$	$0.0463$	$0.9766$	$17.6$
A-AVR	$118.84$	$165.71$	$31.29$	$11.24$	$0.1438$	$0.9308$	$34.2$
W-AVR	$1133.92$	$1158.77$	$18.30$	$10.55$	$0.0492$	$0.9752$	$34.2$
GM	$189.32$	$261.06$	$19.39$	$9.85$	$0.0552$	$0.9720$	$34.2$
MED	$173.24$	$153.39$	$34.70$	$14.26$	$0.1769$	$0.9170$	$34.3$

Table 5: Performance on the Beijing air quality sensor network (

M=12

stations,

K=4

Figure 3: Performance on Beijing multi-site air quality data.

Figure 4: Performance on Year prediction MSD data.

Mixture-of-Experts Models

Mixture-of-experts models use input-dependent routing to activate specialized subnetworks, increasing model capacity without running every parameter for each input.

See all 454 papers

Towards a Statistical Understanding of Mixture-of-Experts
Sep 3, 2026Siyuan He, Bokai Yang, Jie Hu +2Mixture of Experts Sparse Mixture-of-Experts
Qiuzhen College, Tsinghua University, Beijing, China. · KLATASDS-MOE, School of Statistics, East China Normal University, Shanghai, China. · Yau Mathematical Sciences Center, Tsinghua University, Beijing, China. +1
Distributionally Robust Mixture-of-Experts Training
Oct 5, 2026Xin Teng, Muxiao Li, Hongyi WenMixture-of-Experts Language Models Sparse Mixture-of-Experts
New York University Center for Data Science, NYU Shanghai
Sparsity Saves Compute, Not Samples: Oracle Inequalities for Dense and Sparse Mixture of Experts
Oct 10, 2026Sagar Ghosh, Huy Nguyen, Alessandro RinaldoSparse Mixture-of-Experts Mixture-of-Experts Models
The University of Texas at Austin

Sep 3, 2026J/K move · Enter open · S save

Towards a Statistical Understanding of Mixture-of-Experts

Siyuan He, Bokai Yang, Jie Hu, Ziwen Gao+1

Qiuzhen College, Tsinghua University, Beijing, China. · KLATASDS-MOE, School of Statistics, East China Normal University, Shanghai, China. · Yau Mathematical Sciences Center, Tsinghua University, Beijing, China.+1

Mixture of Experts Sparse Mixture-of-Experts Mixture-of-Experts Models Statistical Learning Theory

Open paper AskPDF arXiv

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.

Abstract

Mixture-of-Experts Models

Explore similar work

Towards a Statistical Understanding of Mixture-of-Experts

Distributionally Robust Mixture-of-Experts Training

Sparsity Saves Compute, Not Samples: Oracle Inequalities for Dense and Sparse Mixture of Experts