Skip to content

FM\(^2\): Unified Federated Foundation Models for Heterogeneous Multimodal Medical Imaging

🕒 Published (v1): 2026-07-15 02:28 UTC · Source: Arxiv · link

Why this paper was selected

Federated multimodal medical VLM; addresses privacy-constrained cross-institution medical AI deployment

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

FM\(^2\) is a federated learning framework that trains medical multimodal foundation models from scratch under Imaging Modality Heterogeneity — the practical regime where different clinical sites hold different imaging modalities. It uses dual Mixture-of-Experts modules to disentangle class-level and domain-level heterogeneity, and leverages GPT-4o-generated diagnostic captions as a privacy-preserving textual semantic bridge to transfer representations across clients with entirely disjoint visual modalities.

Problem

Existing Federated Foundation Models (FFMs) for medical imaging either fine-tune natural-image pretrained models (e.g., ViT, CLIP, SAM), which transfer poorly to the medical domain, or train from scratch on a single imaging modality for a single task. Neither approach addresses Imaging Modality Heterogeneity: clients in real clinical federations may hold (i) Overlapped modalities with non-IID label distributions, or (ii) fully Non-overlapped (disjoint) modalities per client. Standard personalized FL methods designed for label skew alone are insufficient for this compound heterogeneity, and no prior work unifies classification, caption-supervised learning, and VQA across these regimes within a single privacy-preserving framework.

Method

FM\(^2\) trains a ResNet-based visual encoder (the Observer \(\mathcal{V}\)) from scratch in a two-stage local protocol. Stage 1 trains \(\mathcal{V}\); its backbone then performs Top-\(K\) Sampling — selecting the \(K\) highest-scoring samples per class — to form a representative subset \(\mathcal{D}_s\).

Stage 2 trains two Mixture-of-Experts modules on \(\mathcal{D}_s\):

  • Class-wise MoE (CMoE) \(\mathcal{M}_c\): \(N_c\) experts (one per local class), kept local for personalization, combining as \(\hat{\mathbf{X}} = \sum_{n=1}^{N_c} W_n E_n(\mathbf{X}_\text{mid})\).
  • Domain-wise MoE (DMoE) \(\mathcal{M}_d\): \(N_d\) experts (one per imaging modality, indexed by a shared global modality vocabulary), globally aggregated via FedAvg to transfer cross-modality representations.

Inference combines both: \(y = \gamma_d \cdot \mathcal{M}_d(\mathcal{V}(\mathbf{X})) + \gamma_c \cdot \mathcal{M}_c(\mathcal{V}(\mathbf{X}))\) with \(\gamma_d + \gamma_c = 1\).

The Heterogeneous Modality Alignment (HMA) regularizer prevents DMoE expert drift by adding a proximal penalty \(\lambda \|\mathcal{M}_d^i - \hat{\mathcal{M}}_d\|^2\) to the local objective, constraining only the domain experts (not the full model as in FedProx/Ditto).

Caption-Enhanced Learning (CEL) extends FM\(^2\) to vision-language alignment: GPT-4o generates diagnostic captions locally via a ReAct-based pipeline; a local PubMedBERT encoder \(\mathcal{T}\) extracts caption embeddings \(\mathbf{t}\); a federated contrastive loss $\(\mathcal{L}_\text{CEL} = -\frac{1}{B}\sum_{j=1}^B \log \frac{\exp(\text{sim}(\mathbf{v}_j, \mathbf{t}_j)/\tau)}{\sum_{k=1}^B \exp(\text{sim}(\mathbf{v}_j, \mathbf{t}_k)/\tau)}\)$ aligns visual features \(\mathbf{v} = \mathcal{V}(\mathbf{X})\) with caption embeddings. Only the visual backbone and DMoE are aggregated; captions and the text encoder remain on-device.

Federated Medical VQA adds a question encoder \(\mathcal{T}_q\) and a lightweight answer head \(\mathcal{H}_\text{vqa}\) with element-wise gating: \(a = \mathcal{H}_\text{vqa}(\mathbf{v} \odot \mathbf{q})\).

The paper proves \(O(1/\sqrt{T})\) convergence under \(L\)-smooth losses, bounded gradient variance \(\sigma^2\), and bounded modality heterogeneity \(\delta^2\) (Theorem 1), and derives a generalization bound showing FM\(^2\)'s risk decomposes as \(R(h) \leq \gamma_d R_d^* + \gamma_c R_c^* + 2\sqrt{2\gamma_d\gamma_c\, d_\text{JS}(P_d \| P_c)} + \epsilon_\mathcal{V}\) (Proposition 2).

Key Contributions

  • Formalizes Imaging Modality Heterogeneity under Overlapped and Non-overlapped regimes and proposes FM\(^2\) as the first unified multi-task federated framework addressing both simultaneously.
  • Dual MoE architecture (CMoE local + DMoE global) with HMA regularizer that selectively constrains only domain experts, paired with \(O(1/\sqrt{T})\) convergence proof and an explicit generalization bound via Jensen–Shannon divergence.
  • Caption-Enhanced Learning (CEL): GPT-4o captions as a shared textual semantic bridge enabling vision-language contrastive supervision across clients with entirely disjoint visual modalities without any caption leaving the device.
  • MIMH benchmark: five MedMNIST modalities (Fundus, Blood Cell, Kidney Cortex, Colon Pathology, Dermatoscope), four configurations (MIMH-5/4/3-Domain overlapping and MIMH-SA non-overlapping), totaling up to 372,273 images across 10 clients.
  • Extensibility to Federated Medical VQA, evaluated on real-world datasets.

Results

  • MIMH-5-Domain (\(\varphi=0.1\), hard non-IID), \(r=100\%\): FM\(^2\)-Large achieves 86.68% vs. FedRep (best baseline) at 74.72% — a +11.96 pp gap.
  • MIMH-5-Domain (\(\varphi=0.1\), \(r=30\%\)): FM\(^2\)-Large 83.79% vs. FedRep 74.72%.
  • MIMH-3-Domain (\(\varphi=0.1\), \(r=100\%\)): FM\(^2\)-Large 90.39% vs. FedProto 85.00%.
  • Scaling from Tiny → Medium → Large consistently improves accuracy across all configurations.
  • MIMH-SA (Non-overlapped): Baselines collapse due to inability to capture cross-modality invariances; FM\(^2\) significantly outperforms even local-only training (E1 result).
  • Out-of-modality generalization (E2): FM\(^2\) fine-tuned on 5% data transfers effectively to three held-out modalities: PneumoniaMNIST (Chest X-Ray), OCTMNIST (Retinal OCT), OrganAMNIST (Abdominal CT).
  • Compared baselines: FedAvg, FedProx, PerFedAvg, FedBN, FedProto, FedRep (all using ResNet-4, SGD lr=\(5\times10^{-3}\), 100 rounds).

Limitations

  • Modality indices are pre-assigned via a shared global vocabulary at federation setup, assuming each client knows its imaging modality; extension to unknown-modality scenarios via online feature clustering is deferred to an appendix.
  • Caption generation via GPT-4o introduces a one-time external API dependency (though captions are retained locally thereafter).
  • The two-stage training protocol (50 Observer + 50 MoE rounds) and single local epoch match baseline budgets but may be sensitive to this split in practice.
  • The backbone is ResNet-4, a relatively small architecture; scalability to ViT-scale backbones is not demonstrated in the main paper.
  • Full VQA and CEL quantitative results are in sections truncated from the provided text, so headline numbers for those tasks cannot be reported here.

Relevance to Vision-Language Models

FM\(^2\) is directly relevant to federated VLM research: its CEL module operationalizes the CLIP-style vision-language contrastive objective under strict privacy constraints, using language as a modality-agnostic alignment signal — a setting where standard centralized VLM training is prohibited. The use of PubMedBERT embeddings of GPT-4o-generated captions as a shared semantic bridge for clients with disjoint visual domains is a principled demonstration that text can substitute for visual feature overlap in cross-modal representation transfer. The Federated Medical VQA extension further shows how VLM-style multimodal reasoning (visual feature + language question fusion) can be decomposed into locally trainable heads sharing a federated visual backbone. For researchers tracking VLMs, this paper provides a rigorous template for privacy-preserving vision-language alignment in specialized domains where data cannot be pooled.