Technical Summary: MetaMoE

Problem Statement

Large foundation models are increasingly fine-tuned on private, domain-specific data across various organizations, resulting in a collection of specialized "expert" models. While unifying these experts into a single Mixture-of-Experts (MoE) model could theoretically provide multi-domain capabilities, this process is hindered by strict privacy constraints that prevent the sharing of raw training data. Existing approaches face significant limitations:

• Federated Learning (FL): Requires costly synchronized optimization and often suffers from performance degradation under heterogeneous data distributions.
• Model Merging (e.g., Model Soup): Directly averages weights, which is fragile when experts diverge significantly in function space.
• Ensembling (e.g., BTM): Combines predictions at inference time but fails to produce a single deployable model for downstream tasks like fine-tuning or RLHF.
• Proxy-Based Unification (e.g., FlexOlmo): Uses public data as proxies to train a router but relies on similarity-based selection, often resulting in redundant, narrowly concentrated proxies that fail to cover the private data manifold. Furthermore, training experts solely on private data creates a "domain isolation" mismatch where experts are never exposed to the proxy data used for routing.

The core challenge is Privacy-Preserving Mixture-of-Experts Unification: How to unify independently trained experts into a deployable MoE model without sharing private data, while ensuring effective coordination between heterogeneous experts.

Methodology: MetaMoE

MetaMoE is a framework that unifies independently trained experts using public proxy data as surrogates for inaccessible private data. The framework operates in three main stages:

1. Diversity-Aware Proxy Selection

To approximate private data distributions without sharing them, MetaMoE selects proxy samples from a public dataset ($D_0$) for each client.

• Relevance-Weighted Determinantal Point Process (DPP): Unlike previous methods that select proxies based solely on similarity (relevance), MetaMoE employs a relevance-weighted DPP. This method augments the standard DPP kernel with client-specific relevance scores ($g(x, D_p)$).
• Mechanism: The selection probability maximizes a log-determinant term that balances two objectives:
  1. Relevance: Favoring samples similar to the client's private domain.
  2. Diversity: Enforcing negative interactions between selected samples to avoid redundancy and ensure broad coverage of the private data manifold.
• Optimization: The selection is performed via efficient greedy MAP inference with Cholesky updates, reducing computational complexity from $O(nm^3)$ to $O(nm)$.

2. Proxy-Aligned Expert Training

To mitigate the mismatch between expert behavior and router supervision, MetaMoE introduces a proxy-aligned training strategy.

• Process: Each client fine-tunes the Feed-Forward Network (FFN) sublayers of their expert model using a combination of their private data ($D_p$) and their client-specific proxy data ($\hat{D}_p$).
• Alignment: By exposing experts to the same proxy data used later for router training, the experts' representations are calibrated to the distribution the router will encounter. This prevents "domain isolation" and enables effective coordination at unification time.

3. Context-Aware Router and Unification

• Router Initialization: The router is initialized with domain-aware vectors derived from the mean embeddings of both private and proxy data.
• Context-Aware Routing: Instead of relying solely on token-level features, the router utilizes a blended representation $\tilde{z}_t = (1-\lambda)z_t + \lambda z_x$, where $z_x$ is a sequence-level embedding capturing global context. This improves expert assignment reliability, especially when tokens have similar surface forms but belong to different domains.
• Final Unification: The FFN sublayers of all experts are merged into MoE layers. The unified model is then jointly fine-tuned on the union of all proxy datasets to refine the router and ensure cohesive expert operation.

Privacy Guarantees

MetaMoE adheres to a strict data residency model. Clients communicate only:

1. Indices of selected public proxy samples (deterministic functions of public data).
2. Final expert weights (FFN sublayers).
3. Routing vectors (mean embeddings).

The paper provides formal privacy analysis showing that the routing vectors have bounded per-sample sensitivity ($O(1/m)$) and that the mean private embedding is unrecoverable without knowledge of the private dataset size ($N$), which is never communicated. This offers strictly stronger privacy guarantees than methods like FlexOlmo, which share mean embeddings derived solely from private data.

Key Contributions

1. Framework Proposal: MetaMoE is proposed as a privacy-preserving framework for unifying independently fine-tuned experts without sharing private data, with formal privacy guarantees.
2. Diversity-Aware Proxy Selection: The introduction of a relevance-weighted DPP addresses the limitations of similarity-only sampling, ensuring proxy datasets are both representative of the client domain and diverse enough to cover the data manifold.
3. Proxy-Aligned Training & Context-Aware Routing: The design of a proxy-aligned expert training strategy and a context-aware router demonstrates that aligning expert training with routing supervision significantly improves coordination and performance.

Experimental Results

MetaMoE was evaluated on Computer Vision (CV) and Natural Language Processing (NLP) benchmarks.

• CV Tasks: Using CLIP ViT-B/32 and ViT-B/16 as seed models across Pets, Flowers, and EuroSAT datasets. MetaMoE achieved average accuracies of 94.52% and 96.24%, respectively. This outperformed the strongest baseline, FlexOlmo, by approximately 1.6 to 2.7 percentage points.
• NLP Tasks: Using LLaMA-3.2-3B and LLaMA-3.1-8B as seed models across CommonsenseQA, CosmosQA, and SocialIQA. MetaMoE achieved average accuracies of 74.42% and 81.59%, consistently outperforming all privacy-preserving baselines (BTM, Model Soup, BTX, FlexOlmo).
• Ablation Studies:
  • Removing the diversity component (using only relevance) significantly degraded performance, confirming the necessity of DPP-based selection.
  • Training experts solely on private data (without proxy alignment) led to poor router coordination.
  • Training experts solely on proxy data (without private data) resulted in substantial accuracy drops, proving that private data is essential for domain expertise while proxies serve primarily for alignment and routing.
  • The method remained robust even when the public dataset had no semantic overlap with client domains.

Significance and Claims

The paper claims that MetaMoE successfully addresses the fundamental trade-off between data privacy and model unification. By leveraging diversity-aware proxy selection and proxy-aligned training, it enables the creation of a single, deployable MoE model that generalizes across heterogeneous client domains without violating data residency constraints.

The authors emphasize that MetaMoE offers a principled alternative to Federated Learning (avoiding synchronization bottlenecks) and simple model merging (preserving expert specialization). The formal privacy analysis demonstrates that the shared artifacts reveal negligible private information, providing a stronger guarantee than existing proxy-based methods. The results suggest that diversity in proxy selection is a critical, previously underutilized factor in privacy-preserving MoE unification.