Distributionally Robust Alignment for Medical Federated Vision-Language Pre-training Under Data Heterogeneity

Motivation

• Federated learning enables privacy-preserving use of multimodal data across medical institutions, addressing VLP’s dependence on large-scale datasets.Medical data-sharing privacy concerns make dataset scaling difficult, while federated learning leverages data across institutions without centralizing it.Source: §1 · p. 1

Challenge

• Heterogeneous client distributions can cause local models to overfit client-specific multimodal correlations, so parameter averaging may yield non-generalizable cross-modal alignment.Biased deep layers can prevent learning transferable and diverse semantics during local training.Source: §1 · pp. 1–2

Method

• FedDRA optimizes alignment over a family of potential testing distributions while dynamically weighting local training to improve transferability.The framework uses distributionally robust optimization to bridge client pre-training samples and downstream domains.Source: §1 · p. 2, Figure 1 · p. 2
• FedDRA further uses global-model anchor representations and two-stage training that tunes deeper layers before the entire network to reduce client-specific overfitting.These constraints guide local updates toward more transferable representations.Source: §1 · p. 2

Results

• Experiments show effectiveness across image-text retrieval, classification, and segmentation, with higher average and worst-client retrieval accuracy than the baseline and stronger results on few-shot tasks.Using ConVIRT as the backbone, the method beats other baseline strategies on few-shot classification and segmentation.Source: §1 · p. 3

Medical Vision-Language Pre-training

• Medical VLP methods align image and language representations using contrastive, fusion, and domain-informed objectives, but generally depend on large pre-training datasets.This dependence limits adaptation to modalities with limited training samples and real-world deployment.Source: §2 · p. 3

Heterogeneity in Self-Supervised Federated Learning

• Federated self-supervised learning addresses heterogeneous decentralized data, but representation-sharing approaches can compromise privacy and unimodal methods do not directly model cross-modal alignment.Multimodal learning must align modalities in addition to learning robust features.Source: §2 · p. 3
• Distributionally robust optimization offers a way to mitigate client-specific spurious correlations by optimizing worst-case group performance, but its application to federated multimodal alignment remains distinct.The related work identifies this framework as adaptable to federated learning tasks while noting the multimodal modality gap.Source: §2 · p. 3

Formulation of Pre-training Dataset and Heterogeneity

• Each of N clients holds paired multimodal samples from its own distribution D_i, and variation among these distributions defines data heterogeneity.The formulation assumes shared latent semantics connect the modalities across clients.Source: §3 · p. 4
• Potential test domains are represented by an uncertainty set Q_c={Q:D_f(Q∥D_C)≤ρ}, where larger ρ includes more unseen distributions.D_C is the distribution formed by grouping all client data, and D_f is an f-divergence.Source: §3 · p. 4

Federated Vision-Language Pre-training

• Federated multimodal pre-training trains client models locally and aggregates them on the server to learn a generalizable model from paired, unlabeled data.FedAvg is given as an example that averages model parameters across communication rounds.Source: §3 · p. 4
• Vision-language encoders project image and text inputs into a shared representation space, with cross-modal alignment measured through a loss such as mutual information.The target objective minimizes expected alignment loss on the testing distribution.Source: §3 · p. 4
• Because local optimization follows D_i, aggregated models may retain client-specific information that fails to generalize across client domains.This is the formulation’s central failure mode under heterogeneous local datasets.Source: §3 · p. 4

4.1 Global Constrained Local Training Objective

• Client heterogeneity can bias local encoders toward distinct domain-specific optima, weakening the aggregated model’s generalization to unseen testing domains.The method explicitly targets the discrepancy between locally trained and server-aggregated models.Source: §4.1 · p. 5
• FedDRA constrains local training with global representations so local models learn features that generalize beyond client-specific domains.The constraint compares local and server-aggregated representations, complementing the mutual-information pre-training loss.Source: §4.1 · p. 6

4.2 Two-Stage Alignment For Mitigating the Deeper Layer Bias

• FedDRA first trains alignment modules with frozen feature encoders, then jointly updates alignment modules and encoders using the same constrained objective.The staged design targets deeper alignment layers before updating the full network.Source: Figure 2 · p. 5, §4.2 · p. 6

4.3 Learning Robust Cross-Modal Alignment via Distributionally Robust Optimization

• The method constructs a family of potential testing distributions and optimizes the model for its worst-case alignment performance within that family.This distributionally robust objective is motivated by unknown, potentially out-of-distribution test domains.Source: §4.3 · p. 6
• FedDRA decentralizes distributionally robust optimization by alternating local model updates with updates to distribution weights under an uncertainty-radius constraint.The server updates weights after local training, and the resulting weights guide subsequent client updates.Source: §4.3 · p. 7
• The robust objective is applied in both stages, optimizing alignment-module parameters first and then the feature encoders together with the alignment modules.The two stages share Llocal but differ in which parameters θ are optimized.Source: §4.3 · p. 7

5.1 Experiment Setting

• Experiments use heterogeneous chest X-ray and ophthalmology clients, with downstream evaluation spanning retrieval, classification, segmentation, and out-of-distribution transfer.Chest X-ray experiments use five LDA-partitioned MIMIC-CXR clients, while ophthalmology experiments use MESSIDOR and BRSET clients and evaluate on shifted MBRSET data.Source: §5.1.1 · p. 8, §5.1.2 · p. 8

5.1.3 Backbones and Baselines

• Across ConVIRT, MGCA, and GLoRIA backbones, FedDRA adapts multi-modal pre-training to heterogeneous federated settings and improves classification and segmentation performance.The experiments compare contrastive VLP backbones with federated and centralized baselines under consistent encoder implementations.Source: §5.1.3 · pp. 8–9, §5.3 · p. 11

5.2 Empirical Finding

• FedAvg outperforms decentralized pre-training on retrieval, but centralized pre-training remains an upper bound, showing substantial room for improvement under heterogeneity.The study evaluates retrieval because it reflects cross-modal mutual-information maximization and alignment.Source: §5.2 · p. 9
• Local retraining lowers averaged accuracy after multiple communication turns, while biased decentralized deep layers continue to limit generalizable feature learning.Retraining shallow layers provides some improvement but leaves a significant gap relative to FedAvg, implicating client-specific information in deeper layers.Source: §5.2 · pp. 9–10

5.3 Main Results

• FedDRA improves both average and worst-client retrieval accuracy and outperforms baselines on few-shot classification and segmentation tasks.These results indicate more robust cross-client features and stronger generalization of the learned representation space.Source: §5.3 · p. 10, Table 2 · p. 10

5.4 Analysis Experiments

• Removing global constraints or first-stage alignment training reduces downstream or retrieval performance, supporting both components of the two-stage strategy.The ablations report particularly lower retrieval accuracy without the global constraint and better alignment after first-stage training.Source: Table 3 · p. 10, §5.4 · p. 11
• DRO weighting improves few-shot transfer and worst-client accuracy by optimizing performance over an uncertain distribution set that bridges local training data and downstream domains.The analysis attributes improved transferability and client-wise robustness to the DRO component.Source: §5.4 · pp. 11–12
• Larger uncertainty radii improve out-of-domain classification and segmentation but can reduce in-domain retrieval, while larger global-constraint values help until excessive values hurt overall performance.Increasing the global constraint from 1 to 5 consistently improves downstream performance, whereas excessively large values decrease it.Source: §5.4 · p. 12

6 Conclusion

• FedDRA mitigates heterogeneity in federated medical vision-language pre-training, with effectiveness verified by comprehensive experiments.Source: §6 · p. 12
• The study prioritizes privacy by avoiding transmission of representations between clients, although this may limit improvement from cross-client representation sharing.Source: §6 · p. 12
• Future work could introduce greater multimodal pre-training diversity while keeping local data private.Source: §6 · p. 12

A.1.1 Pre-training setup

• MIMIC-CXR provides 227,835 image-text pairs from 65,379 patients for multimodal pre-training using only image-text pairs.Source: §A.1.1 · p. 17
• Five client datasets are constructed by applying LDA to MIMIC-CXR disease-category attributes, with heterogeneity degree 1 in the main experiments and 5 in analysis experiments.Source: §A.1.1 · p. 17
• The disease categories are converted into a 25-category multiclass label from five binary disease indicators before LDA partitioning.Source: §A.1.1 · p. 17
• Figure 5 illustrates the connection between a latent variable and the text modality.Source: Figure 5 · p. 17
• Main experiments use five heterogeneous client datasets, 25 communications, uncertainty radius ρ = 0.1, and µ = 5, with 50 randomly sampled batches per communication.Source: §A.1.1 · p. 17

A.1.2 Downstream tasks

• Transferability is evaluated through few-shot classification, medical image segmentation, and image retrieval.Source: §A.1.2 · p. 17
• Few-shot classification uses RSNA Pneumonia Detection and Covidx, with a linear layer fine-tuned on 1% and 10% of training data and accuracy reported.Source: §A.1.2 · pp. 17–18
• Segmentation uses RSNA masks, a frozen pre-trained image encoder, a fine-tuned U-Net decoder, and Dice scores under 1% and 10% data regimes.Source: §A.1.2 · p. 18
• Image retrieval evaluates whether text queries retrieve their paired images using top-1 and top-5 recall accuracy on local-client validation splits.Source: §A.1.2 · p. 18

A.2.1 Pre-training setup.

• Ophthalmology pre-training uses MESSIDOR and BRSET from France and Brazil, assigning them to two clients with captions derived from diabetic retinopathy and edema records.Source: §A.2.1 · p. 18

A.2.2 Downstream tasks.

• MBRSET evaluates transferability under a significant distribution shift because it was collected in low-income areas using portable devices.Source: §A.2.2 · p. 18
• Few-shot classification predicts diabetic retinopathy and edema status using 10%, 20%, and 100% of training data, with classification accuracy reported.Source: §A.2.2 · p. 18

• FedAvg achieves better downstream transferability and in-domain image-text retrieval than decentralized pre-training, yet remains below centralized pre-training.The remaining gap is attributed to lower batch diversity in federated training, which may limit robust alignment learning.Source: B Additional Experiment Results · p. 19, Table 6 · p. 19

• The appendix includes detailed ablation and empirical-study results corresponding to Fig. 4 and Section 5.2.Source: C Detailed Experiment Results · p. 19
• Additional experiments evaluate ablations over client number, uncertainty radius, and global constraint degree.Source: Table 7 · p. 19, Table 8 · p. 19, Table 9 · p. 20
• Client-level analyses report retrieval accuracy and server-versus-local performance after communication rounds.Source: Table 10 · p. 20, Table 11 · p. 20
• Further tables compare client test accuracy for FedAvg baselines, retrained local models, centralized pre-training, and decentralized models.Source: Table 12 · p. 21, Table 13 · p. 21

D.1 Derivation of Proposition 1

• Lemma 1 upper-bounds contrastive loss using batch-average and dataset-maximum L2 representation distances under bounded-distance assumptions.Source: §D.1 · p. 22
• The derivation uses the bound Dmax ≤ L · D with D = 1/bz, substitutes it into the preceding inequality, and then sets αi = 5α.Source: §D.1 · pp. 22, 24
• The proof assumes R_T(·) is the contrastive loss and concludes Proposition 1 after successive inequality substitutions.Source: §D.1 · pp. 22–24
• The proof expresses target-domain generalization error as expected contrastive loss and applies Lemma 1 to derive an upper bound.Source: §D.1 · p. 23
• The resulting bound decomposes error into representation discrepancies, cross-client terms, and a client-specific constant.Source: §D.1 · pp. 23–24