MedMix: Specialization-Consistent Federated Sparse MoEs under Modality Heterogeneity

2026-08-13

AI Summary

This study addresses the challenges of sparse Mixture-of-Experts (MoE) expert specialization conflicts and aggregation misalignment caused by modality heterogeneity in federated multimodal medical AI. We propose MedMix, a novel framework featuring a modality-context-aware semantic alignment mechanism. By integrating modality-aware routing, consensus anchor alignment, and adaptive expert aggregation, MedMix achieves cross-client routing consensus and precise matching of functionally similar experts. Experiments on real-world medical datasets demonstrate that MedMix attains state-of-the-art average F1 scores, with particularly significant performance gains under severe modality heterogeneity and missing data scenarios. These results confirm its effectiveness in resolving semantic alignment difficulties during federated collaborative training, thereby enhancing robustness in heterogeneous medical environments.

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.

Problem

• Federated Learning
• Multimodal Medical AI
• Modality Heterogeneity
• Sparse Mixture-of-Experts
• Expert Specialization Misalignment

Innovation

• Federated Sparse MoE
• Modality Heterogeneity
• Semantic Alignment
• Consensus-Guided Routing
• Client-Adaptive Aggregation