MOSAIC: Modality-agnostic Spectral Alignment for Federated Image-level Weakly Supervised Tumor Segmentation under Client-specific Missing Modalities

Abstract: Trustworthy multimodal fusion in clinical settings requires handling incomplete and heterogeneous modality subsets across institutions, where privacy constraints prohibit centralized data sharing. Federated learning (FL) mitigates data-sharing constraints but suffers from client-specific missing modalities, where institutions possess incomplete multimodal subsets, degrading fusion quality and segmentation performance. While FL and weak supervision have been studied separately, their joint use with image-level labels under heterogeneous missing modalities remains unaddressed. We propose \textbf{MOSAIC}, the first modality-agnostic federated framework for weakly supervised binary tumor segmentation under client-specific missing modalities. We introduce a client-specific modality-alignment module that fuses available channels into a shared latent space without prior knowledge of modality identity, a spectral prototype alignment loss that reconciles cross-client distribution shift using compact non-invertible frequency-domain statistics, and a dedicated federated refinement network that denoises the resulting CAM pseudo-labels into accurate masks, breaking the accuracy ceiling of weak supervision. Experiments on three multi-institutional brain tumor benchmarks (FeTS2022, BraTS-MEN, and BraTS-SSA) demonstrate significant improvements over all image, box, and point-supervised baselines, approaching fully supervised accuracy using only image-level labels and reaching 0.84 Dice on FeTS2022. Dynamic new client addition enables previously unseen institutions to join an already-trained federation within 0.01-0.04 Dice without retraining. Code is available at https://github.com/Tarun2201/MOSAIC.

• The paper introduces MOSAIC, an end-to-end federated framework that combines client-specific modality adapters, spectral prototype alignment, and pseudo-label refinement for weakly supervised tumor segmentation.
• MOSAIC achieves Dice scores of 0.84, 0.81, and 0.78 on FeTS2022, BraTS-MEN, and BraTS-SSA, respectively, outperforming competing weakly supervised methods and remaining robust for single-modality clients.
• The paper shows that compact spectral statistics reduce modality-induced feature shift while preserving communication efficiency, although binary 2D MRI segmentation, formal privacy guarantees, and broader modality validation remain open challenges.

Problem setting and contribution

MOSAIC addresses a configuration that, according to the authors' survey of the literature, no prior method targets jointly: federated learning (FL), image-level (weak) supervision, and client-specific missing modalities in multimodal tumor segmentation. Each client $k$ holds a private set of 2D axial MRI slices with only a binary image-level label indicating tumor presence and a client-specific modality subset $\mathcal{M}_k \subseteq \mathcal{M}$; no modality is guaranteed at every client. Standard federated optimizers (FedAvg, FedProx, FedBN, MOON) assume dense supervision and complete modality availability; existing federated weakly supervised segmentation methods (FedDM, FedICRA, FedLPPA) use denser supervision forms and do not handle missing modalities; missing-modality methods require full pixel labels and modality-specific encoders. MOSAIC is positioned as the first end-to-end framework covering all three challenges simultaneously.

The framework rests on three components: (i) a count-based client-specific modality-alignment module, kept strictly local; (ii) a spectral prototype alignment (SPA) loss exchanging compact non-invertible frequency-domain band energies; and (iii) a two-phase pipeline in which federated CAM pseudo-labels are refined by a dedicated segmentation network.

Method

Modality-agnostic alignment. A lightweight per-client U-Net adapter $A_{\phi_k}$ maps the available $|\mathcal{M}_k|$ channels into a fixed $D$-channel latent space, conditioned only on channel count rather than modality identity. The parameters $\phi_k$ are never aggregated, adding zero communication overhead while keeping the shared downstream backbone modality-agnostic. This design directly enables dynamic federation membership: clients with previously unseen modality combinations can join without retraining or metadata bookkeeping.

Spectral prototype alignment. For each feature tensor, the client computes per-channel 2-D DFTs and partitions the frequency plane into $\Omega$ radial bands, producing a prototype $E^k_{\text{loc}} \in \mathbb{R}^{D \times \Omega}$ of band-wise mean amplitudes. Clients align local prototypes to an EMA-maintained global consensus via MSE; only these statistics are transmitted. The authors argue privacy from non-invertibility: radial mean-pooling over spatial positions is many-to-one, so feature maps cannot be reconstructed from the shared statistics. SPA is applied at two depths (alignment output and intermediate backbone features), and a region-conditioned variant (RCSA) aligns foreground and background spectra separately to avoid suppressing diagnostically relevant tissue contrast.

Phase 1: CAM generation. A multi-exit classifier produces CAMs that an aggregation network converts to binary pseudo-masks, trained with FedProx plus SPA. Notably, the paper identifies mislocalization—not under-segmentation—as the characteristic failure mode under missing modalities, since absent contrast channels shift activations toward high-intensity artifacts.

Phase 2: refinement. A student U-Net with 2.5D input is trained against a soft target formed by soft-voting three cues: a trimap-denoised CAM, an EMA teacher prediction, and a prototype cue computed by cosine similarity to federated foreground/background feature prototypes. Per-pixel reliability weights derive from cue agreement, gated further by the trimap. Aggregation uses reliability-weighted FedAvg ($r_k = \bar{w}^k$). The full objective combines weighted Focal-Tversky loss, prototype alignment hinge losses, teacher-student consistency, presence classification BCE (which also gates false positives on tumor-free slices at test time), multi-scale intensity-gated CRF loss, and both spectral losses. Side-channel communication totals 548 bytes per client per round.

Experimental results

Evaluation spans FeTS2022 (glioma, four clients), BraTS-MEN (meningioma, four clients), and BraTS-SSA (Sub-Saharan African glioma, three clients), each with deliberately heterogeneous modality assignments including single-modality clients, against sixteen baselines spanning fully supervised, box-, point-, and image-level supervision. All baselines receive the same modality-alignment module without SPA, isolating differences to supervision strategy and alignment.

Dataset	MOSAIC Dice / IoU / HD95	Best competing weak baseline	Fully supervised ref.
FeTS2022	0.84 / 0.79 / 32.6 mm	FL-W3S 0.75 Dice	FedBN/FedAvg 0.88
BraTS-MEN	0.81 / 0.79 / 56.5 mm	AME-CAM/FedNorm+ 0.78	MOON 0.84
BraTS-SSA	0.78 / 0.73 / 51.6 mm	ScoreCAM 0.70	FedBN 0.82

All improvements are statistically significant under paired Wilcoxon tests ($p<0.01$) across all metrics on all datasets except HD95 versus FedDM on BraTS-MEN. Three results stand out:

• Surpassing richer supervision. MOSAIC exceeds the best box-supervised method (FedDM) by 0.09 Dice on FeTS2022 and point-supervised FedICRA by 0.12, using only image-level labels—contradicting the implicit assumption that denser supervision yields better federated performance.
• Exceeding centralization. On FeTS2022, MOSAIC surpasses a centralized image-level reference trained on pooled data with the complete modality set by 0.02 Dice (0.84 vs. 0.82); on BraTS-SSA the centralized reference performs markedly worse (0.58 Dice) than the federated model. The authors attribute this to site-specific fusion adaptation avoiding representation confusion from concatenated incomplete inputs—a notable claim that privacy-constrained federated fusion can be qualitatively superior to data pooling.
• Single-modality robustness. On the FLAIR-only Client 4 of FeTS2022, MOSAIC reaches 0.82 Dice versus 0.58–0.78 for existing image-level approaches, supporting the claim that SPA compensates for severely incomplete fusion inputs.

Phase 1 alone already outperforms all competitors (0.80/0.79/0.75 Dice), establishing modality-robust CAM generation as the primary driver; Phase 2 adds +0.04/+0.02/+0.03 Dice, with the largest gains precisely where missing FLAIR causes worst mislocalization (e.g., Client 2 on SSA: 0.66→0.75).

Component validation

Alignment objective comparison. Replacing SPA with drop-in alternatives under identical communication protocols yields a strict ordering: std alignment (0.65 Dice) < contrastive InfoNCE (0.70) < mean alignment (0.73) < histogram-KL (0.78) < SPA (0.80). Since all five share equivalent privacy properties, SPA's advantage derives from the alignment signal itself: radial band energies capture modality-specific textural signatures that spatial moments and marginal histograms miss.

Direct divergence measurement. On a patient-aligned evaluation set (identical anatomy across clients, isolating modality mismatch), SPA reduces mean pairwise MMD$\mathcal{M}_k \subseteq \mathcal{M}$0 from 1.39 to 1.10 (−21%) and Fréchet distance from 346.8 to 167.4 (−52%) at the alignment output, with consistent reductions at the bottleneck (MMD$\mathcal{M}_k \subseteq \mathcal{M}$1 0.37→0.23, −38%). Because patient content is held constant, this constitutes direct evidence that SPA causally reduces modality-induced distribution shift.

Spectral resolution. Performance is non-monotonic in $\mathcal{M}_k \subseteq \mathcal{M}$2: peak at $\mathcal{M}_k \subseteq \mathcal{M}$3 (0.84 Dice), sharp degradation at $\mathcal{M}_k \subseteq \mathcal{M}$4 (0.79), attributed to overfitting high-frequency client noise. All tested values remain within 0.025 Dice of optimum.

Leave-one-out ablation. Every component contributes positively. Presence classification is most critical (removal costs 0.04 Dice and +6.5 mm HD95 via lost false-positive gating); trimap denoising follows (−0.02 Dice). The core-only variant (adapter + SPA, no refinement) achieves 0.82 Dice, already beating all competitors, confirming the core novel components account for most of the gain.

Dynamic client addition. A fifth unseen client joins within 0.01–0.04 Dice of base clients across all four modality configurations, with local-adapt and direct-join strategies differing by less than 0.01 Dice for most subsets.

Uncertainty. Test-time augmentation analysis shows Dice stability below 0.035 everywhere, low background entropy (≤0.0047 average), and highest foreground entropy consistently on the {T1ce, T2} subset lacking FLAIR—the uncertainty tracks the information content of available modalities, which the authors present as calibrated rather than overconfident behavior.

Limitations and open questions

The paper concedes several constraints explicitly. The framework handles binary segmentation only; extending to nested multi-class sub-regions is non-trivial because image-level labels become ambiguous when multiple classes coexist in one slice, and the authors suggest multi-label or coarse spatial supervision as possible remedies without validating them. The implementation is 2D slice-based; a 3D formulation would likely reduce boundary errors but requires stronger pseudo-label refinement since volumetric CAMs are noisier. Performance remains bounded by CAM quality—the refinement stage mitigates but does not remove this ceiling—and the authors note that foundation-model-derived or text-guided localization could raise it, again without evaluation. Validation is confined to multi-parametric brain MRI; whether frequency-domain alignment generalizes to genuinely cross-modal federated settings combining MRI with PET, CT, or clinical time-series remains untested. Finally, the privacy argument rests on non-invertibility of band-energy pooling alone; no formal guarantee (e.g., differential privacy bounds on the exchanged prototypes and model updates) is provided.

Conclusion

MOSAIC demonstrates that image-level supervision, the cheapest annotation form, suffices for accurate federated tumor segmentation even when clients hold disjoint modality subsets, reaching 0.84 Dice on FeTS2022 within 0.05 Dice of dense-supervision references and exceeding centralized pooled training. Its empirical support—statistically significant gains across three structurally distinct benchmarks, a controlled comparison of alignment objectives, and causal divergence measurements on patient-aligned data—is unusually thorough for the federated WSS setting. The principal open questions concern extension beyond binary masks, beyond 2D slices, and beyond homogeneous imaging modalities.