OT-RED Merges Distributed AI Models Across Hospitals With One Upload
A one-shot optimal transport method aggregates independently trained mixture-of-experts models into a single global estimator without pooling raw data.
- New optimal transport method aggregates independently trained mixture-of-experts models in one communication round.
- Solves label switching, covariate-dependent gating, and heterogeneous local expert counts that break naive averaging.
- Server runs a majorization-minimization loop plus IRLS on a small support sample of inputs.
- Matches centralized training accuracy on Beijing air quality and YearPredictionMSD benchmarks with lower compute.
- Statistical consistency of the aggregated estimator proven under standard assumptions on local estimators.
- MATLAB code with all baselines available on GitHub under MIT license.
Merging Distributed Mixture-of-Experts Models in One Upload
Distributed mixture-of-experts training becomes difficult when records are split across hospitals, sensors, or regulated data stores. In an arXiv paper, Faïfcel Chamroukhi and Nhat Thien Pham propose OT-RED, an aggregation method that combines independently trained MoE models after a single upload from each worker. Training records and outcomes can remain local while the server builds a global approximation to a model trained on pooled data.
The work covers classical statistical MoE models, where several regression experts model different regions of the input space and a softmax gating function assigns each input to those experts. For input x and outcome y, the conditional model has the form p(y|x) = Σ_k π_k(x)f_k(y|x). Each π_k(x) is an input-dependent gate, and each f_k is an expert distribution or regression model.
| Property | OT-RED |
|---|---|
| Worker task | Fit an independent local MoE |
| Communication | One parameter upload per worker |
| Server inputs | Local model parameters, support covariates, and target expert count K |
| Server output | One global MoE with aligned experts and a fitted gate |
| Local capacities | May vary across workers, so K_m can differ from K |
| Statistical result | Consistency under the paper’s stated regularity assumptions |
| Privacy | Reduced data movement and communication, without a formal privacy guarantee |
Why Averages Scramble the Model
Coordinate-wise parameter averaging assumes that matching coordinates describe matching components. Mixture models violate that assumption because their likelihood is unchanged when component labels are permuted. Two workers can learn equivalent structures while assigning different numbers to corresponding experts.
- Label switching: expert 1 on one machine may correspond to expert 3 on another, causing an average to combine unrelated regressions.
- Input-dependent gates: mixture weights vary with
x, so methods designed for unconditional mixtures lose information about where each expert applies.
This story is for Pro members
You've reached the end of the free preview. Upgrade to AlphaSignal Pro to read the full article - and everything else behind the paywall.
Comments
