OmniFM: Toward Modality-Robust and Task-Agnostic Federated Learning for Heterogeneous Medical Imaging

Meilin Liu; Jiaying Wang; Jing Shan

cs.CVMar 23, 2026

OmniFM: Toward Modality-Robust and Task-Agnostic Federated Learning for Heterogeneous Medical Imaging

Authors: Meilin Liu, Jiaying Wang, Jing Shan

Organizations: School of Software, Shenyang University of Technology

Abstract

Federated learning (FL) has become a promising paradigm for collaborative medical image analysis, yet existing frameworks remain tightly coupled to task-specific backbones and are fragile under heterogeneous imaging modalities. Such constraints hinder real-world deployment, where institutions vary widely in modality distributions and must support diverse downstream tasks. To address this limitation, we propose OmniFM, a modality- and task-agnostic FL framework that unifies training across classification, segmentation, super-resolution, visual question answering, and multimodal fusion without re-engineering the optimization pipeline. OmniFM builds on a key frequency-domain insight: low-frequency spectral components exhibit strong cross-modality consistency and encode modality-invariant anatomical structures. Accordingly, OmniFM integrates (i) Global Spectral Knowledge Retrieval to inject global frequency priors, (ii) Embedding-wise Cross-Attention Fusion to align representations, and (iii) Prefix-Suffix Spectral Prompting to jointly condition global and personalized cues, together regularized by a Spectral-Proximal Alignment objective that stabilizes aggregation. Experiments on real-world datasets show that OmniFM consistently surpasses state-of-the-art FL baselines across intra- and cross-modality heterogeneity, achieving superior results under both fine-tuning and training-from-scratch setups.

Figures & tables

Figure 1: Conventional federated learning for medical image is task-specific and assumes homogeneous imaging modalities. Our OmniFM offers a single reusable pipeline that supports diverse backbones and remains robust under heterogeneous modalities for multiple tasks.

Figure 2: Optimization dynamics under heterogeneous imaging modalities, where MRI in Client 1 and Microscope in Client 2.

Figure 3: Overview of our OmniFM . Client-side image feature embedding and FFT-derived spectral embeddings are fused with retrieved global spectral prototypes via embedding-wise cross-attention and prefix-suffix prompting. Spectral embeddings are periodically uploaded to update a global knowledge bank, while a spectral-proximal alignment term mitigates modality-induced drift during optimization.

Figure 4: Spectrum from different imaging modalities.

Method	Scenario 1 (Hard)						Scenario 2 (Medium)						Scenario 3 (Extra)
	Acc (%)			F1-Score			Acc (%)			F1-Score			Acc (%)			F1-Score
	20%	50%	100%	20%	50%	100%	20%	50%	100%	20%	50%	100%	20%	50%	100%	20%	50%	100%
FedAvg [ 39 ]	56.31	67.31	84.21	0.234	0.380	0.474	56.44	67.11	84.49	0.229	0.334	0.460	24.24	56.84	99.49	0.165	0.522	0.977
FedProx [ 33 ]	27.04	58.79	93.92	0.060	0.293	0.491	30.17	38.33	30.14	0.119	0.124	0.131	24.07	55.79	97.53	0.130	0.369	0.587
Ditto [ 32 ]	28.67	65.31	97.56	0.096	0.385	0.649	30.23	69.23	96.11	0.123	0.486	0.820	25.09	58.66	99.24	0.170	0.447	0.931
SCAFFOLD [ 29 ]	24.09	16.56	19.85	0.046	0.063	0.082	24.04	18.01	20.59	0.053	0.067	0.078	52.31	49.99	99.41	0.241	0.180	0.846

Table 1: Results on federated cross-modality classification under three heterogeneous scenarios (more details in Appendix). We employ ResNet-18 [ 23 ] as the local backbone, train for 100 rounds with 5 local epochs per round. Bold : the best, underline : the second best.

Scale	Client	Scenario 1										Scenario 2
		PSNR $\mathbf{\uparrow}$					SSIM $\mathbf{\uparrow}$					PSNR $\mathbf{\uparrow}$					SSIM $\mathbf{\uparrow}$
		OmniFM	FedAvg	FedProx	FedPer	Local	OmniFM	FedAvg	FedProx	FedPer	Local	OmniFM	FedAvg	FedProx	FedPer	Local	OmniFM	FedAvg	FedProx	FedPer	Local
$\mathbf{\times 2}$	Client 1	42.93	35.90	20.68	40.24	36.97	0.9784	0.9526	0.5426	0.9537	0.9115	43.11	41.67	24.82	41.42	40.38	0.9811	0.9399	0.8230	0.9610	0.9583
	Client 2	41.95	36.07	22.42	39.32	35.09	0.9634	0.9444	0.5462	0.9478	0.9231	41.95	41.70	24.09	41.98	42.23	0.9690	0.9582	0.8152	0.9631	0.9651
	Client 3	41.75	35.87	22.44	39.00	34.81	0.9661	0.9543	0.5513	0.9567	0.9533	41.85	41.12	24.20	39.22	34.92	0.9634	0.9661	0.8076	0.9527	0.9472
	Avg.	42.21	35.95	21.84	39.52	35.63	0.9693	0.9505	0.5467	0.9527	0.9293	42.30	41.50	24.37	40.87	39.17	0.9711	0.9548	0.8152	0.9590	0.9569

Table 2: Results on super-resolution with three independent clients across three scales. Bold : the best, Underline : the second best.

Method	F-C		F-L		F-CL		F-2stage
Method	IID	Non-IID	IID	Non-IID	IID	Non-IID	IID	Non-IID
FedAvg [ 39 ]	0.783	0.775	0.806	0.802	0.823	0.827	0.811	0.814
FedProx [ 33 ]	0.734	0.750	0.800	0.780	0.816	0.796	0.773	0.785
FedAdam [ 43 ]	0.741	0.735	0.783	0.771	0.777	0.774	0.782	0.777
FedAvgM [ 25 ]	0.754	0.747	0.789	0.786	0.784	0.768	0.793	0.794
FedYogi [ 43 ]	0.745	0.736	0.782	0.769	0.783	0.774	0.785	0.782

Table 3: Results on Task 1 under IID and non-IID. F-C and F-L denote connector-only and LLM-only tuning, respectively; LC and 2-stage represent joint one-stage and sequential connector–LLM tuning. Best are in bold and second-best are underlined . Following, we adopt LLaVA-1.5 with a CLIP ViT-B/32 visual encoder and LLaMA-3.2-3B text backbone, and train LLM components via LoRA [ 26 ] (

r

= 8,

\alpha

= 32) for parameter-efficient fine-tuning.

Method	Task-Specific Performance								Avg.
	C1	C2	C3	C4	C5	C6	C7	C8
FedAvg	88.67	72.59	70.86	85.93	71.09	89.45	66.62	76.48	77.71
FedProx	88.92	73.05	71.03	86.10	71.40	89.20	66.81	76.32	77.85
FedAdam	89.10	74.22	71.58	85.70	71.85	89.80	67.20	76.90	78.17
FedDyn	88.75	73.90	71.92	86.35	71.55	90.02	67.45	76.70	78.08
FedPer	89.43	74.10	72.21	85.95	72.30	89.90	67.82	77.52	78.40

Table 4: Results on Task 2. C1: Computed Tomography, C2: Ultrasound, C3: Optical Coherence Tomography, C4: Fundus Photography, C5: Light Microscopy, C6: Histopathology, C7: Dermatoscopy, C8: Chest X-ray. Training protocol is consistent with Table 3 . Bold : the best, and underline : the second best.

Method	SLAKE	VM2019	VM2020	VM2021	VR	Average
Full Tuning
$\triangleright$ FedAvg	77.45	67.25	15.06	22.00	41.52	44.66
$\triangleright$ FedProx	77.47	67.00	14.87	22.12	41.00	44.33
$\triangleright$ FedAdam	77.00	67.89	15.12	22.44	42.42	44.77
$\triangleright$ FedDyn	78.24	69.00	16.21	21.65	42.02	45.02
$\triangleright$ FedPer	79.23	69.20	16.77	22.54	43.43	46.03

Table 5: Results on Task 3. Training protocols follow Table 3 . Bold indicates the best and underline denotes the second best.

Figure 5: Sensitivity analysis of the alignment coefficient

\lambda

and the top-

k

value in GSKR across representative tasks. Default configurations are indicated in red. Due to the paired-image input design, the task of multimodal fusion is evaluated separately.

Method	Group 1 (9 Clients)		Group 2 (6 Clients)		Group 3 (5 Clients)		Group 4 (3 Clients)		Avg.
Method	Dice	IoU	Dice	IoU	Dice	IoU	Dice	IoU	Dice	IoU
Local	45.24	29.33	44.21	25.00	41.25	27.22	46.23	30.00	44.23	27.89
FedAvg	44.02	28.56	44.91	25.32	40.93	26.10	46.78	30.95	44.16	27.73
FedProx	42.32	25.23	44.00	24.43	38.99	25.32	45.35	29.76	42.66	26.19
FedPer	71.23	66.99	75.24	62.12	72.43	57.34	76.54	62.99	73.86	62.36
OmniFM	79.84	74.93	80.82	68.74	78.62	66.01	78.03	65.31	77.41	68.75

Table 6: Results on segmentation across 4 intra-modality heterogeneous client groups. Bold : the best, Underline : the second best.

Metric	Client	FedAvg	FedProx	Ditto	FedDyn	MOON	OmniFM
Metric	Client	[ 39 ]	[ 33 ]	[ 32 ]	[ 1 ]	[ 31 ]	Ours
VIF $\uparrow$	CT-MRI	0.221	0.185	0.196	0.182	0.195	0.217
	PET-MRI	0.331	0.242	0.275	0.239	0.275	0.309
	SPECT-MRI	0.284	0.235	0.249	0.245	0.249	0.280
	Avg	0.279	0.221	0.240	0.222	0.240	0.269
SSIM $\uparrow$	CT-MRI	0.739	0.684	0.718	0.716	0.718	0.739

Table 7: Results on multi-modal fusion across heterogeneous client configurations. Bold : the best, Underline : the second best.

Variant	CLS	SR	SEG	VQA-T1	VQA-T2	VQA-T3
OmniFM	97.82	35.79	75.04	0.800	79.27	44.21
w/o GSKR	95.82	34.88	73.42	0.775	78.91	42.62
w/o ECA	96.51	35.04	74.03	0.783	79.15	43.01
w/o PSP	96.92	35.12	74.51	0.787	79.42	43.19
w/o SPAlign	97.14	35.21	74.09	0.781	77.64	42.37

Table 8: Ablation on representative tasks. CLS: Cross-modality classification; SR: Super-resolution; SEG: Segmentation.

Explore similar work

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
Federated Learning for Cross-Modality Medical Image Segmentation via Augmentation-Driven Generalization
Feb 24, 2026Sachin Dudda Nagaraju, Ashkan Moradi, Bendik Skarre Abrahamsen +1Cross-Modality Medical Image Segmentation Data Augmentation
Department of Circulation and Medical Imaging, Norwegian University of Science and Technology, Trondheim, Norway. · Central Staff, St. Olavs Hospital, Trondheim University Hospital, Trondheim, Norway.
OmniMed-FL: A Robust Multimodal Federated Learning Framework for Clinical Diagnosis
Sep 9, 2026Ayush Debnath, Ruelia Saha, Sudip MisraMultimodal Federated Learning Medical Diagnosis
Indian Institute of Technology Kharagpur, India · KTH Royal Institute of Technology, Sweden · SRM University-AP, India +1

Jul 15, 2026J/K move · Enter open · S save

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

Shengchao Chen, Ting Shu

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

Multimodal Federated Learning Cross-Modal Alignment Medical VQA Medical Image Classification Multimodal Foundation Models Federated Learning+1

Open paper AskPDF arXiv

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 FM

^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. FM

^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/\sqrt{T})

convergence and generalization guarantees. FM

^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.

Abstract

Explore similar work

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

Federated Learning for Cross-Modality Medical Image Segmentation via Augmentation-Driven Generalization

OmniMed-FL: A Robust Multimodal Federated Learning Framework for Clinical Diagnosis

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