MedMix: Specialization-Consistent Federated Sparse MoEs under Modality Heterogeneity

Adiba Orzikulova; Dong Min Kim; Jaehong Yoon; Sung-Ju Lee

cs.LGAug 14, 2026

MedMix: Specialization-Consistent Federated Sparse MoEs under Modality Heterogeneity

Authors: Adiba Orzikulova, Dong Min Kim, Jaehong Yoon, Sung-Ju Lee

Organizations: KAIST · NTU Singapore

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.

Figures & tables

Figure 1: Fragility of federated MoE specialization under modality heterogeneity. (a) Routing divergence: normalized top- $k$ expert activation frequencies. Clients show divergent routing preferences under the same observed modality configuration. (b) Expert misalignment: joint t-SNE of expert-layer outputs from two clients, with points colored by the top-1 routed expert. Rearranged color and cluster structure indicates that same-index experts acquire different client-specific semantic roles. (c) Aggregation impact: validation F1 change after aggregation relative to the pre-aggregation local model, $\Delta\mathrm{F1}=\mathrm{F1}_{\mathrm{post}}-\mathrm{F1}_{\mathrm{pre}}$ , averaged over communication rounds and 5 seeds. Aggregation is beneficial in the balanced split but becomes less reliable under heterogeneous modality distributions, reflecting reduced cross-client specialization consistency.

Table 1: Performance comparison under modality heterogeneity. Results are F1 scores (macro-F1 for ADNI and binary F1 for MIMIC-IV; mean

\pm

std over five seeds). Balanced/Medium/High denote increasing inter-client modality heterogeneity; + Intra-Client Missingness adds record-level modality dropout; Mod-Excl. denotes modality-exclusive clients.

Panel of Figure 2 — Figure 2: Cross-client specialization consistency analysis. (a) Each small dot denotes one shared modality-mask pattern, and the large dot denotes the support-weighted mean JS divergence; MCR substantially reduces routing disagreement, while CRA further improves cross-client routing consistency. (b) Index-wise global aggregation without CEA causes larger cosine drift from the target client’s local expert, whereas CEA maintains substantially lower drift throughout training, indicating improved semantic compatibility of expert specialization.

ADNI
Metric	Attention	Transformer	FuseMoE	FlexMoE	FedDUET	PEPSY	FedMoE	FedAlign-MoE	MedMix
Params (M) $\downarrow$	3.79	4.17	10.84	5.02	4.65	4.46	10.84	10.84	4.90
MFLOPs $\downarrow$	55.8	90.3	324.7	63.5	53.8	63.3	324.7	324.7	73.9
Time (ms) $\downarrow$	0.82	0.94	1.85	2.70	0.72	1.51	1.92	1.91	2.37
Avg. F1 $\uparrow$	.453	.473	.485	.513	.473	.424	.469	.464	.530
MIMIC-IV

Table 2: System cost comparison. Params denote trainable per-client inference parameters. FLOPs and time are measured per sample on a balanced test split. Avg. F1 reports the average performance across all evaluated heterogeneity settings.

Variant	Balanced	Medium	High	Mod-Excl.	+Intra Bal.	+Intra Med.	+Intra High	Avg.
w/o CEA,CRA,MCR	.567	.527	.535	.446	.541	.513	.533	.523
w/o CEA, CRA	.571	.521	.524	.467	.544	.508	.543	.523
w/o CEA	.554	.524	.544	.467	.548	.514	.549	.529
MedMix Complete	.560	.521	.546	.481	.540	.519	.540	.530

Table 3: Ablation study on ADNI. MCR denotes modality-context-aware routing, CRA denotes consensus-guided routing alignment, and CEA denotes client-adaptive expert aggregation.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

\ell(f_{\theta}(\mathbf{x},\tilde{\mathbf{r}}),y)+\lambda_{\mathrm{bal}}\mathcal{L}_{\mathrm{bal}}+\lambda_{\mathrm{CRA}}\mathcal{L}_{\mathrm{CRA}}.

Appendix

Algorithm 2 ClientUpdate for MedMix

Dataset	$M$	Setting	Beta $(\alpha,\beta)$	$\mathbb{E}[p_{a}]$	$\mathbb{E}[\#\mathrm{avail}]$
ADNI	4	Balanced	—	—	$4.00$
		Medium	$(17,3)$	$0.85$	$3.40$
		High	$(6,14)$	$0.30$	$1.20$
MIMIC-IV	3	Balanced	—	—	$3.00$
		Medium	$(15,5)$	$0.75$	$2.25$
		High	$(6,14)$	$0.30$	$0.90$

Appendix

Table 4: Inter-client modality-availability settings. Balanced uses the natural stratified split without additional structural modality drops. Medium and High use Beta–Bernoulli client-level structural masking. The +Intra-Client Missingness variants additionally use

p_{\mathrm{intra}}=0.3

Dataset	Method	Learning rate	Batch size
ADNI	Attention [ 1 ]	$4\times 10^{-5}$	8
	Transformer [ 25 ]	$1\times 10^{-4}$	8
	FuseMoE [ 8 ]	$1\times 10^{-4}$	16
	FlexMoE [ 30 ]	$1\times 10^{-4}$	8
	FedDUET [ 21 ]	$1\times 10^{-4}$	16
	PEPSY [ 19 ]	$4\times 10^{-5}$	8

Appendix

Table 5: Federated baseline training hyperparameters used in our experiments.

Dataset	Experts	Top- $k$	Heads	Layers	Classifier hidden/layers	LR	Batch size
ADNI	16	2	4	1	256 / 2	$1\times 10^{-4}$	8
MIMIC-IV	32	4	4	1	320 / 3	$4\times 10^{-5}$	16

Appendix

Table 6: MedMix model and training hyperparameters.

Figure 3: Additional routing-divergence examples. Normalized top-2 expert activation frequencies for four different observed modality patterns. Here, C, B, and G denote clinical, biospecimen, and genomic modalities; for example, CB / C refers to the clinical modality token when only clinical and biospecimen modalities are observed for the sample. Darker green indicates more frequent expert selection. When routing the same modality pattern, clients activate different expert subsets, extending Figure 1 (a).

Figure 4: Expert misalignment across clients. t-SNE visualizations of final-round expert outputs for all clients in the High and Mod-Excl. heterogeneity settings. Each row uses a shared embedding across clients, and colors denote top-1 routed experts. Structure varies across clients, extending Figure 1 (b).

Figure 5: Aggregation impact over training. Round-wise aggregation impact,

\Delta\mathrm{F1}=\mathrm{F1}_{\mathrm{post}}-\mathrm{F1}_{\mathrm{pre}}

, for four different heterogeneity settings. Lines show means over clients and seeds. Shaded regions denote standard error over seeds. Negative values indicate that aggregation hurts the just-trained local model.

Explore similar work

ProMoE-FL: Prototype-conditioned Mixture of Experts for Multimodal Federated Learning with Missing Modalities
Jul 7, 2026Aavash Chhetri, Bibek Niroula, Eduard Vazquez +4Multimodal Federated Learning Missing-Modality Learning
NepAl Applied Mathematics and Informatics Institute for research, Nepal · Fogsphere (Redev.AI Ltd), UK · University of Lausanne, Switzerland +4
FM $^2$ : Unified Federated Foundation Models for Heterogeneous Multimodal Medical Imaging
Jul 15, 2026Shengchao Chen, Ting ShuMultimodal Federated Learning Cross-Modal Alignment
School of Artificial Intelligence, Shenzhen University Shenzhen, China · Australian AI Institute, University of Technology Sydney Sydney, Australia
Prototype-guided Bilateral Alignment Multimodal Federated Learning
Sep 30, 2026Tianchi Liao Tianchi_Liao, Lele Fu, Sheng Huang +3Multimodal Federated Learning Federated Learning Aggregation
Department of Computer Science, Hong Kong Baptist University, Hong Kong SAR, China · School of Computer Science and Engineering, Sun Yat-Sen University, Guangzhou, China

Jul 7, 2026, cs.CVJ/K move · Enter open · S save

ProMoE-FL: Prototype-conditioned Mixture of Experts for Multimodal Federated Learning with Missing Modalities

Aavash Chhetri, Bibek Niroula, Eduard Vazquez, Yash Raj Shrestha+3

NepAl Applied Mathematics and Informatics Institute for research, Nepal · Fogsphere (Redev.AI Ltd), UK · University of Lausanne, Switzerland+4

Multimodal Federated Learning Missing-Modality Learning Medical Imaging Prototype Learning Mixture-of-Experts Models Federated Learning

Open paper AskPDF arXiv

In this paper, we address the problem of multimodal federated learning with missing modality. Existing methods utilize an additional public dataset or perform naive feature synthesis that is based solely on the available modality. To address these limitations, we propose ProMoE-FL, a Prototype-conditioned Mixture-of-Experts framework for robust missing-modality feature synthesis in multimodal federated learning. ProMoE-FL builds a global client-aware prototype bank that captures clinically meaningful modality priors across institutions. Our Mixture of Experts is conditioned on these prototypes and modality indices to enable direction-aware expert routing for dynamically synthesizing missing features. We perform extensive quantitative and qualitative evaluations on four public chest X-ray datasets (MIMIC-CXR, NIH Open-I, PadChest, and CheXpert) and demonstrate that ProMoE-FL consistently outperforms state-of-the-art methods in both homogeneous as well as the more challenging heterogeneous settings.

Abstract

Explore similar work

ProMoE-FL: Prototype-conditioned Mixture of Experts for Multimodal Federated Learning with Missing Modalities

FM2^22: Unified Federated Foundation Models for Heterogeneous Multimodal Medical Imaging

Prototype-guided Bilateral Alignment Multimodal Federated Learning

FM $^2$ : Unified Federated Foundation Models for Heterogeneous Multimodal Medical Imaging