MedMix: Specialization-Consistent Federated Sparse MoEs under Modality Heterogeneity

Abstract: Federated multimodal medical AI faces modality heterogeneity at both the client and sample levels: clients may systematically lack access to specific modality types, while individual records within the same client may contain different partial modality subsets. Sparse Mixture-of-Experts (MoE) architectures are a promising remedy for modality-adaptive computation, but their use in federated learning is fragile under cross-client modality heterogeneity, where locally learned routing policies can diverge across clients and drive experts toward incompatible specializations. Different clients may assign the same observed modality configuration to different experts, or train similarly indexed experts on different missing-modality configurations, causing standard aggregation to misalign or overwrite the expert specialization that sparse MoEs are intended to learn. To address this challenge, we propose MedMix, a semantic-alignment framework for federated multimodal sparse MoEs that coordinates cross-client routing and expert specialization using modality context. At the client side, MedMix uses modality-context-aware routing to guide expert selection using each token's modality identity, position, and incompleteness context. Across clients, it uses consensus-guided routing alignment to construct server-side consensus anchors for shared modality patterns and align local routing distributions across clients. Complementing these routing mechanisms, client-adaptive expert aggregation leverages client-specific modality-pattern prototypes to match and aggregate functionally similar experts across clients. Experiments on real-world multimodal medical datasets show that MedMix achieves the best average F1 across diverse modality heterogeneity and modality incompleteness settings, with especially clear gains under severe heterogeneity.

• The paper introduces MedMix, which combines modality-context-aware routing, consensus-guided alignment, and client-adaptive expert aggregation to preserve consistent expert specialization across federated medical clients.
• MedMix achieves the best average F1 on ADNI (0.530) and MIMIC-IV (0.567), with its largest advantage under severe heterogeneity, including a 0.039 F1 gain over the strongest ADNI baseline.
• The method reduces cross-client routing divergence while using roughly 2.2–7.9× fewer parameters or FLOPs than dense-expert federated MoE baselines, though privacy leakage and performance on real-world heterogeneity remain open questions.

Problem and motivation

Multimodal medical AI in federated settings faces what the authors call dual-level modality heterogeneity: institutions differ systematically in which modalities they collect (cross-client heterogeneity), while individual patient records within an institution contain only partial modality subsets (intra-client incompleteness). Sparse Mixture-of-Experts (MoE) fusion layers are a natural fit for this setting because they can route tokens from different modality configurations to specialized experts, and centralized methods such as FlexMoE and FuseMoE have demonstrated robustness to arbitrary modality combinations. The paper's central claim is that this robustness does not survive federated aggregation: when clients observe different modality-pattern distributions, locally learned routers diverge, and experts with the same index acquire different functional roles on different clients. Standard FedAvg then averages parameters that are not semantically comparable, degrading or even reversing the benefit of aggregation.

The paper supports this claim with diagnostic experiments on a standard MoE-FL baseline. Under the same observed modality configuration, clients exhibit distinct normalized top-$k$ expert activation profiles; joint t-SNE embeddings of expert-layer outputs show that same-index experts develop client-specific semantic roles; and the validation F1 change induced by aggregation ($$\Delta\mathrm{F1} = \mathrm{F1}_{\text{post}} - \mathrm{F1}_{\text{pre}}$$) is consistently positive under balanced modality splits but unreliable—sometimes negative—under heterogeneous splits. A notable implication of these diagnostics is that the standard load-balancing regularizer is insufficient: balanced expert utilization does not imply cross-client semantic compatibility of expert specialization. This is a somewhat contrarian position relative to common practice in federated MoE work, and the paper's design follows from it directly.

Method

MedMix couples three components, each addressing a distinct failure mode identified in the diagnostics.

Modality-context-aware routing (MCR). Standard routers score experts from token representations alone. MCR instead constructs a context-aware router input by concatenating the token embedding with learned embeddings of the modality identity, the full modality-availability mask, and the token position, passed through a lightweight projection. This lets the router distinguish tokens that are representationally similar but arise from different observation contexts (e.g., a clinical token observed alone versus alongside biospecimen data), producing more transferable expert assignments across clients.

Consensus-guided routing alignment (CRA). After each round, clients summarize mean routing distributions per layer and per modality pattern $p = (\mathbf{r}, m)$ over validation tokens. The server aggregates these into consensus anchors, weighting client summaries by pattern support and routing confidence (negative entropy), and retains anchors only for patterns with sufficient cross-client support. Clients then minimize a KL-divergence loss aligning their batch-level routing distributions to the broadcast anchors, with the loss zero when no anchored pattern is present in the batch. To increase anchor exposure under skewed local distributions, an anchor-exposure augmentation samples feasible submasks of observed modalities proportionally to the global mask distribution, exposing clients to globally supported subpatterns without fabricating unobserved modalities.

Client-adaptive expert aggregation (CEA). Rather than averaging same-index experts, CEA performs function-aware matching. Each client computes pattern-level prototypes of its expert-input distributions (without sharing raw data); the server scores every source expert by its functional response on the target client's prototypes, weighted by pattern support and routing probability, converts scores into aggregation weights via temperature-scaled softmax, and forms personalized experts as a residual-weighted blend of the target client's own expert and matched source experts. Non-expert components are aggregated globally with sample-weighted FedAvg.

The local objective combines task loss, load balancing, and CRA alignment. The framework is privacy-preserving in the sense that only routing summaries and aggregate prototypes are exchanged; the paper does not, however, analyze the privacy leakage of these statistics, which is an assumption worth noting for deployments under strict regulatory constraints.

Experimental results

The evaluation covers ADNI (four modalities: MRI, genomics, clinical assessments, biospecimen biomarkers; three-class Alzheimer's diagnosis) and MIMIC-IV (labs, clinical notes, diagnosis codes; one-year mortality prediction). Federated splits are constructed with a Beta–Bernoulli process for client-level structural masking, plus record-level dropout ($p_{\text{intra}} = 0.3$) and an extreme modality-exclusive setting. All splits are monotone with respect to the natural dataset mask—modalities are only removed, never fabricated. Baselines span centralized fusion architectures adapted to FL (Attention, Transformer, FuseMoE, FlexMoE), FL-native incomplete-modality methods (FedDUET, PEPSY), and federated sparse MoE methods (FedMoE, FedAlign-MoE).

Dataset	Method	Avg. F1	High heterogeneity	Mod-Excl.
ADNI	MedMix	0.530	0.546	0.481
ADNI	FlexMoE (best baseline)	0.513	0.507	0.471
MIMIC-IV	MedMix	0.567	0.571	0.505
MIMIC-IV	FuseMoE / FedAlign-MoE	0.564	0.566 / 0.562	0.497 / 0.504

MedMix achieves the best average F1 on both datasets and ranks top-2 in every ADNI setting. Gains are concentrated in severe heterogeneity: on ADNI, the margin over the strongest baseline is 0.039 F1 under High inter-client heterogeneity, versus 0.017 under Balanced. On MIMIC-IV the improvements are smaller and MedMix is not uniformly best per setting, which the authors attribute to the larger per-client sample count and fewer modalities allowing baselines to learn more stable specialization. This asymmetry is honestly reported and supports the paper's thesis that its mechanisms matter most when local modality distributions diverge sharply.

The mechanistic analyses corroborate the design rationale. MCR reduces the support-weighted Jensen–Shannon divergence between clients' routing distributions over shared patterns by more than an order of magnitude, and CRA reduces it further. CEA keeps prototype-weighted cosine drift between a client's local expert and its received aggregated expert substantially lower than index-wise global aggregation throughout training. The system-cost comparison is favorable: MedMix attains the highest average F1 with roughly $2.2\times$ fewer parameters and $4.4\times$ fewer FLOPs than the dense-expert federated MoE baselines on ADNI ($2.8\times$ and $7.9\times$ on MIMIC-IV), and is comparable in cost to FlexMoE—suggesting the gains stem from specialization consistency rather than added capacity. The ablation shows each component contributes: MCR alone improves Mod-Excl. from 0.446 to 0.467, CRA adds gains under intra-client missingness, and CEA lifts Mod-Excl. to 0.481, though the total average improvement over the base variant (0.523 to 0.530) is modest relative to the per-setting variance.

Limitations and open questions

The paper concedes several constraints. The heterogeneity scenarios are synthetic constructions layered on natural missingness and may not capture real deployments where modality availability correlates with label distributions and site-specific clinical factors. Evaluation is restricted to a single shared prediction task; multi-task and task-heterogeneous federated settings remain untested. The CRA mechanism depends on validation-token pattern support thresholds and a quorum requirement, so behavior under extremely sparse or highly skewed pattern support is not characterized. Privacy analysis of the exchanged routing summaries and expert-input prototypes is not provided. Finally, the ablation gains are small in absolute terms on some settings, leaving open how much of the benefit under moderate heterogeneity is attributable to each component versus the overall architecture.

Conclusion

MedMix identifies and addresses a concrete failure mode of federated sparse MoEs: client-specific modality distributions induce divergent routing policies and index-incompatible expert specializations that standard aggregation cannot reconcile. Its three mechanisms—context-aware routing, consensus-anchored routing alignment, and function-matched personalized expert aggregation—yield the best average F1 on ADNI and MIMIC-IV across a spectrum of modality heterogeneity settings, with the largest gains under severe heterogeneity, at a computational cost comparable to or below competing federated MoE baselines.