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

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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- kk 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, ΔF1=F1post−F1pre\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.
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.
MetricAttentionTransformerFuseMoEFlexMoEFedDUETPEPSYFedMoEFedAlign-MoEMedMix
ADNI
Params (M) ↓\downarrow3.794.1710.845.024.654.4610.8410.844.90
MFLOPs ↓\downarrow55.890.3324.763.553.863.3324.7324.773.9
Time (ms) ↓\downarrow0.820.941.852.700.721.511.921.912.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.
VariantBalancedMediumHighMod-Excl.+Intra Bal.+Intra Med.+Intra HighAvg.
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

 ℓ(fθ(x,r~),y)+λbalLbal+λCRALCRA.\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
DatasetMMSettingBeta (α,β)(\alpha,\beta)E[pa]\mathbb{E}[p_{a}]E[#avail]\mathbb{E}[\#\mathrm{avail}]
ADNI4Balanced——4.004.00
Medium(17,3)(17,3)0.850.853.403.40
High(6,14)(6,14)0.300.301.201.20
MIMIC-IV3Balanced——3.003.00
Medium(15,5)(15,5)0.750.752.252.25
High(6,14)(6,14)0.300.300.900.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 pintra=0.3p_{\mathrm{intra}}=0.3 .
DatasetMethodLearning rateBatch size
ADNIAttention [ 1 ]4×10−54\times 10^{-5}8
Transformer [ 25 ]1×10−41\times 10^{-4}8
FuseMoE [ 8 ]1×10−41\times 10^{-4}16
FlexMoE [ 30 ]1×10−41\times 10^{-4}8
FedDUET [ 21 ]1×10−41\times 10^{-4}16
PEPSY [ 19 ]4×10−54\times 10^{-5}8
Appendix
Table 5: Federated baseline training hyperparameters used in our experiments.
DatasetExpertsTop- kkHeadsLayersClassifier hidden/layersLRBatch size
ADNI16241256 / 21×10−41\times 10^{-4}8
MIMIC-IV32441320 / 34×10−54\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, ΔF1=F1post−F1pre\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

CardsList
Jul 7, 2026cs.CV

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

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.

Aavash Chhetri, Bibek Niroula, Eduard Vazquez +4

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

Multimodal Federated LearningMissing-Modality LearningMedical ImagingPrototype LearningMixture-of-Experts ModelsFederated Learning
AskPDF
Jul 15, 2026cs.CV

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

Building foundation models for medical imaging requires pooling data across institutions, yet privacy regulations prohibit centralized aggregation. Existing Federated Foundation Models either fine-tune natural-image models with poor medical-domain transfer, or train from scratch within a single modality, lacking the flexibility to unify tasks. We identify an under-explored challenge, Imaging Modality Heterogeneity, where clients operate under two structural regimes: Overlapped (shared modalities with heterogeneous label distributions) and Non-overlapped (fully disjoint modalities per client). We propose FM2^2, a unified framework that trains the core backbone from scratch to preserve medical domain fidelity while optionally incorporating biomedical pretrained encoders for vision-language alignment. FM2^2 equips each client with dual Mixture-of-Experts modules (a Class-wise MoE for personalized category knowledge and a Domain-wise MoE for shared cross-modality representations), coupled with a Heterogeneous Modality Alignment (HMA) regularizer that explicitly aligns modality-specific expert parameters, admitting provable O(1/T)O(1/\sqrt{T}) convergence and generalization guarantees. FM2^2 further incorporates Caption-Enhanced Learning (CEL), where locally retained GPT-4o-generated captions serve as a textual semantic bridge enabling representation transfer across clients with disjoint modalities, and demonstrates extensibility to Federated Medical VQA. Experiments on our MIMH benchmark (classification and CEL) and real-world medical VQA datasets confirm consistent superiority over state-of-the-art federated baselines and strong out-of-modality generalization across all three tasks.

Shengchao Chen, Ting Shu

School of Artificial Intelligence, Shenzhen University Shenzhen, China · Australian AI Institute, University of Technology Sydney Sydney, Australia

Multimodal Federated LearningCross-Modal AlignmentMedical VQAMedical Image ClassificationMultimodal Foundation ModelsFederated Learning+1
AskPDF
Sep 30, 2026cs.AI

Prototype-guided Bilateral Alignment Multimodal Federated Learning

Multimodal federated learning (MFL) has emerged as a pivotal paradigm for leveraging distributed data to enhance model performance. However, existing methods predominantly rely on idealized assumptions of model homogeneity and balanced modality distributions, rendering them ill-suited for practical scenarios characterized by heterogeneous client architectures and severe modality imbalance. To address these challenges, we propose a \textbf{M}ultimodal \textbf{Fed}erated learning Prototype-guided Bilateral Alignment (MFedPBA) framework. MFedPBA facilitates robust knowledge synergy through a dual alignment mechanism: (i) at the feature level, it aligns heterogeneous feature spaces via a projection encoder optimized by contrastive learning and the Gromov-Wasserstein distance; (ii) at the decision level, it employs an entropy-weighted aggregation of naturally aligned logit prototypes. This novel design achieves robust MFL by jointly tackling heterogeneous feature spaces and collectively aggregating decisions. Extensive experiments demonstrate that our method significantly outperforms state-of-the-art baselines under conditions of model heterogeneity and modality imbalance.

Tianchi Liao Tianchi_Liao, Lele Fu, Sheng Huang +3

Department of Computer Science, Hong Kong Baptist University, Hong Kong SAR, China · School of Computer Science and Engineering, Sun Yat-Sen University, Guangzhou, China

Multimodal Federated LearningFederated Learning AggregationHeterogeneous Federated Learning
AskPDF