PaperlayerView exact revision on arXiv

Source-linked AI summary

MedMix: Specialization-Consistent Federated Sparse MoEs under Modality Heterogeneity

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

arXiv:2608.13911v1cs.LG

TL;DR

Federated multimodal sparse MoEs struggle when clients and samples have different missing-modality patterns, which can misalign routing and expert specialization. MedMix aligns routing and aggregates functionally similar experts using modality context, achieving the best average F1 across diverse heterogeneity settings on ADNI and MIMIC-IV, with larger gains under severe heterogeneity.

  • Problem

    Federated multimodal learning must handle cross-client differences in modality availability and within-client variation in patient-level modality incompleteness.

  • Method

    MedMix uses modality-context-aware routing, consensus-guided routing alignment, and client-adaptive expert aggregation to coordinate routing and specialization across clients.

  • Results

    MedMix achieves the best average F1 on both ADNI and MIMIC-IV across diverse modality heterogeneity settings, including an average binary F1 of 0.567 on MIMIC-IV.

  • Takeaways & Limitations

    MedMix improves robustness across diverse modality heterogeneity and modality incompleteness settings in federated multimodal MoEs.

  • Takeaways & Limitations

    The constructed heterogeneity scenarios may not capture hospital deployments where modality availability correlates with label distributions and other site-specific clinical factors.

Abstract

from arXiv · show

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.

1 Introduction

MedMix addresses the dual-level modality heterogeneity that makes federated multimodal sparse MoEs difficult to align and aggregate. It combines modality-aware routing, consensus alignment, and adaptive expert aggregation, and improves performance across heterogeneous settings on ADNI and MIMIC-IV.

  • Motivation: Privacy-sensitive medical records remain institutionally siloed, while clinical records exhibit pervasive modality incompleteness, motivating federated learning for multimodal collaboration.These constraints create practical barriers to centralizing raw patient records.
  • Problem Setting: Dual-level modality heterogeneity combines cross-client differences in modality availability with within-client variation in patient-level incompleteness.This captures both institution-level and sample-level missing-modality variation.
  • MedMix: MedMix coordinates cross-client routing and expert specialization through modality-context-aware routing, consensus-guided routing alignment, and client-adaptive expert aggregation.Routing uses modality identity, position, and incompleteness context; alignment constructs consensus anchors for shared modality patterns; aggregation matches experts using modality-pattern prototypes.
  • Challenge: Under dual-level heterogeneity, client-specific modality distributions can induce inconsistent routing and expert specialization, weakening semantic compatibility during federated aggregation.Sparse MoEs are appealing because they allocate inputs to specialized experts, but their federated interaction is fragile.
  • Evaluation: MedMix is evaluated on ADNI and MIMIC-IV across diverse modality heterogeneity scenarios, showing consistent improvements over strong federated multimodal and MoE baselines, especially under severe heterogeneity.ADNI integrates MRI, genomic, clinical, and biospecimen modalities, while MIMIC-IV integrates laboratory measurements, clinical notes, and diagnosis codes.

2 Background

This section defines multimodal federated learning with incomplete modality observations and introduces sparse MoEs as modality-adaptive, capacity-efficient models. It also shows that modality heterogeneity makes federated MoE routing and expert specialization inconsistent across clients, motivating specialization-consistent coordination.

  • Multimodal Federated Learning: Multimodal federated learning involves K clients, M possible modalities, labels, and per-sample modality-availability masks that define each observed modality set.The learning objective is a global model under modality incompleteness.
  • Sparse Mixture of Experts: Sparse MoEs increase capacity by activating only a subset of parameters, while multimodal variants support incomplete inputs and arbitrary observed-modality combinations.Each observed modality is encoded separately before sparse routing.
  • Sparse Mixture of Experts: The MoE layer selects the top-s experts using routing probabilities and combines their outputs with routing weights, while load balancing prevents expert collapse.The load-balancing regularizer encourages balanced expert utilization during local training.
  • Fragility of MoEs in FL under Modality Heterogeneity: Dual-level modality heterogeneity breaks cross-client semantic comparability because clients can learn different routing policies and assign different roles to same-index experts.This fragility directly challenges standard federated aggregation of sparse MoE components.
  • Fragility of MoEs in FL under Modality Heterogeneity: Clients exhibit distinct expert activation profiles for shared modality sets, demonstrating disrupted routing consistency under standard cross-client aggregation assumptions.The resulting fragility motivates mechanisms that preserve specialization consistency during routing and aggregation.

3 MedMix

MedMix coordinates modality-aware routing, cross-client routing consistency, and functionally aligned expert aggregation for federated sparse MoEs under modality heterogeneity. Its components use modality context, consensus anchors, and client-specific prototypes to preserve specialization without sharing raw data.

  • Overview: MedMix combines modality-context-aware routing, consensus-guided routing alignment, and client-adaptive expert aggregation to coordinate sparse-MoE specialization across heterogeneous clients.The three components guide expert selection, align routing distributions, and match functionally similar experts using modality-pattern information.
  • Modality-context-aware routing: Modality-context-aware routing injects modality identity, availability, and token position so similar tokens from different observation contexts receive more semantically consistent expert assignments.This addresses token-only routing’s inability to distinguish modality-availability patterns.
  • Consensus-guided routing alignment: Consensus-guided routing alignment aggregates reliable client routing summaries into anchors for shared modality patterns and aligns local batch-level routing only when anchored patterns are observable.Anchors require sufficient cross-client support, and the alignment loss is zero when no anchored pattern appears in a batch.
  • Consensus-guided routing alignment: Anchor-exposure augmentation maps observed samples to feasible, globally supported subpatterns, increasing opportunities for consensus-guided alignment without introducing unobserved modalities.The augmented mask is sampled according to the empirical global mask distribution when feasible.
  • Client-adaptive expert aggregation: Client-adaptive expert aggregation compares experts by functional responses on client-specific modality-pattern prototypes, allowing differently indexed experts to be matched and personalized.A residual contribution preserves the target client’s original expert specialization, while non-expert components are aggregated globally.
  • Federated training procedure: Together, MedMix trains with context-aware routing, consensus targets, load balancing, and personalized experts, promoting specialization-consistent routing and aggregation without sharing raw data.After local training, clients send model updates and pattern-level routing summaries for server-side aggregation and next-round anchor construction.

4 Experiments

MedMix is evaluated on two real-world multimodal medical datasets under federated modality heterogeneity, using diverse centralized-fusion and federated incomplete-modality baselines. Experiments follow a standardized federated training and validation protocol with specified MedMix hyperparameters.

  • Datasets: MedMix is evaluated on ADNI for three-class Alzheimer’s diagnosis and MIMIC-IV for one-year mortality prediction, with up to four and three modalities per sample, respectively.ADNI includes MRI, genomics, clinical assessments, and biospecimen biomarkers; MIMIC-IV includes laboratory, discharge-note, and diagnosis/procedure-code features.
  • Baselines: Baselines include centralized multimodal architectures adapted to federated learning and federated methods designed for inter-client heterogeneity and intra-client incompleteness.The listed baselines include Attention, Transformer, FuseMoE, FlexMoE, FedDUET, and PEPSY.
  • Implementation and evaluation protocol: All methods use the same PyTorch federated simulation with up to 150 communication rounds, one local epoch per round, full participation, and Adam optimization.Each client uses a 70/15/15 train/validation/test split; checkpoints are selected by validation F1 with smoothed early stopping.
  • Implementation and evaluation protocol: MedMix uses one sparse MoE fusion layer, gate load-balancing weight 0.01, CRA weight 0.02, anchor-exposure probability 0.3, and CEA residual weight 0.5.Experiments were conducted on NVIDIA GeForce RTX 3090 GPUs, with additional implementation details in Appendix D.

5 Results

MedMix delivers the strongest average performance across datasets and modality-heterogeneity settings, with especially pronounced gains in challenging regimes. Routing alignment and client-adaptive expert aggregation improve cross-client specialization consistency while maintaining a moderate computational footprint.

  • Primary Results: MedMix achieves the strongest average performance on both datasets, including ADNI macro-F1 and MIMIC-IV binary F1 of 0.567.On ADNI, it ranks top-2 in every setting and gains 0.546 vs. 0.507 in High inter-client heterogeneity and 0.481 vs. 0.471 in Mod-Excl., outperforming FlexMoE.
  • Primary Results: Centralized multimodal MoE models and federated MoE baselines become less consistent as client modality distributions diverge, whereas MedMix promotes specialization consistency.Incomplete-modality methods address incompleteness but lack sparse expert specialization mechanisms for highly heterogeneous modality contexts.
  • Cross-Client Routing Alignment and Expert Compatibility: MCR reduces cross-client routing disagreement, while CRA further improves routing consistency across shared modality-mask patterns.The analysis uses support-weighted Jensen–Shannon divergence between clients’ routing distributions.
  • Cross-Client Routing Alignment and Expert Compatibility: CEA keeps personalized experts closer to target clients’ local experts than index-wise aggregation, which produces increasing cosine drift.This indicates that same-index experts are not necessarily semantically compatible across clients.
  • System Cost: MedMix achieves the highest Avg. F1 on both datasets with a moderate computational footprint.Compared with FuseMoE, FedMoE, and FedAlign-MoE, it uses about 2.2× fewer parameters and 4.4× fewer FLOPs on ADNI.
  • Ablation Study: MedMix improves ADNI average F1 from 0.523 without its components to 0.530 with all components.MCR alone improves Mod-Excl. from 0.446 to 0.467, especially under strongly divergent modality patterns.

6 Discussion and Conclusion

The discussion identifies divergent routing and incompatible expert specialization as key fragilities of federated multimodal MoEs under modality heterogeneity. MedMix addresses these issues through context-aware routing, consensus-guided alignment, and client-adaptive aggregation based on functional similarity, while noting limits to the evaluation scenarios.

  • Limitations and societal impacts: The evaluation constructs inter-client heterogeneity and intra-client incompleteness scenarios to systematically study modality heterogeneity alongside natural modality incompleteness in real-world medical datasets.
  • Limitations and societal impacts: These constructed scenarios may not capture full hospital-deployment complexity, where modality availability can correlate with label distributions and other site-specific clinical factors.
  • Conclusion: Federated multimodal MoEs are fragile under modality heterogeneity because client-specific modality distributions can induce divergent routing and weaken semantic compatibility across expert specializations.
  • Conclusion: MedMix promotes specialization consistency through context-aware routing, consensus-guided routing alignment, and client-adaptive expert aggregation based on functional similarity.

A Algorithm

MedMix federated training alternates server coordination with client-side modality-aware sparse-MoE updates. It aligns routing through consensus anchors and matches experts using pattern-level routing summaries and expert-input prototypes.

  • Federated training: Each round selects clients, sends shared parameters and client expert states, aggregates updates, and constructs CRA anchors from routing summaries.The server returns shared parameters, CRA anchors, and client-specific CEA expert states.
  • Client update: Clients combine shared parameters with local expert states, expose sampled modality patterns, apply MCR routing, and optimize task, load-balancing, and CRA losses.Anchor-exposure augmentation samples ˜r ∼ Aρ(·|r) before modality encoding and routing.
  • Client update: After local training, clients collect pattern-level routing summaries and expert-input prototypes, then return shared parameters and local expert parameters.Routing summaries and prototypes are collected for each modality pattern and support server-side alignment and expert matching.
  • Routing alignment: Routing summaries record support, confidence, and mean routing distributions for layer-level modality patterns p = (r, m).These summaries provide the pattern-level information used to construct CRA anchors.
  • Expert aggregation: CEA uses aggregate expert-input representations for the same modality patterns to perform functional expert matching across clients.The prototypes are collected locally and used by CEA during expert aggregation.

B Related Work

Prior work addresses missing modalities, sparse MoE fusion, and federated multimodal learning, but these lines of research leave unresolved the interaction between modality heterogeneity and federated MoE specialization. MedMix targets this gap in end-to-end multimodal medical prediction with incomplete modalities and structured availability patterns across clients and samples.

  • Multimodal Learning with Missing Modalities: HeMIS, modality-specific teachers, and SMIL handle incomplete multimodal data through shared latent representations, teacher learning, or Bayesian meta-learning without relying on noisy imputation.These approaches motivate treating missing modalities as part of model design rather than preprocessing.
  • Research Gap: MedMix extends this literature to sparse MoE-based multimodal medical prediction, where incomplete modality patterns shape routing and expert specialization across federated clients.Its setting combines sample-level incompleteness with systematic cross-client modality heterogeneity, creating additional specialization failures before aggregation.
  • Sparse Mixture-of-Experts for Multimodal Fusion: Sparse MoE models provide conditional multimodal fusion by activating only a small subset of parameters, while prior work emphasizes routing stability and expert load balancing.Sparsely-gated MoE uses trainable gating to select experts per example, and Switch Transformer simplifies sparse routing.
  • Federated Multimodal Learning with Incomplete Modalities: Federated multimodal learning must address both inter-client modality differences and intra-client missingness, beyond standard label or feature distribution shifts.Existing approaches primarily focus on representation alignment, reconstruction, uncertainty modeling, or personalization in non-MoE fusion architectures.
  • Federated Mixture-of-Experts: Federated MoE methods use specialized ensembles, client-specific sub-MoEs, routing-distribution alignment, or semantic-aware expert aggregation to address client heterogeneity.FedMix emphasizes client-level ensemble selection, while FedMoE and FedAlign-MoE address modular expert exchange and expert semantic mismatch.

C Federated Modality-Missingness Construction

The evaluation constructs modality-heterogeneity splits along inter-client structural availability and intra-client record-level missingness, while respecting each dataset’s natural modality mask. It also includes a modality-exclusive stress test with disjoint modality patterns across clients.

  • Construction: Split construction varies inter-client modality heterogeneity and intra-client missingness, removing only modalities absent from the dataset-provided natural mask.Records with no remaining observed modality are removed.
  • Inter-client heterogeneity: Balanced, Medium, and High settings represent increasing inter-client heterogeneity through stratified partitioning or Beta–Bernoulli client-level structural modality masking.Balanced preserves label and natural modality-mask distributions without additional structural modality drops.
  • Intra-client missingness: The +Intra-Client Missingness variants apply independent record-level modality dropout with p_intra = 0.3 after inter-client masking while preserving at least one observed modality per record.This dropout is applied to each remaining observed modality.
  • Modality-exclusive setting: The Mod-Excl. setting assigns each sufficiently frequent modality combination to a distinct client, creating disjoint modality patterns for severe specialization stress testing.If structural masking drops all modalities for a client, one modality is randomly retained to avoid a degenerate client.

D Implementation Details

Federated experiments use a shared training and evaluation protocol with client-level data splits, validation-F1 model selection, and consistent local optimization. Missing modalities are preserved through observed-modality masks and zero-valued placeholders, while baseline and MedMix configurations are tuned separately for their respective settings.

  • Shared federated protocol: All federated methods train for up to 150 communication rounds with one local epoch per round, full client participation, and Adam optimization.Data are split 70/15/15 within each client.
  • Shared federated protocol: Model selection uses validation F1 with smoothed early stopping, coefficient 0.1, minimum improvement 10−3, and patience equal to 20% of the training budget.These settings are applied within the shared federated protocol.
  • Missing-modality handling: Unavailable modalities use zero-valued placeholders with correct input shapes, while masking-capable methods receive each record’s observed-modality mask during fusion.Simulated missingness uses partition-specific modified masks, whereas natural splits retain original dataset-level masks.
  • Baseline hyperparameters: Baseline architectures follow centralized fusion configurations, with federated learning rates and batch sizes tuned separately; FedMoE and FedAlign-MoE are adapted to supervised multimodal prediction.Table 5 summarizes baseline learning rates and batch sizes.
  • MedMix hyperparameters: MedMix uses one sparse MoE fusion layer, one router, gate load-balancing weight 0.01, and router hidden dimension 256 across both datasets.MCR embedding dimensions are 32/64/16, while CRA uses support 160, quorum 2, alignment weight 0.02, and starts from round 2; targeted modality dropout is 0.3 with at least one modality retained.

E Additional Preliminary Experiment Details · F License of Assets

The supplementary material details preliminary diagnostics of routing, expert misalignment, and aggregation impact using ADNI and a sparse MoE backbone with FedAvg. It also documents restricted-access dataset use and licensing practices for ADNI and MIMIC-IV.

  • E Additional Preliminary Experiment Details: Preliminary diagnostics use ADNI, MedMix’s sparse MoE fusion backbone, standard FedAvg, final local models, and client validation sets.These settings apply unless otherwise stated.
  • E Additional Preliminary Experiment Details: For each modality pattern p = (r, m), top-2 router expert selections are counted and normalized within each client.Additional examples show diverging routing distributions across clients for identical modality patterns.
  • E Additional Preliminary Experiment Details: Figure 4 visualizes expert misalignment across clients in High and Mod-Excl. heterogeneity settings using shared embeddings and top-1 expert colors.The visualized structure varies across clients, extending Figure 1(b).
  • E Additional Preliminary Experiment Details: Aggregation impact is measured as ∆F1 = F1post − F1pre by evaluating each client before and after FedAvg on the same validation set.Negative ∆F1 values indicate that aggregation hurts the just-trained local model.
  • E Additional Preliminary Experiment Details: Figure 5 reports round-wise aggregation impact across four heterogeneity settings, averaging over clients and five seeds with standard-error shading.Lines show means over clients and seeds.
  • E Additional Preliminary Experiment Details: Expert-output t-SNE visualizations jointly fit one embedding across all five clients in each split and facet it by client.Points are colored by their top-1 routed expert.
  • F License of Assets: The work uses two restricted-access third-party clinical research datasets: ADNI and MIMIC-IV.The datasets are cited as ADNI and MIMIC-IV.
  • F License of Assets: The authors did not create, redistribute, sublicense, or release participant-level data and conducted experiments under applicable access agreements.ADNI data were obtained through the LONI Image and Data Archive.
Loading 2608.13911v1…