Optimal Transport Aggregation for Distributed Mixture-of-Experts

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.