Distributionally Robust Alignment for Medical Federated Vision-Language Pre-training Under Data Heterogeneity
Authors: Zitao Shuai
Chenwei Wu
Zhengxu Tang
Liyue Shen
Venue: Accepted by TMLR ()

Abstract: Vision-language pre-training (VLP) has emerged as an effective scheme for multimodal representation learning, but its reliance on large-scale multimodal data poses significant challenges for medical applications. Federated learning (FL) offers a promising solution to scale up the dataset for medical VLP while preserving data privacy. However, we observe that client data heterogeneity in real-world scenarios could cause models to learn biased cross-modal alignment during local pre-training. This would limit the transferability of the federally learned representation model on downstream tasks. To address this challenge, we propose Federated Distributionally Robust Alignment (FedDRA), a framework for federated VLP that achieves robust vision-language alignment under heterogeneous conditions. Based on client datasets, we construct a distribution family that encompasses potential test-time domains, and apply a distributionally robust framework to optimize the pre-trained model's performance across this distribution space. This approach bridges the gap between pre-training samples and downstream applications. To avoid over-fitting on client-specific information, we use anchor representation from the global model to guide the local training, and adopt a two-stage approach to first tune deeper layers before updating the entire network. Extensive experiments on real-world datasets demonstrate FedDRA’s effectiveness in enhancing medical federated VLP under data heterogeneity. Our method also adapts well to various medical pre-training methods.

Published in Transactions on Machine Learning Research (06/2025)
Distributionally Robust Alignment for Medical Federated Vision-Language Pre-training Under Data Heterogeneity
Zitao Shuai, Chenwei Wu, Zhengxu Tang & Liyue Shen
University of Michigan
(liyues)Oumich.ee edu
Reviewed OpenReview: httppp///ppaaariiee..ee///orrrtt
Abstract
Vision-language pre-training (VLP) has emerged as an effective scheme for multimodal repre- sentation learning, but its reliance on large-scale multimodal data poses significant challenges for medical applications. Federated learning (FL) offers a promising solution to scale up the dataset for medical VLP while preserving data privacy. However, data heterogencity in real-world scenarios could cause models to learn biased cross-modal aligmment during local model training. This would limit the transferability of the federally learned representation model on downstream tasks. To address this challenge, we propose Federated Distributionally Robust Aligument (FedDRA), a framework for federated VLP that achieves robust vision-language alignment under heterogeneous conditions. Based on client datasets, we construct distribution family that includes potential test-time domains, and apply distributionally robust framework to optimize the pro-trained model's perfor- mance across this space, thus bridging the distribution gap between pre-training data and downstream applications. To avoid over-fitting on client-specifie information, we use anchor representation from the global model to guide the local training, and adopt a two-stage ap- proach to first adapt deep layers before updating the entire network. Extensive experiments on real-world datasets demonstrate FedDRA's effectiveness in enhancing medical federated VLP under data heterogeneity. Our method also adapts well to various medical pre- traiiingn methods. Code is available: Official Implementation.
observe that client
1 Introduction
Vision-language pro-training (VLP) learns transferable multinodal representations by extracting latent se- mantics from large-scale image-text pairs, where the dataset scale largely impacts the performance of the learned model (Oquab et al., 2023). However, scaling up multimodal pre-training datasets is a non-trivial challenge especially for modical applications, due to privacy concerns and regulations of patient data shar- ing (Ladbury et al. 2023). Recent work has explored federated learning as solution to leverage data across multiple medical institutions while preserving privacy (Lu et al. 2023).
However, in real-world scenarios, datascts collected from different institutes always heterogencous. For example, hospitals in tropical regions receive high proportion of pneumonia patients, whereas those in colder climates may see more pncumothorax cases (Mendogni et al., 2020). This data heterogeneity is not only a long-standing problem in conventional federated learning (Ghosh et al.) 2019 Huang et al. 2022) but also a practical challenge that impedes the deployment multimodal learninng in federated setting, such as medical languae pre-training. Current medical VLP methods often focus on learning a modality-
shared latent space, where the multimodal training data pairs are well-aligned. However, such learned cross-modal alignments may not be transferable lo data from unseen distributions. As shown in Fig. D this
*corespanding author

---

Published in Transactions on Machine Learning Research (06/2025)
(1). Demonstration Foderated Pre-traiming under data beterogencity
(a). Local pre-trsin with unloboled mulli-model pairod data
[b). Server regeropation
(c) Berativv communicato
Encocler
Clicnt
Encoder
Client
P:
rining
Server
Agaragane socil mode
breakae
Clecst transizz
Com, tum
6f
ofX
Com mars
Sarvar 06
LaIa
Local Model 0g
Pre-train Goot: Cross-modal align 1y)
LnnIa
Local 0x
Pre-train Goal:
broedean
Coos--undal aligns
Hetero.
Encoder
Hetern. Dy
Encoder Y
0
FedAug: of
Fllnng
Cors. Clentllcaalminn triiiin g of
Yf
plane Sohake
Global Model
Com tam 1 Server of..0 & booaddon
Modility
(2), Illstratioon our motivation through the mutual information (MI) view
(4). Representatonn space before ccermuu tum
(bi. With baseline methods
(e). With our FedDRA method
Mindalit
Midality
Modality X
Modulitt Y
Modality X Modality Y
oopec
MI visteal los dorunn
segrz.
Y2
Y2
clien Mi
X2
1/2
X2
Y2
Y22
rrpm. poir
Liss Aligned
Update with equal step-siee Local training
Perfonn putential distribution aggregation
Optimiiee with higher weight After trnining
1yx,,2y)
Less Aligned
aggr-gation
Figure 1: (1) We aim to tackle the data-hungry problem in VLP via federally utilizing heterogeneous multi- modal paired data. During local pre-training, local models pre-trained with naive methods may capture biased client-specific patterns, and simply aggregating them may fail to obtaín a robust model. (2). While some clients' data can be easily fitted, cross-modal alignments in other clients are difficult for local models to learn. The simply averaged model would perform worse when the testing distribution is close to less aligned distributions. To balance the fitting of local clients and handle the disparity between training and testing data, our method employs the DRO framework, optimizing model performance over a family of potential testing distributions by dynamically adjusting the update step size each client.
will harm the performance of the federally pre-trained model, which is aggregated from client models trained on heterogeneous local datasets.
We start by investigating how data heterogeneity affects the performance of federally pre-trained vision- language models. In classic medical VLP (Wang et al. 2022 Bannur al. 2023), the model learns cross- modal alignment through maximizing the mutual information of the two modalities based on the observed training data. In federated settings, this approach often learns local models that overfit client-specific information. As retenstrated in Fig. [ simply averaging parameters of biased client models, may obtain model that performs better on certain over-fitted client data distributions, but fails to robustly generalize to the varying testing domain. Secondly, biased deep layers, which overfit multi-modal correlations of local datasets, would prevent the model from learning transferable and diverse semantics during local training.
In this paper, we propose a Federated Distributionally Robust Alignment (FodDRA) framework, learn transferable cross-modal aligument under data heterogeneity. Our key idea is maximize gross-modal mutual information with distributional robustness. Specifically, to bridge the gap between the downstream testing domain and pre-training data domain, we construct a set of distributions based on client distributions. Then we employ a decentralized distributionally robust optimization method to iteratively improve the pre- trained model's performance on this constructed set. To alleviate the negative eflect. of over-fitting client- specific information, we maintain a global model to provide anchor representations for guiding local training and utilize two-stage training scheme finctune deep layers before updating the whole network.
Our contributions primarily focus
2

---

Published in Transactions on Machine Learning Research (06/2025)
We for the first time tackle the problem of medical VLP under the federated setting by utilizing het- erogeneous multi-modal datasets from different institutes. We conduct empirical studies to analyze the influence of data heterogeneity on federated multi-modal learning.
We propose FedDRA to address the data heterogeneity challenge in federated VI.P to obtain trans- ferable cross-modall alignment. It iteratively optimizes model performance on distribution family and uses a two-stage global-guided local training strategy to reduce over-fitting on elient-specific patterns.
Experiment results show the effectiveness of our method in learning multi-modal representations under the federated setting for various downstream tasks, including image-lexi retrieval, image classification and segmentation.
2 Related Work
Medical Vision-Language Pre-training. Pro-training multi-modal models large-scale datasets and then transferring learned knowledge to downstream tasks has become popular approach to leverage diverse
semantics contained in multi-modal unlabeled data (Li et al. 2022b Bao et al. 2022 Radford et al. 2021
Current works aim to learn a shared latent space to connect the representations of each modality, leveraging
a wide range of self-supervised learning methods, i.c., contrastive learning (Radford et al.
2021 Chen ct al.
2020) and multi-modal fusion (Li et al., 2021a] Chen et al.) 2022). Medical multi-mnoda
pre-training lasks
are often conducted on vision-based datasets, especially vision-language pre- treiirinnngg
Zhang et al.
(2022)
first utilizes image-text contrastive loss to align visual and language representations.
Huang ct al.
2021
aligns fine-grained cross-tnodal representations through a word-patch contrastive loss, and has improved the
performance on fine-grained vision tasks. Furthermore, recent work (Wang et al. 2022 Bannur et al. 2023)
incorporates medical domain knowledge to mitigate the misalignment during pre-training. However, almost all of the current methods still rely on large-scale pre-training datasets, which impede their adaptaion to modalities with limited training samples and deployment. real-world scenarios.
Heterogeneity in Self-Supervised Federated Learning. Federated self-supervised learning aims to leverage diverse semanties in local unlabeled datasets in decentralized and privacy-preserved way. One of the key challenges of federated learning is data heterogeneity (Li & Wang 2019 Collins et al.) 2021 Li et al.
2021b Tong et al., 2023). which has been long discussed in the uni-modality scenarios. Typically, Zhang et al. (2023 Huang et al. b002 Li et al. (2022a) employ additional communications on local data representations to increase sample diversity. Such methods fail to protect data privacy, which is a vital concern in medical applications. On the other hand, Zhhang et al. (2021): Li et al. (2021b) utilize server model constrain the update of local models, and Yan et al. (2023) utilizes the mask-autoencoodr handle heterogeneity. Zhuang et al.] (2022); Li et al. (2021b): Han el. al.
2022 consider distilllation-baseed methods yet ignore the direct modeling of cross-modk alignment. However, these uni-modal self-supervised learning methods have not accounted for the modality gap (Zhang et al. 2024c between multi-modal data. While uni-modal self- supervised learning aims to learn robust features Radford el al.] 2021), multi-mddal learning needs to align input modalities to maximize the mutual information between their representations (Su et al. 2023). Recent advances (Lu et al.] 2023) have verified that federated learning can be utilized to scale up the pre-training dataset. However, this work hasn't considered the harm of data heterogeneity issue (Ghosh et al. 2019
which may result. in learning biased local model over-reliance on spurious correlations (Saab al. 2022)
that are client-specific. To tackle this, distributionally robust optimization (Deng et al.
2020) framework
can alleviate these issues by optimizing the group-wise worst-case performance on given objeclive (Liu et al.
2022). which can be flexibly adapted to various of federated learning tasks (Han et al.
2023 Rehman et al.
2023
Capitani et al. 2024
3 Problem Formulation
Formulation of Pre-training Dataset and Hetcrogeneity. In this paper, we consider the multi-modal datasets with two modalities X,Y, e.g., image and text modalities. Following (Su al.) (003)), we assume sample of modality X and sample y of modality Y are generated from latent semantics through implicit
3

---

Published in Transactions on Machine Learning Research (06/2025)
Table 1: Comparison between related works and our proposed FedDRA for federated vision-language pre-
training. We organize the related works based on task settings and technical points:(1) Heterogeneous client
datasets. (2) Multi-modal paired datasets. (3) cnnrer-compuaation-io free. (4) No communication on training
(5ta.(5) Introduce global constraints.
(6) Distributionally robust framework.
Related Work
Task Settings
Related Technical Points
Hettero.
Multi-Modal
Server Comp.-Free
Feat. Commu.-Fre
Global Coust.
Dist. Robust
Zhuang ct al.
2021
Zhuang ct
2022
Zhang et sl. 2023
Yan et al.
2023
Lu et al.
2123
Ours
mappings that are consistent across all clients. For instance, disease labels of a given X-ray image and radiology-report pair, are latent semantic variables that connect these two modalities. That's because these labels determine the pathology region of the radiology image and corresponding description in the diagnosis report. In federated learning, we consider N clients, each has its own local dataset, collectively forming a group C, which represents the set of clients. We assume that each client has a corrosponding data distribution D E [I....,N], NI, and data samples are given as (z, y) ~ D,. In real-world scenarios, the distributions D, of local datasets vary across clients, introducing data heterogeneity that can negatively impact federated learning performance.
We aim to obtain generalizable model that performs well on the potential testing domain Dr. In real- world scenarios, testing datasets often exhibit domain shifts (Zhang et al. 2024b). with samples drawn from distributions distinct from the pre-training data. For example, a medical multi-modal model might be pre-trained on data from routine clinical practice and then transferred to tasks utilizing datasets collected during the COVID-19 pandemic. The overall data distribution D could be formed by client datasets D; with different contribution weights X; [l....... . when the entire data is grouped by C. Changing these weights alters the aggregated distribution DC However, these distributions are limited to the pre-training domains and cannot account for unseen cases that may arise during testing. Therefore, we introduce uncertainty set Q², which is distribution family building on the client data distribution while incorporating potential domain shifts. Qf consists of distributions that are close to DC, and can be formally written as: (Q) & p), where Dy is the f-divergence of two distributions. Here, E R+ is the uncertainty radius, and a larger P introduces more unseen distributions.
Pre-training. Given N clients and their local datasets, federated learning aims to utilize the elient dataset to train generalizable model in privacy-preserved way. Tt iteratively trains local models the client side and aggregates (e.g., FedAvg strategy simply averages model parameters) them on the server. For each communication lurn r, each client learns a local model through E update steps, and sends it to the server, and overwrites the local model with the aggregated model sent back from the server. Specifically, foderated multi-modal pro-training aims to effectively leverage paired and unlaboled multi-modal data from local clients to learn a generalizable model f.
Federated Vision-Languan
Multi-modal pre-training utilizes paired data from multiple modalities to learn model f that can well repre- sent data samples. In this paper, we consider the pre-training task image and text modalities. We focus a conventional setting in vision-language pre-training, where the model consists of two modality-speeiffe encoders for vision input X and text input Y, and well-aligned encoders should project their inputs into a shared representation space Z. For example, a good pre-trained model that can encode an image of a running dog and its text description "a photo of running dog" to a shared representation space Z, which is called cross-modal alignment in (Castrejon et al. 2016 Gao et al. 2024).. Suppose the quality of the representation space of pre-trained models can be measured by a loss objective R gddd + E (e.g., mutual information between representations X and Y with dimension d), effective federated multi-modal pre-training method should train a model that has a lower risk (( [R( (a, v))] on the potential testing domain DT-
In federated learning scenarios, varying weights can he reflected in the different update step sizes the sampling of different data batches.
4

---

Published in Transactions on Machine Learning Research (06/2025)
Server Estimate Worst-case Distribution
(2) Local Pre-training with Global Constraint
Serversside Operations
Worst-case ostimation Update weighting param,
Server averages of to
Vision up Encoder
aligner
Local model B.
Local model B
Server
A1
AM
0,2c
Global model
Id
data flow
Global $ model
Ri Rienn Eafe
loss flow
L
27 = Proj(
sarver op.
Global Repre
Grouped cliient distributions
clicnt op.
Client
APP
AuPa
Ø Encader
aligner 4
P
"plrunal offiriin worsetingg
P2
ApP2
Update 2 stage
Update in both stages
Worst-case distribution
Local triining loss: Liuest = Lier + Bl... O - 8, 17,, Voloooos
Figure 2: (1) Our proposed FedDRA method optimizes performance on a distribution family Qe built on client data distributions. After each round local Iraining, the server performs an additional step to obtain the weight X₂ for each client. The combination of A and the grouped client distributions constructs worst- case distribution from Q where the model is expected to perform the worst. The server then transmits À lo each client to adjust the local update step size.
(2) During local pre-training, the aggregggg models to obtain global representations 2 and zy, encouraging local models to learn generalizable patterns. Our method adopts a two-stage pre-training schema, which first trains the deep layers, and then jointly optimizes the entire model in the second stage.
use frozen copies of
In federated setting, f is aggregated from fr. which are learned by minimizing B(ra))) [R(f( (, v))l during local training. f, may capture clicnt-specific information that may not be generalizable across client dataset domains and will affect the performance of the aggregated f if local datasets are heterogencous,
Table I provides comparison of the most similar previous works, highlighting the distinctions between their tasks and ours, well as the technical differences between their approaches and ours.
4 Method
4.1 Global Constrained Local Training Objective
During local training, the pre-trained model would capture client-specific information that can not generalize other data domains, as shown in Sec. 5.2 Here, we will provide an in-depth analysis of this phenomenon and propose a global constraint term alleviate it.
In classical vision-language pre-training setting, the vision- model f is composed of two encoders, fo X' + Z for image modality X, and model fo y - 2 for text modality Y, as shown in Fig. 2 Given an image-text paired data (I, v), these models project the input to representations 2e = (x), by (v). The goal of vision-langung, pre-training is to learn robust cross-modal alignment from paired data, and thus obtain a generalizable representation space, where image representation Ze and text representations zy are well-aligned. It could be viewed as maximizing the mutual information of reprosentations of the two modalities (Su et al. 2023). Therefore, we can measure the cross-modal alignment degree with a mutual- information-based loss objective R, which is approximated by InfoNCE [Liu et al. 2021 Lu et al. 2024 loss in this paper.
Current multi-modal pre-training methods often encourage the model to maximize mutual information of the training pairs, neglecting the potential data heterogeneity problem in the federated learning scenarios. As the distribution D, varies across clients, each client datasct corresponds to distinct optimal model f.. which is induced from the distribution of the client dataset. For local training in client i, given that only (x, y) D; available, the locally learned encoders f tends to move towards fi, and might capture
5

---

Published in Transactions on Machine Learning Research (06/2025)
harmful elient-specific information. In federated pre-training, the model is expected to capture patterns that are transferable across clients and potential testing domains, and client-specific information might harm the model's generalization ability. Therefore, it is crucial to explicitly force the model to learn generalizable knowledge. Previous mothods such as FedAvg (McMahan ct al., 2017 Lu ct al. 2023). which do not account for this distinction, may result in learning biased local models fi, and diminish the generalization ability of the aggregated model f.
Our goal is to lederally learn a vision-language model that better aligns image and lext modalities. Let fo : X' Z be the encoder for the image modality X, and let fa : y + Z be the encoder for the text modality y. Suppose the federally learned encoders are denoted as fo : X + Z and so - + 2, and that for each data domain Di, there exist optimal encoders fv; and fa, and learned encoders fús and fo. Formally, let L denotes the InfoNCE, the generalization error on testing domain Dr can be written as RT( = D (x), fo(U)),, which is upper-bounded as stated in proposition
Proposition 1. Let (Di, Ja Vi and DTT be the distributions and optimal encoders for each client data domains and the testing domain, respectively. Given mized weights { E W/ = =,, wa 2 0, federally learned model f, and temperature T in InfoNCE loss. The generulization error RT (f) follows:
RT( S B( w [ (x) & (1) + &f (V) + Io, (v) + +c..
where Ci, ag are client-specific constants, and (,,a(x) = (l9(2) - hxx) denotes the distance of embeddings of a data sample x projected by 9(-) and h(.).
Here, fo (x) and Jo (2) mcasure the discrepancy between the locally trained models and the optimal models of the local data domain. These discrepancies are minimized during local training, leading the local models to learn client-specific information. Another two terms feefe, (x) and forfo, (x) capture the
discrepancy between the server-nggregated models and the locally trained models. Minimizing these terms can not only help reduce the negative effects of data heterogeneity, but also encourage the local models to learn generalizable patterns.
Motivated by this, we can directly minimize the terms 'foffe (x) and forfo, (a) to encourage the models to learn generalizable features. By projecting the inputs x and y through fo and fo. we obtain the global representations = and my, respectively. Therefore, the constraint loss can be defined as Le 1| =+ [ - =ill- The total loss objective for local training could be:
Liocal LMI + pLC.
(1)
where μ is the hyper-parameler that adjusts the constrained degree, LMt is the pre-training loss term based on infoNCE loss (e.g., image-text contrastive loss).
4.2 Two-Stage Alignment For Mitigating the Deeper Layer Bias
Furthermore, as discussed in Sec. 5.2 deep layers that contain biased client-specific information can impede pre-trained model to learn generalizable representations. This observation is similar to the findings in the supervised federated learning domain (Legate et al. 2024), where training from a better initialized last layer, can less capture biased information client. local datasets. Motivated by this, we model the deep layers of the encoding functions of modality X,Y Y as aligmment modules fu, fo, and aim to obtain generalizable alignment modules. In practice, we add additional blocks as the alignment module for simplicity, instead of dividing each encoder to two separate parts. For data flow of an input pair (z, y), the image encoder and text encoder fv. fo firstly take y to obtain intermidiate features žz, &y, respectively. Then the corresponding alignment modules fu. will project غz, žy to aligned representation ūz, Zy. which are finally used in computing Liocal-
To mitigate the negative impact of the biased alignment module on the generalization ability of the pre- trained model, we first adapt alignment modules fv and f4. During the adaptation, the encoders fwr fo initialized from weights pre-trained in the natural domain are kept frozen. Then we use the finetuned fy and fo to enhance the training of the feature encoders fo and fo- To encourage alignment. modules learn
6

---

Published in Transactions on Machine Learning Research (06/2025)
to extract general features, in the first stage, we train them with frozen feature encoders using the learning objective Eq.(1 ① Since these feature encoders are less biased from client-specific information, alignment modules are encouraged to learn less biased mappings from & to By and ε to Σz. Then in the second stage, we train both the alignment module and feature encoders to enhance their capability of extracting medical leatures, with the same learning objective as the first stage. The complete pipeline is illustrated Fig.
Algorithm FedDRA for Federated V-L Pre-training
A Algorithm 2 1st Stage Training
1: Input: Init. encoder param, o(0); Init. aligmment module
1: Input: Learning rate 1f Local iteration steps F;; Lo- cal Param, gr GrE; Global Param. grEe, GrEx; Weight A
param. 600;; Client weight ( fi=1s Uncertainty radius
Client datasets {Dr Local iterations E; Step-size 7
Set (0F8; (r = Rr
2: r r R-
1 do
3: Set (0r"x;0rRv (ypX*)== = F8"F, R)
3: Server broadcasts ((+R), (y(-R),(rrd ((r)), ((rE),
4: for t= :rr.....f (r.....(r..... I)E E do
4: for client i= I,.... N do
5: Sample Et from D, uniformly
5: Utilize Algorithm 2 on D, to get 0(++1)#
e
6: end for
Server computes: ((++1)B # EN 6 ((++
7: X1/VV, Ltocnt (0) 0r
7:
8: end for
8:
Server broadcasts (
9:
for client i I,..., N do
Utilize Algorithm& 3 on Dj to get 0f 0(r+1)E end for
10:
11:
Algorithm 2nd Stage Training
Server computes: (((++1)))= N > 2i=1 0 ((()+1)))
12:
1: Input: Learning rate Local iteration steps E; Lo- cal Param. gr Rr;; Global Param. gr Rr Weight d
13:
for client i= 1...,, do
Compute loss of (f++1) of model (0(r+1) ((r+1)))) D; end for
14:
15:
E 2: Set (0r#; ( ( 0(++1)8)
16:
3: Set (erB";e((+t))) = E 0(r+11))
17:
for client i= 1,...,
........ N do x(+1l= = Proj
\ (
(11
E iimi (10))oftt)
4: for t= i=rE...... 1)
do
18:
5: Sample e(0) from D, uniformly
19: end for
(0(+1,0(+)) = (0): 0)) - X
20:
end for
7: () Ltccai 0 0r 0r (((-1))=;;;))
21: return (8(RE), ((RE))
8: end for
9: The client sends ofr+1) of to the server
4.3 Learning Robust Modsl Alignment via Distributionally Robust Optimization
In real-world scenarios, Dy is unknown during pre-training and is typically different from the training data distribution, as demonstrated in Sec. 3 To address this issue, we assunee the distribution of the testing data close to that of the entire training data (Rahimian & Mehrotra 2019: Levy el al., 2020), and thus constructing a distribution family that covers potential testing distributions. As discussed in Sec. 3 the uncertainty set Q comprises distributions that lic within small distance of the training distribution, making it a suitable choice. We aim to obtain model that performs well on the whole set of potential testing distributions. Instead of optimizing the model parameter 0 on a distribution randomly selected from OC, we optimize 0 over a worst-case distribution from this set, where the pre-trained model performs the worst in aligning the two modalities drawn from it. Formally, this objective can be written as R(f) := supQeo= {T<..)>~ (0,
Inspired by this, we introduce the Distributtinnall--Robuss Ropusttpptpiiizzation (DRO) our federated multi-modal pre-training task to optimize such worst-case performance. DRO first constructs family of potential testing distributions º as shown in Sec. 3 and optimize the model's performance on the worst-case distribution, where the model performs the poorest among distributions in Ω However, during federated learning, the server has no access to the distribution of the entire data. Motivated by (Zhang et al. (024a), we introduce a de-centralized form of the DRO problem. The optimization object could be written as:
sup AEAcO
R(O,A)= = Σ (A.R((0ss... s.t. s,(C AII(1, I,.., ..... 1)) =p].
(2)
where R+(0): :- (( D, Liocat (0: (x. (2. is the empirical risk on client data D., p is the uncertain radius as mentioned in Sec.
7

---

Published in Transactions on Machine Learning Research (06/2025)
Then, we can optimize Eq. 2 by alternatively optimize the weights À and model parameters A. Specifically, we optimize the parameter 0, of local models on each iteration by o((+1) = - where rj the learning rate. Following the mirror gradient ascent of weight proposed in (Zhang et al., 2024a], we update
À with X Ei x iE [], where is the step-size hyper-parameter e4 is the loss calculated
on D.. Then we compute X+1 by projecting A into the set (A: D2 |I 1)ll S p} to fit the constraints of the uncertainty set. In practice, we update at the end of cach commmication turn.
We apply the proposed DRO in both the first stage of training the alignment module and the second stage of training the feature encoder. In both stages, we use the same objective, Liocal, to encourage the model to learn gencralizable information and mitigate the impact of data heterogeneity information ((z,, Fy). The key difference between the two stages is, in the first stage, the oplinization larget 0 in Eq. 2 corresponds to the parameters of the alignment modules & and V, whereas in the second stage, represents the parameters of the feature encoders Ø g6 and alignment modules $, V. The pseudo-code of the whole algorithm can be seen in Algorithm 1.
maximizing the mutual
5 Experiment
5.1
Experiment Setting
We focus on adapting medical vision-language pre-training methods to heterogeneous federated learning settings. We employ the framework of image-text contrastive learning with two modality-speciffc encoders, a fundamental design in multi-modal pre-training. We use vision-language pro-training tasks on Chest X-ray datasets and ophthalmology image datasets to evaluate the effectiveness of our FedDRA method.
5.1.1 Experiment Set-up of Pre-training on Chest X- X-Ray datasets
Bigolin Lanfredi et al. 2022) dataset for medical vision-language pre-training. Following Yan et al. 2023 we employ the Latent Dirichlet Allocation (LDA) (Blei et al., 2003) to divide the MIMIC-XRR datasct based on discase labels to construct 5 helerogeneous client datasets. We set the heterogeneily degree in the LDA algorithm to be 1. Each divided dataset consists of train splits and test splits based on the notation of the MIMIC-CXR We use the train split for pre-training, and test split to evaluate pro-trained model's image-text retrieval performance. Here, we only divide the raw data into subsets, because vision-language pre-training requires a large batchsize and data-consuming,, thus we need to guarantee each client has 10k to 50k paired data.
Pre-training setup. Following WWagg et al. 2022) we utilize the MIMIC-CXR
For main experiments, we set the number of communication turns to 25, and randomly sample 50 batches of data for local training at each turn. Here, we choose relatively small number of communccation turns compared to classical supervised federated learning. That's 's because VLP needs large local optimization steps per turn to learn cross-modal alignment.
Downstream tasks, Following WWagg et al. 2022; we conduct the following downstream tasks to evaluate the transferability and generalization ability of the pre-trained model. (1) Few-shot classification. We test their performance on multiple image classification benchmarks: RSNA Pneumonia Detection (RSNA) (Shih) et al. 2019), and Covidx (Wang et al.) 2020). We fine-tune pre-trained model with an additional linear layer 1%, 10% percentage of the training dataset, and evaluate the classification accuracy. (2) Medical image segmentation. We conduct medical image segmentation experiments on the RSNA (Wang et al. 2020) benchmark. We freeze the encoder and fine-tune a U-Net decoder using 1%, 10% of the training data, and then use the Dice score for evaluation. The datasets we have used for the fine-tuning based tasks are out-of-distribution, so that we can evaluate the transfcrability of the pre-trained model. (3) Image-text retrieval. We utilize the test splits of client datasets for evaluation, these datasets are unseen in pre-training, and can be viewed as in-domain samples. We report the top-1 and top-5 recall accuracy in data batch.
8

---

Published in Transactions on Machine Learning Research (06/2025)
5.1.2 Experiment Set-Up of Pre-training Ophthalmology Datasets
Pre-training setup. We conduct vision-language multi-modal pre-training using retinal image datasets from different institutes to simulate a more real-world sctting. These retinal datascts are from different in- stitutions of low-income and high-income countries, and are highly heterogeneous real-world scenes. Specif- ically, we utilize MESSIDOR (Decencière et al. 2014] from France and BRSET NNkayama et al. 2023 from Brazil as pre-training datasets, and assign them to two clients. These datasets include both images and tabular EHR records indicating Diabetic Retinopathy (DR) status and edema risk. For implementation, we transform these tabular data into text captions.
Downstream tasks. We evaluate the transferability of the models on few-shot classification tasks using the MBRSET (Wu et al., 2025 Nakayama al.) dataset. Unlike the pre-training datasets, MBRSET was collected by portable devices, resulting in significant distribution shift. We perform few-shot classification tasks diabetic retinopathy and edema status using this dataset. We fine-tune the model with an additional linear layer on 10,, 20% and 100% of the training data, and report classification accuracies.
5.1.3 Backbones and Baselines
We focus on enabling medical multi-modal pre-training methods to be applied to heterogeneous feder- ated learning scenes. We have considered the generalization ability of our method on different backbone VLP methods, and adopted contrastve-learning-
methods: simple language-image contrastive align- ment (ConVIRTT) (Zhang et al.] 2022 Radford al., 2021), global-local language-image contrastive align- ment (GLoRIA) (Huang et al.] 2021), and Multi-Granularity Cross-modal Alignment (MGCA) (Wang et al., 2022 All of the loss objectives of these pre-training methods contain contrastive loss term, which act the infoNCE loss to maximize the mutual information between two modalities. And we take this loss term for computing the client weights in the DRO part.
For baseline federated learning strategies, we have adapted FedMAE Yan et al. 2023). FedEMA Zhuang et al. 2022), FodMOON (Li et al., 2021b). FedX (Han et al. 2022). FodU (Zhuang et al.) 2021 FedL. DAWA (Rehman et al. 2023) for comparison. These are sell-supervised learning methods in the federated learning domain which also focus on tackling the data heterogeneity. For basic federated learning base- lines, we consider simple averaging (FedAvg) (McMahan et al. 2017], decentralized training, and centralized training. For baselines pre-trained in Local strategy, we report the averaged performance of the local models.
For fair comparisons, we re-implemented all baseline methods using the same backbones. To adapt uni-modal self-supervised learning baselines to our setting, we added an image-text contrastive loss, applying the same hyperparameterss as in our method for consistency. We use ViT-base (Dosovitskiy et al., 2020) as the vision encoder and Bert-base (Devlin et al. 2018) as the text encoder, with input pre-processing following (Wang et al. 2022)- Additionally, we employ extra transformer block from ViT-base and Bert-base the alignment module for vision and language, respectively. These modules are initialized from the checkpoint pre-trained on natural data.
5.2 Empirical Finding
In this section, we will demonstrate our key empirical findings for federated multi-modal pre-training under heterogencous client datascts, which actually motivates us to propose our FodDRA. We conduct experiments on the image-text retrieval lask, which reflects the ability maximize the mutual information and learn cross- modal alignment. In the following studies, mainly compare performances of naive federated (FedAvg) pre-trained model, decentralized pre-trained model, and centralized pre-trained model.
Federated learning enhances pre-training by leveraging more samples in a privacy-preservel manner, while data heterogeneity can affect the effectiveness of the FedAvg. Figure3 a) presents the retrieval accuracics of the models under consideration. Despite the heterogeneity of local datasets, the FedAvg strategy significantly outperforms the decentralized pre-training approach. However, the centralized pre-trained model remains an upper bound, indicating substantial room for improvement.
9

---

Published in Transactions on Machine Learning Research (06/2025)
(4), FL loverage diverse semantics in local elicats (b), Local training can leam harmful client-spocific info.
(e). Binsed deep layers hinder the alignment of encoders
00.0
25-
10
22.0
PedAig Contrrllto Enitlel
Fellag
1-5
Rec.g:1
1-5 Avg.
Rec/e01
10
Fed/og
Tedday
1-5 Ang.
Rec.065
Ang.
Rec.ge5
Trstial
local avp
Cleet3
Clean
Figure 3: (a). The comparison of retrieval acc. on each client denoted as {C,) of centralized, FodAvg, and averaged acc. of decentralized pre-trained models. (b). The performance of the server model after 25 com. turns and the averaged performance of corresponding client models after 25 and 26 com. turns, each client denoted [Client;]ā=r (c). The averaged acc. each client. We show the acc. of centralized and FedAvg pre-trained baselines, and de-centralized pre-trained models shown as {C retrained on the union of training splits of client datasets.
Local training can learn harmful client-specific information, degrading the performance of the pre-trained model. After several commuication turns, re- re-training the aggregated server model on local datasets may lead to a performance drop, as shown in Figure 3 b). We considered a set of models retrained on a server model with local datasets respectively. The server model learned through 25 communication turns. Compared to the starting server model, the averaged accuracy of local retrained models is significantly lower. This degradation may be because local training focus learning domain-speeific information in the late communication rounds, which would affect the aggregated model's overall performance.
Decentralized pre-trained deep layers can hinder the learning of a generalizable feature extrac- tor. We re-trained the first four shallow layers of the decentralized pre-trained model on the combined local datasots. While this led to some performance improvements, a significant gap still remains compared to the FedAvg pre-trained baselines, as shown in Figure 3 c(c). This gap indicates that the biased frozen deep layers prevent the model from learning more diverse semantics from the combined dataset. We hypothesize that these deep layers may contain biased, client-specific information, which obstructs the cross-modal alignment process. Our findings align with observations (Legate al. 2024 in supervised federated learning.
Overall, from empirical findings, we conclude that federated multimodal pre- tre-training is sensitive to data heterogeneity, and simply averaging local model weights does not solve this issue essentially. Furthermore, performance clasely tied to the generalization ability of the final layers in pre-trained models. Thus, we are motivated to propose our method.
5.3 Main Results
Our method learns robust and enriched cross-modal alignment and has better transfcrability. Table 2 has shown results of downstream lasks, here we utilize the ConVIRT as the backbone pre-training method. In the image-text retrievel task, both average and worst-client accuracies our method are higher than basclinc's, which means our model can capture more robust cross-client features. In the few-shot classification and segmentation, our method beats other baseline strategics on each task, which demonstrate the higher generalization ability of the representation space learned by method.
10

---

Published in Transactions on Machine Learning Research (06/2025)
Table 2: Downstream task performance. We report the few-shot classification accuracy on Covid and RSNA, the in-domain retricval accuracy, and the Dice score for segmentation on RSNA.
Strategy
Backbone
RSNA (cls.)
Covid (cls.)
RSNA (seg.)
In-domain Imago-Tex Text Retricval
1%
10%
1%
10%
1%
10%
Rec.01
Rec.05
Wst.@1
Wst.
FedEMA
ConVIRT
82.8
83.1
79.2
86.5
70.9
73.6
24.0
67.0
21.9
62.4
FedMOON
ConVIRT
82.5
83.2
77.8
89.2
69.0
71.3
27.8
70.9
25.3
67.2
FedAvg
ConVIRT
83.1
83.3
78.0
88.5
69.6
71.5
28.8
72.1
25.3
66.7
FedDRA (Ours)
ConVIRT
83.2
83.7
81.0
90.3
71.7
74.1
30.2
73.2
27.0
68.9
FedX
GLoRIA
82.7
83.4
78.3
88.5
71.0
72.1
28.5
72.2
25.9
68.0
FedU
GLoRIA
83.0
83.5
78.7
89.3
71.2
72.6
29.2
73.0
27.6
69.5
FedAvg
GL.RIA
83.2
83.3
77.5
89.0
71.4
72.4
29.9
73.8
27.8
69.5
FedDRA (Ours)
GLoRIA
83.6
84.1
79.4
89.8
72.0
73.2
31.1
74.3
28.2
70.2
FedLDAWA
MGCA
82.4
83.5
78.1
88.5
70.4
72.6
29.0
73.5
27.0
68.9
FedAvg
MGCA
82.6
83.5
75.8
88.2
70.1
71.4
29.3
73.7
26.8
70.4
FedDRA (Ours)
MGCA
83.1
83.8
79.3
89.1
71.0
72.8
29.8
74.1
27.4
70.6
Table 3: We have conducted ablation experiments to verify the effectiveness of our key technical designs.
We report the few-shot classification accuracy and the retrieval accuracy.
Two-stage
Global Constraint
DRO- Weighing
Covid (cls.)
RSNA (cls.)
In-domain Image=Text Retrieval
1%
10%
1%
10%
Rec.01
Rec.05
Wst.
Wst.@5
83.0
83.4
80.5
89.6
29.4
72.7
26.2
68.1
82.8
83.0
79.8
88.6
28.3
71.9
26.0
67.9
80.2
29.7
72.8
25.6
67.3
83.2
83.7
81.0
90.3
30.2
73.2
27.0
68.9
Table 2 has shown the performance of adapted sell-supervised federated learning methods which focus on single-modaligg From the results, we observe that baselines have shown better transferability on visual downstream tasks, compared to the naive FedAvg strategy. However, in the multi-modal retrieval task, our method beats these baselines by large margin, which indicates that previous single-modality methods cannot be easily adapted directly for multi-modal data. Furthermore, Fed is still competitive baseline in multi-modal retrieval tasks compared to other adapted methods, as we observed in the experiments. We conjecture that's 's because FedAvg only focuses on maximizing the in-domain mutual information, and doesn't 't introduce additional loss terms which would hurt the learning of enriched cross- modal alignment. However, this may lead lower generalization ability on few-shot downstream tasks as discussed before.
Our method can be transferred to multiple various multi-modal pre- training methods. Table 2 shows the downstream task performance of the MGCA and GLoRIA backbone pre- training methods when combined with our strategy. Our method has success- fully adapted MGCA and GLoRIA to the heterogeneous federated multi-modal pre-training scenario, demonstrated by the significant improvement in classifica- tion and segmentation tasks.
Table 4: Downstream task performance. We report the few- shot classification accuracy of Diabctic Retinopathy and Edema Risk classification lasks on the MBRSET dataset.
Strategy
Diabetic Retinopathy (cls.)
Risk of Edema (cls.)
10%
20%
100%
10%
20%
100%
Decentzalized
79.7
81.1
92.5
93.8
Red
79.4
S0.2
82.3
92.8
93.6
94.2
PsIMAE You et pl 2023
79.2
80.3
82.0
92.4
53.3
94.0
FedX
Han
81.6
93.5
RodU
Zhnsing of al. 721
79.7
80.5
81.7
92.8
93.4
94.1
RsADILA
(Ours
80.6
81.5
83.1
93.4
94.1
94.9
FodGlobal
94.7
95.8
5.4 Analysis Experiments
The two-stage pre-training strategy and global constraints can enhance the learning of cross- modal alignement. We remove the global constraint loss Lg of our method, and compare the pre-trained model's performances with those of the original version. As shown in Table the downstream performance, particularly the image-text retrieval accuracy, is significantly lower in the modified version. Similarly, we
11

---

Published in Transactions on Machine Learning Research (06/2025)
#a02 Rand
#ram #ramd) araml
Raurd
90
CI Sin Tex Ratio
0.45
0.45
CI5 loc Iallo
0.55
0.35
70
0.25
1125
05
9.15
0.15
85
3.5
1485
50
(b) Analysis study on uncertainty radius. (c) Analysis study global constraint degree. (d) Improvement of image-text embedding similarity per.commu turn. (e) Image-text embedding similarity curve of our method. (f) Averaged similarity curves.
Figure 4: (a) Analysis study on number of clients.
remove the first-stage pre-training on aligument modules to verify the role of the two-stage pre-training strategy. We have found that the first stage pre-training can help learn better cross-moda alignment and achieve high image-text retrieval accuracies, as shown in Table 3
DRO weighing can reduce the domain gap and improve the downstream performances. We remove the DRO weighing part and compare the performances of the model pre-trained with the original method. As shown in Table 3 removing the DRO-weighing part leads to a large performance drop in few- shot classification performances, thus adding DRO-wigghing component can improve the transferability of the pre-trained model. The client-wise worst accuracies of the original model are much higher than those of the modified version. This means DRO has succesfully bridged the gap between local training data and downstream dataset, by optimizing model performance on the constructed uncertain set of distributions.
Table 5: Results on client datasets with different heterogeneity degrees, across different pre-training methods
Strategy
a Backbone
RSNA (cls.)
Covid (cls.)
RSNA (seg.)
Backbone
RSNA (cls.)
Covid (cls.)
RSNA (seg.)
1%
10%
1%
10%
10%
1%
10%
10%
1%
10%
FedAvg FedDRA (ours)
ConVIRT
82.2
83.3
78.0
88.5
69.6
71.5
GLoRIA
83.2
83.8
78.5
89.0
71.4
72.4
ConVIRT
83.2
83.7
81.0
90.3
71.7
71.1
GLeRIA
81.8
82.9
78.5
89.6
69.3
72.5
FedAvg
ConVIRT
80.7
82.3
77.6
88.2
68.5
71.7
GLoRIA
81.8
52.9
78.5
80.6
69.3
72.5
FedDRA (ours)
ConVIRT
81.8
82.9
78.5
89.6
69.3
72.5
GLoRIA
82.2
83.0
78.4
89.4
71.5
72.9
Contralized
ConVIRT
83.4 84.6
82.5
02.0
72.6
76.4
GLoRIA
84.0
84.7
82.2
91.8
73.5
73.7
FedDRA dynamically schedules updating stepsize for each client, and therefore optimizes the worst-case performance. We select two clients C1, C2 in the federated pre-training. For each client, we calculate the average cosine similarity between image and text embeddings, using the server-aggregated model at each communication turn. In Fig. d( and Fig. [ee, we plot curves of these similarities, which reflect the cross-modal alignment degree. At each communication turn, when the similarity of a client is relatively high, its similarity in next turn would get a smaller improvement. That's because our FedDRA can assign higher updating stepsize to client where the cross-modal alignment is less extracted by the model.
Our FedDRA can alleviate over-fitting client-specific information, and learn better cross-modal alignment For our FodDRA and the FodAvg method, we plot the similariiis of text and image embedding averaged across clients at each communication turn. As shown Fig [f) Fed requires fewer communication rounds to converge, but results in fluctuations after certain communication turns. This aligns with findings in Sec..22 where local retraining model that are trained after multiple communication rounds, would introduce harmful client-specific information and distort the learned representation space. In contrast, our FedDRA gradually extract cross-modal alignment from local training in distributionally robust manner, and learns a stronger representation space.
Analysis on global constraint ated pro-training to enhance performance on less optimized client data domains, leading to smaller disparity on image-text retrieval performance on each client domain. As we increase μ from 1 to 5, the downstream performance consistently increases. However, excessively large μ can decrease overall performance.
μ.. As shown in Fig. 4c), a larger i encourages feder-
12

---

Published in Transactions on Machine Learning Research (06/2025)
A larger uncertainty radius p improves transferability in downstream tasks. Fig. (b) shows the downstream performance of models pre-trained with different uncertainty radii in the DRO process. As larger would bring higher performance in few-shot classification and segmentation tasks on out-of-domain datascts. We also observed that a smaller better supports cross-modal alignment learning, achicving belter image-lext retrieval performance on in-domain datasets, as shown in Table [ in the Appendix. This is because the larger uncertainty radius would incorporate more potential out-of-distribution cases, which can enhance the model's transferability.
Robust check heterogeneity degree of client datasets. We changed the which adjusts the heterogeneity degree of the LDA allocated client datasets, to check the robustness of our method under different heterogeneity degree. As shown in Table our method consistently enhance pre-training methods performances under client datasets with different heterogeneity degrees.
Robust check on number of clients. We adjusted the number of clients involved federated pre-training. As shown in Fig. [a), increasing the number of clients introduces greater diversily, which can enhance the downstream performance of the pre-trained model.
6 Conclusion
Data limitation is long-standing problem in the multi-modal learning domain. Despite federated learning leveraging datasets from multiple sources while guarantecing privacy issues, its performance would be damaged by data heterogeneity. Inspired by our empirical findings on the impact of heterogeneily on federated multi-modal learning, we propose the FedDRA framework to mitigate heterogeneity for federated medical vision-language pre-training. The effectiveness of our method has been verified by comprchensive experiments. While introducing representation from other clients might bring larger improvement, we still consider the most privacy-preserve setting where representations are not transmissible. Further work could explore how to generalize to more diverse multi-modal pre-training data in federated multi-modal learning, while preserving local data privacy.
7 Acknowledgement
The authors acknowledge support from U-M MIDAS PODS Grant and U-M MICDE Catalyst Grant, and computing resource support from NSF ACCESS Program. This work used NCSA Delta GPU through alloca- tion CIS230133 and ELE230011 from the Advanced Cyberinfrastructure Coordination Ecosystem: Services & Support (ACCESS) program, which is supported by National Science Foundation grants 2138259, 2138286, 2138307, 2137603, and 2138296.
13

---

Published in Transactions on Machine Learning Research (06/2025)
References
Shruthi Bannur, Stephanie Hyland, Qianchu Liu, Fernando Perez-Garcia, Maximilian Nsse, Daniel C Castro, Benedikt Boocking, Harshita Sharma, Kenza Bouzid, Anja Thieme, ct al. Learning to exploit temporal structure for biomedical vision-language processing. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 15016 15016-15027, 2023.
Hangbo Bao, Wenhui Wang, Li Dong, Qiang Liu, Owais Khan Mohammed, Kriti Aggarwal, Subhojii Som, Songhao Piao, and Furu Wei. Vlmo: Unified vision-language pre-training with mixture-of-modalite experts. Advances in Neural Information Processing Systems, 35:32897 32912, 2022.
Ricardo Bigolin Lanfredi, Mingyuan Zhang, William F Auffermann, Jessica Chan, Phuong-Anh T Duong, Vivek Srikumar, Trafton Drew, Joyce D Schroeder, and Tolga Tasdizen. Reflacx, dataset of reports and eye-tracking data for localization of abnormalities in chest x-rays. Scientific data, 9(1):350, 2022.
David M Blei, Andrew Y Ng, and Michael Jordan. Latent dirichlet allocation. Journal of machine Learning research, 3(Jan):993 1022, 2003.
Giacomo Capitani, Federico Bolelli, Angelo Porrello, Simone Calderara, and Elisa Ficarra. Clusterfix: A cluster-based debiasing approach without protected-group supervision. In Proceedings of the IEEE/CVP Winter Conference on Applications of Compuler Vision, pp. 4870 4879, 2024.
Lluis Castrejon, Yusuf Aytar, Carl Vondrick, Hamed Pirsiavash, and Antonio Torralba. Learning aligned cross-modal representations from weakly aligued data. In Proceedings of the IEEE conference on compuler vision and pattern recognition, pp. 2940 2949, 2016.
Liqun Chen, Zhe Gan, Yu Cheng, Linjie Li, Lawrence Carin, and Jingjing Liu. Graph optimal transport for cross-domain alignment. In International Conference on Machine Learning, pp. 1542 1553, PMLR, 2020.
Zhihong Chen, Yuhao Du, Jinpeng Hu, Yang Liu, Guanbin Li, Xiang Wan, and Tsung- Hui Chang. Multi- modal masked autoencoders for medical ssiion-and-aanguage pre-training In International Conference on Medical Image Computing and Assisle Intervention, pp. 679 689. Springer, 2022.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. In International conference on machine learning, pp. 2089 2099. PMLR, 2021.
Etienne Decencière, Xiwei Zhang, Guy Cazuguel, Bruno Lay, Béatrice Cochener, Caroline Trone, Philippe Gain, John-RRcchate Varela, Pascale Massin, Ali Erginay, et al. Feedback a publicly distributed image database: the messidor database. Image Analysis & Stercology, pp. 231-234, 2014.
Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Distributionally robust federated aver- aging. Advances in neural information processing systems, 33:15111 15122, 2020.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. urXiv preprint xrXiv:1810.04800 2018
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa. Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth Transformers for image recognition at scale. arXiv preprint arXiv:2010. 11929, 2020.
16x16 words:
Yuting Gao, Jinfeng Liu, Zihan Xu, Tong Wu, Enwei Zhang, Ke Li, Jie Yang, Wei Liu, and Xing Sun. Softelip: Softer cross-modal alignment makes clip stronger. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 1860 1868, 2024.
Avishek Ghosh, Justin Hong, Dong Yin, and Kannan Ramchandran. Robust federated learning in hetero- geneous environment. arXiv preprinl ar Xix:: 1906.06629,, 2019.
Peixuan Han, Zhenghao Liu, Zhiyuan Liu, and Chenyan Xiong. Distributionally robust unsupervised dense retrieval training web graphs. orXiv preprint arXiv:2300. 2023.
14

---

Published in Transactions on Machine Learning Research (06/2025)
Sungwon Han, Sungwon Park, Fangzhao Wu, Sundong Kim, Chuhan Wu, Xing Xie, and Meeyoung Cha. Fedx: Unsupervised federated learning with cross knowledge distillation. In Europeon Conference Computer Vision, pp. 691-701 707. Springer, 2022.
Shih-Cheng Huang, Liyue Shen, Matthew P Lungren, and Serena Yeung, Gloria: multimodal global- local representation learning framework for label-efficient medical image recognition. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp.
2021.
Wenke Huang, Mang Ye, and Bo Du. Learn from others and be yourself in heterogeneous federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pottern Recognition, pp. 10143- 10153, 2022.
Colton Ladbury, Arya Amini, Ameish Govindarajan, Isa Mambetsariev, Dan J Raz, Erminia Massarelli, Terence Williams, Andrei Rodin, and Ravi Salgia. Integration of artificial intelligence lung cancer: Rise of the machine. Cell Reports Medicine, 2023.
Gwen Legate, Nicolas Bernier, Lucas Pagc-Caccia, Edouard Oyallon, and Eugene Belilovsky. Guiding the last layer in federated learning with pre-trained models. Advances in Neural Information Processing Systems, 36, 2024.
Daniel Levy, Yair Carmon, John C Duchi, and Aaron Sidford. Large-scale methods for distributionally robust optimization. Advances in Neural Information Processing Systems, 33:8847 8860, 2020.
Daliang Li and Jumpu Wang. Fedmd: Heterogenous federated learning via model distillation. arXiv preprint arXiv:1910.0880 2019.
Jingtao Li, Lingjuan Lyu, Daisuke Iso, Chaitali Chakrabarti, and Michael Spranger. Mocosfl: enabling cross-client collaborative self-supervised learning. In The Eleventh International Conference on Learning Representations, 2022a.
Junnan Li, Ramprasaath Selvaraju, Akhilesh Gotmare, Shafiq Joty, Caiming Xiong, and Steven Chu Hong Hoi. Align before fuse: Vision and language representation learning with momentum distillation. Advances in neural information processing systems, 34:9694-9700 2021a.
Junnan Li, Dongxu Li, Caiming Xiong, and Steven Hoi. Blip: Bootstrapping language-image pre-training for unified vision-language understanding and generation. In International Conference on Machine Learning, pp. 12888 12900. PMLR, 2022b.
Qinbin Li, Bingsheng He, and Dawn Song. Model-cont rastive federated learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10713 10722, 2021b.
Chen Liu, Yanwei Fu, Chengming Xu, Siqian Yang, Jilin Li, Chengjie Wang, and Li Zhang. Learning few-shot embedding model with contrastive learning, In Proceedings of the AAAI conference artificial intelligence, volume 35, pp. 8635-8643, 2021.
Jiashuo Liu, Zheyan Shen, Peng Cui, Linjun Zhou, Kun Kuang, and Bo Li. Distributionally robust learning with stable adversarial training. IEEE Transactions on Knowledge and Data Engineering, 2022.
Siyu Lu, Zheng Liu, Tianlin Liu, and Wangchunshu Zhou. Scaling-up up medical aaston-and-aanguae repre- sentation learning with federated learning. Engineering Applications of Artificial Intelligence, 126:107037, 2023.
Yiwei Lu, Guojun Zhang, Sun Sun, Hongyu Guo, and Yaoliang Yu. f-micl: Understanding and generalizing infonce-based contrastive learning. arXiv preprint XrXiv:2402.100 10150, 2024.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communicttion-offon learning of deep networks from decentralized data. In Artificial intelligence and slatistics, pp. 1273-12 1282. PMLR, 2017.
15

---

Published in Transactions on Machine Learning Research (06/2025)
Paolo Mendogni, Jacopo Vannucci, Marco Ghisalberti, Marco Anile, Beatrice Aramini, Maria Teresa Con- gedo, Mario Nosotti, Luca Bertolaccini, on behalf of the Italian Society for Thoracic Surgery (endorsed by the Italian Ministry of Health) Collaborators of the Pneumothorax Working Group Collaborators of the Pneumothorax Working Group, Ambra Enrica D' D'Ambrosio, et al. Epidemiology and management of pri- mary apontaneous pneumothorax: a systematic review. Interuclive curdiovascular and thoracic surgery, 30(3):337-345, 2020.
Luis Filipe Nakayama, Mariana Goncalves, T. Zago Ribeiro, Helen Santos, Daniel Ferraz, Fernando Malerbi, Leo Anthony Celi, and Caio Regatieri. A brazilian multilabel ophthalmological dataset (brset), PhysioNet hllps://doi. org/10, 13026, 2023.
Luis Filipe Nakayama et al. mbrset, a mobile brazilian retinal dataset.
Maxime Oquab, Timothée Darcet, Th6o Moutakanni, Huy Vo, Mare Szafranice, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint arXiv:2304.07198 2023.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8748-8763. PMLR, 2021.
Hamed Rahimian and Sanjay Mchrotra. Distributionally robust optimization: A review. arXiv preprint arXiv:1908.0560 2019.
Yasar Abbas Ur Rehman, Yan Gao, Pedro Porto Buarque De Gusmão, Mina Alibcigi, Jiajun Shen, and Nicholas D Lane. L-dawa: Layer-wise divergence aware weight aggregation in federated self-supervised visual representation learning. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 16464-16473, 2023.
Khaled Saab, Sarah Hooper, Mayee Chen, Michael Zhang, Daniel Rubin, and Christopher Ré. Reducing reliance spurious features in medical image classification with spatial specificity. In Machine Learning for Healtheare Conference, pp. 000-784. PMLR, 2022.
George Shih, Carol C Wu, Safwan Halabi, Mare D Kohli, Luciano M Prevedello, Tessa S Cook, Arjun Sharma, Judith K Amorosa, Veronica Arteaga, Maya Calperin-Aizenber et al. Augmenting the national institutes of health chest radiograph dataset with expert annotations of possible pneumonia. Radiology: Artificial Intelligence, 111)=e18004, 2019.
Weijie Su, Xizhou Zhu, Chenxin Tao, Lewei Lu, Bin Li, Gao Huang, Yu Qiao, Xiaogang Wang, Jie Zhou, and Jifenn Dai. Towards all-in-one pre-training via maximizing multi-modal mutual information. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11888-15899, 2023.
Yunze Tong, Junkun Yuan, Min Zhang, Didi Zhu, Keli Zhang, Fei Wu, and Kun Kuang. Quantitatively measuring and contrastively exploring heterogeneity for domain generalization. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, KDD '23. ACM, August 2023. doi: 00555/5550005333599555 URL htttp://dii.org/ 0tg/10.:/00 1145/3580355 /5509/9988830533998811
Fuying Wang, Yuyin Zhou, Shujun Wang, Varut Vardhanabhuti, and Lequan Yu. Multi-granularity cross- modal alignment for generalized medical visual representation learning. Advances in Neural Information Processing Systems, 35:33536-3335 -35449, 2022.
Linda Wang, Zhong Qiu Lin, and Alexander Wong. Covid-net: tailored deep convolutional neural network design for detection of covid-19 cases from chest x-ray images. Scientific reports, 10(1):19549, 2020.
Chenwei WU, David Restrepo, Luis Filipe Nakayama, Lucas Zago Ribeiro, Zitao Shuai, Nathan Santos Barboza, Maria Luiza Vieira Sousa, Raul Dias Fitterman, Alexandre Durao Alves Percira, Caio Vini- cius Saito Regatieri, et al. portable retina fundus photos dataset for clinical, demographic, and diabetic retinopathy prediction. Scientific Data, 12(1):323, 2025.
16

---

Published in Transactions on Machine Learning Research (06/2025)
Rui Yan, Liangqiong Qu, Qingyue Wei, Shih-Cheng Huang, Liyue Shen, Daniel Rubin, Lei Xing, and Yuyin Zhou. Tabel-efficient sell-superviedd federated learning for tackling data heterogeneity in medical imaging. IEEE Transactions on Medical Imaging, 2023.
Fengda Zhang, Kun Kuang, Long Chen, Zhaoyang You, Tao Shen, Jun Xiao, Yin Zhang, Chao Wu, Fei Wu, Yucting Zhuang, al. Federated unsupervised representation learning. Frontiers of Information Techmology & Electromic Engineering, 24(8):1181-IH -1993 2023.
Fengda Zhang, Zitao Shuai, Kun Kuang, Fei Wu, Yucting Zhuang, and Jun Xiao. Unified fair federated learning for digital healthcare. Pallerns, 5(1), 2024a.
Jianqing Zhang, Yang Hua, Jian Cao, Hao Wang, Tao Song, Zhengui Xue, Ruhui Ma, and Haibing Guan. Eliminating domain bias for federated learning in representation space. Advances in Neural Information Processing Systems, 36, 2024b.
Yilan Zhang, Yingxue Xu, Jianqi Chen, Fengying Xie, and Hao Chen. Prototypical information bottlenecking and disentangling for multimodal cancer survival prediction. orXiv preprint arXiv:2401.016/ 2024c.
Yuhao Zhang, Hang Jiang, Yasuhide Miura, Christopher D Manning, and Curtis P Langlotz. Contrastive learning of medical visual representations from paired images and text. In Machine Learning for Healthcare Conference, pp. 2-25. PMLR, 2022.
Hong- Yu Zhou, Chenyu Lian, Liansheng Wang, and Yizhou Yu. Advancing radiograph representation learn- ing with masked record modeling. aiv preprint arXiv:2301.00 131155 2023.
Weiming Zhuang, Xin Gan, Yonggang Wen, Shuai Zhang, and Shuai Yi. Collaborative unsupervised visual representation learning from decentralized data. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 4912-4921, 2021.
Weiming Zhuang, Yonggang Wen, and Shuai Zhang, Divergence-aware federated self-supervised learning. arXiv preprint arXix:2204.044 2022.
17


Reviews:
[{"review_id": "WrENOBDERf", "reviewer": "Reviewer_osd7", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n- The paper addresses a significant and timely problem in medical VLP: data heterogeneity in federated learning settings.  This is a relatively unexplored area with high practical relevance for real-world medical applications.    \n- The proposed FedDRA method appears technically sound, employing a distributionally robust optimization framework and a two-stage training approach to address the challenges of data heterogeneity and overfitting.    \n- The paper provides extensive experiments on real-world medical datasets, demonstrating the effectiveness of FedDRA in enhancing medical federated VLP under data heterogeneity.  The results show consistent improvements over baseline methods. \n\n**Weaknesses**\n- Clarity and Presentation: The paper could benefit from improved clarity and presentation. Some technical details are not clearly explained, and the figures could be more informative. More specific examples and clearer illustrations would enhance the paper's readability and impact.\n- Limited Scope: The paper focuses on a specific type of VLP (image-text contrastive learning) and a limited set of medical applications. The generalizability of the proposed method to other VLP approaches and medical domains is not fully explored.\n- Computational Cost: The paper does not provide a detailed analysis of the computational cost of FedDRA. The two-stage training approach and the distributionally robust optimization framework could potentially increase the computational burden, especially for resource-constrained clients.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes FedDRA (Federated Distributionally Robust Alignment), a VLP (Vision-Language Pre-training) framework for medical applications that uses federated learning to scale up datasets while preserving privacy. FedDRA addresses data heterogeneity challenges in real-world scenarios by constructing a distribution family encompassing potential test-time domains and using a distributionally robust framework. The framework includes a two-stage approach to avoid overfitting on client-specific information, first tuning deeper layers before updating the entire network. The server optimizes the final model towards the target set distribution, with mirror gradient descent to set the according weights $\\lamba_i$ for training the local modules.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Can the authors specifically address the following?\n- Address the computational cost that may be significant compared to the other baselines such as FedAvg\n- Can the authors comment in the paper on the generalizability of this method to other domains?\n- Can the authors also address on cases when the target set distribution is unknown (which is often the case in FL) and how effective the proposed method may be in this case?\n- It would be nice if the paper was reorganized, especially the proposed method part since the method is quite complex, and reading through the paper was quite difficult with many things to parse.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "EVZUlGMqvb", "reviewer": "Reviewer_ecgm", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The experiments are extensive. It covers classification, segmentation and image retrieval task. The proposed design choices seem reasonable and are backed by proper ablation studies.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work addresses data-heterogeneity challenge in federated learning for medical multi-modal learning task. It proposes a two-stage global-guided local training strategy along-with Distributionally-Robust-Optimization (DRO) weighting to learn robust and enriched\ncross-modal alignment, and better generalization.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The performance gain across all the tasks (even on multi-modal task) with the other baseline methods (table 2) seem really small. Have the authors done multiple runs to report the standard error, it is important to see if the performance improvement is significant as compared to other methods such as FedAvg. \n\nEven in ablation studies table, it is difficult to observe the gains,  for example absence of two-stage component consistently shows less than 1 % performance drop.\n\nAuthors should also perform a computational overhead (training time) comparison with the other baselines.", "broader_impact_concerns": "No ethical concern.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Pd0WZbpJgS", "reviewer": "Reviewer_QCdX", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength\n\nThe author proposes a novel federated learning method for training vision and language pretraining model in the server client setting.\n\nThe author experimentally demonstrates the ineffectiveness of the existing average model under heterogeneity. The proposed method shows better accuracy than the methods in the most settings.\n\nThe author conducts ablation study to show the effectiveness of the proposed modules. Further hyperparameter analysis is conducted.\n\n\nWeakness\n\nThe paper says that DRO loss is used In both the first stage and second stage, but it looks like the loss is different between Algorithm 2 and Algorithm 3. I think it would be better to match the losses in the Algorithms if they are the same. Further, I think it would be better to include the DRO process in the Algorithm since it is a little complex and involves server-client communication.\n\nIn some case, the performance change is the order of 0.1% accuracy. I would like to see the standard deviation of the scores to verify the significance of the result.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a novel federated learning method for vision and language pretraining aimed for the medical application. To make the learned model robust to the unknown test domain, the author proposes to (i) separate the model into encoder and aligner, (ii) constrain the local weight using the global model and (iii) train them in the two-stage training setting and train the model with robust loss. The author experimentally demonstrates the effectiveness of the proposed method on the medical multi-modal data.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Modification of the Algorithms would help increase the clarity of the paper.\n\nAdding the score deviation would be important to verify the significance of the score increase.", "broader_impact_concerns": "Federated learning is aimed for the privacy-preserving model training. Emphasizing that the information communicated between server and clients does not harm to the client privacy would help user adopt the method with safety.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This submission proposes Federated Distributionally Robust Alignment (FedDRA), a Vision-Language Pre-training framework for medical applications that uses federated learning to scale up datasets while preserving privacy. The proposed method addresses data heterogeneity challenges in real-world scenarios by constructing a distribution family encompassing potential test-time domains and using a distributionally robust framework. The framework includes a two-stage approach to avoid overfitting on client-specific information. 

In consideration of the above contributions, all the reviewers are basically positive to this work, which are strengthened by the authors' replies on concerns about clarity, presentation, computational cost, DRO loss, and marginal accuracy gain. The paper can be further improved by minor revisions based on the reviewers' comments.
Decision: Accept (poster)