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

During local training, the pre-trained model would capture client-specific information that can not generalize to 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 to alleviate it.

In classical vision-language pre-training setting, the vision-language model $f$ is composed of two encoders, $f_{\psi}:\mathcal{X}\to\mathcal{Z}$ for image modality $X$, and model $f_{\phi}:\mathcal{Y}\to\mathcal{Z}$ for text modality $Y$. Given an image-text paired data $(x,y)$, these models project the input to representations $z_{x}=f_{\psi}(x),z_{y}=f_{\phi}(y)$. The goal of vision-language pre-training is to learn a cross-modal alignment from unpaired data, and thus obtain a generalizable representation space, where image representation $z_{x}$ and text representations $z_{y}$ are well-aligned. It could be viewed as maximizing the mutual information of representations of the two modalities (Su et al., 2023). Therefore, we can measure the cross-modal alignment degree with a loss mutual-information-based loss objective $\mathcal{R}$, which is approximated by InfoNCE (Liu et al., 2021; Lu et al., 2024) in this paper.

Current multi-modal pre-training methods often encourage the model to maximize mutual information $I(z_{x},z_{y})$ of the training pairs, neglecting the potential data heterogeneity problem in the federated learning scenarios. As the distribution $D_{i}$ varies across clients, each client dataset corresponds to a distinct optimal model $f_{i}$, which is induced from the distribution of the client dataset. For local training in client $i$, given that only $(x,y)\sim\mathcal{D}_{i}$ are available, the locally learned encoders $\hat{f_{i}}$ tends to move towards $f_{\psi_{i}}$ and $f_{\phi_{i}}$, and might capture some harmful client-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 methods such as FedAvg (McMahan et al., 2017; Lu et al., 2023), which do not account for this distinction, may result in learning biased local models $\hat{f_{i}}$, and diminish the generalization ability of the aggregated model $\hat{f}$.

Given the distribution $D_{\mathcal{T}}$ of testing data domain, we aim to minimize the generalization error of the federally learned model on $D_{\mathcal{T}}$. Formally, let $\mathcal{F}_{\psi}\subseteq\{f_{\psi}:\mathcal{X}\rightarrow\mathcal{Z}\}$ to be a hypothesis space defined on input space of image modality $X$ and $\mathcal{F}_{\phi}\subseteq\{f_{\phi}:\mathcal{Y}\rightarrow\mathcal{Z}\}$ of text modality $\mathcal{Y}$. Suppose encoders $\hat{f_{\psi}}\in\mathcal{F}_{\psi}$ and $\hat{f_{\phi}}\in\mathcal{F}_{\phi}$ are the federally learned models, and there exist optimal $f_{\psi_{i}}\in\mathcal{F}_{\psi}$ and $f_{\phi_{i}}\in\mathcal{F}_{\phi}$ for each data domain $D_{i}$. Denote $\epsilon_{g,h}(x)=\|g(x)-h(x)\|^{2}_{2}$ to be the error of two models $g(\cdot),h(\cdot)$ on data sample $x$, $\mathcal{L}$ as the InfoNCE loss. We have an upper bound of the generalization error $\mathcal{R}_{T}(\hat{f})=\mathbb{E}_{(x,y)\sim D_{\mathcal{T}}}[\mathcal{L}(\hat{f_{\psi}}(x),\hat{f_{\psi}}(y))]$.

where $C_{i},\alpha_{i}$ are client-specific constants.

Here, $\epsilon_{f_{\psi_{i}},\hat{f_{\psi_{i}}}}(x)$ and $\epsilon_{f_{\phi_{i}},\hat{f_{\phi_{i}}}}(x)$ measure 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 $\epsilon_{\hat{f_{\psi}},\hat{f_{\psi_{i}}}}(x)$ and $\epsilon_{\hat{f_{\phi}},\hat{f_{\phi_{i}}}}(x)$ capture the discrepancy between the server-aggregated models and the locally trained models. Minimizing these terms can not only help reduce the upper bound of the generalization error, but also encourage the local models to learn patterns that generalize well to unseen testing data domains.

Motivated by this, we can directly optimize the terms $\epsilon_{\hat{f_{\psi}},\hat{f_{\psi_{i}}}}(x)$ and $\epsilon_{\hat{f_{\phi}},\hat{f_{\phi_{i}}}}(x)$ to encourage the models to learn generalizable features. By projecting the inputs $x$ and $y$ through $\hat{f_{\psi}}$ and $\hat{f_{\phi}}$, we obtain the global representations $z^{*}_{x}$ and $z^{*}_{y}$, respectively. Therefore, the constraint loss can be defined as $L_{C}=||z_{x}-z^{*}_{x}||^{2}+||z_{y}-z^{*}_{y}||$. The total loss objective for local training could be:

where $\mu$ is the hyper-parameter that adjusts the constrained degree, $L_{MI}$ is the pre-training loss term based on infoNCE loss (e.g., image-text contrastive loss).

Furthermore, as discussed in Sec. 5.2, deeper 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 in client local datasets. Motivated by this, we model the deep layers of the encoding functions of modality $X,Y$ as alignment modules $f_{\Psi},f_{\Phi}$, 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 a data flow of an input pair $(x,y)$, the image encoder and text encoder $f_{\psi},f_{\phi}$ firstly take $x,y$ to obtain intermidiate features $\tilde{z}_{x},\tilde{z}_{y}$, respectively. Then the corresponding alignment modules $f_{\Psi},f_{\Phi}$ will project $\tilde{z}_{x},\tilde{z}_{y}$ to aligned final representation $z_{x},z_{y}$.

To mitigate the negative impact of the biased alignment module on the generalization ability of the pre-trained model, we first train generalizable alignment modules $f_{\Psi}$ and $f_{\Phi}$, and then use them to enhance the training of the feature encoders $f_{\psi}$ and $f_{\phi}$. To encourage alignment modules learn 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 unbiased mappings from $\tilde{z}_{y}$ to $z_{y}$ and $\tilde{z}_{x}$ to $z_{x}$. Then in the second stage, we train both the alignment module and feature encoders to enhance their capability of extracting medical features, with the same learning objective as the first stage. The complete pipeline is illustrated in Fig. 2.

In real-world scenarios, $\mathcal{D}_{\mathcal{T}}$ is unknown during pre-training and is typically out-of-distribution. A common approach to address this issue is to assume the distribution of the testing data domain is near the distribution of overall training data (Rahimian & Mehrotra, 2019; Levy et al., 2020), and construct a set $\mathcal{Q}^{\mathcal{C}}$ that covers potential testing distributions. By optimizing the model’s parameter $\theta$ over this set with the loss objective $\mathcal{R}(\theta):=\sup_{Q\in\mathcal{Q}^{\mathcal{C}}}\left\{\mathbb{E}_{(x,y)\sim Q}\left[L_{\text{local}}(\theta,x,y)\right]\right\}$, we can pre-train a model that generalizes well to the whole set of potential testing distributions. Here, the maximum loss implies a worst-case distribution in $\mathcal{Q}^{\mathcal{C}}$, where the pre-trained model performs the worst in aligning the two modalities.

Inspired by this, we aim to optimize the loss objective on the worst-case distribution, and introduce the Distributionally-Robust-Optimization (DRO) to our federated multi-modal pre-training task. DRO first construct a family of testing distributions $\mathcal{Q}^{\mathcal{C}}$ 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 $\mathcal{Q}^{\mathcal{C}}$. However, during federated learning, the server has no access to the distribution of the entire data. Motivated by (Zhang et al., 2023), we introduce a de-centralized form of the DRO problem. The optimization object could be written as:

where $R_{i}(\theta):=\mathbb{E}_{(x,y)\sim D_{i}}[L_{local}(\theta;(x,y))]$ is the empirical risk on client data $D_{i}$, $\rho$ is the uncertain radius as mentioned in Sec. 3.

Then, we can optimize Eq. 2 by alternatively optimize the weights $\lambda$ and model parameters $\theta$. Specifically, we optimize the parameter $\theta_{i}$ of local models on each iteration $t$ by $\theta_{i}^{(t+1)}=\theta_{i}^{(t)}-\eta\lambda_{i}^{t}\nabla_{\theta}L_{local}$, where $\eta$ is the learning rate. Following the mirror gradient ascent of weight proposed in (Zhang et al., 2023), we update $\lambda$ with $\lambda^{t+1}=\frac{\lambda_{i}^{t}e^{\gamma v_{i}^{t}}}{\sum_{i=1}^{|\mathcal{C}|}\lambda_{j}^{t}e^{\gamma v_{j}^{t}}}$. Then we compute $\lambda^{t+1}$ by projecting $\tilde{\lambda}^{t+1}$ into the set $\{\lambda:D_{f}(|\mathcal{C}|)\cdot\|\lambda\|\|(1,1,\dots,1)\|\leq\rho\}$ to fit the constraints of the uncertainty set. In practice, we update $\lambda$ after the local training of each communication 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, $L_{local}$, to encourage the model to learn generalizable information and mitigate the impact of data heterogeneity on maximizing the mutual information $I(z_{x},z_{y})$. The key difference between the two stages is, in the first stage, the optimization target $\theta$ in Eq. 2 corresponds to the parameters of the alignment modules $\Phi$ and $\Psi$, whereas in the second stage, $\theta$ represents the parameters of the feature encoders $\phi,\psi$ and alignment modules $\Phi,\Psi$. The pseudo-code of the whole algorithm can be seen in Algorithm 1.

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-specific encoders, a fundamental design in multi-modal pre-training. We use vision-language pre-training tasks on Chest X-ray datasets and ophthalmology image datasets to evaluate the effectiveness of our FedDRA method.

In this section, we will demonstrate our key empirical findings for federated multi-modal pre-training under heterogeneous client datasets, which actually motivates us to propose our FedDRA. We conduct experiments on the image-text retrieval task, which reflects the ability to maximize the mutual information and learn cross-modal alignment. In the following studies, we 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-preserved manner, while data heterogeneity can affect the effectiveness of the FedAvg. Figure 3(a) presents the retrieval accuracies 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.

Local training can learn harmful client-specific information, degrading the performance of the pre-trained model. After several communication turns, 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 is 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 on learning domain-specific information in the late communication rounds, which would affect the aggregated model’s overall performance.

Decentralized pre-trained deeper layers can hinder the learning of a generalizable feature extractor. We re-trained the first four shallow layers of the decentralized pre-trained model on the combined local datasets. 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). 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 et al., 2024) in supervised federated learning.

Our method learns robust and enriched cross-modal alignment and has better transferability. Table 2 has shown results of downstream tasks, here we utilize the ConVIRT as the backbone pre-training method. In the image-text retrievel task, both average and worst-client accuracies of our method are higher than baseline’s, which means our model can capture more robust cross-client features. In the few-shot classification and segmentation, our method beats other baseline strategies on each task, which demonstrate the higher generalization ability of the representation space learned by our method.

Table 2 has shown the performance of adapted self-supervised federated learning methods which focus on single-modality. 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 a large margin, which indicates that previous single-modality methods cannot be easily adapted directly for multi-modal data. Furthermore, FedAvg is a strong baseline in multi-modal retrieval tasks compared to other adapted methods, as we observed in the experiments. We conjecture that’s because FedAvg only focuses on maximizing the in-domain mutual information, and doesn’t introduce additional loss terms which would hurt the learning of enriched cross-modal alignment. However, this may lead to lower generalization ability on few-shot downstream tasks as discussed before.

The two-stage pre-training strategy and global constraints can enhance the learning of cross-modal alignement. We remove the global constraint loss $L_{C}$ of our method, and compare the pre-trained model’s performances with those of the original version. As shown in Table 3, the downstream performance, particularly the image-text retrieval accuracy, is significantly lower in the modified version. Similarly, we remove the first-stage pre-training on alignment 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-modal 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-weighing 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.

FedDRA dynamically schedules updating stepsize for each client, and therefore optimizes the worst-case performance. We select two clients $C_{1},C_{2}$ 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. 4(d) and Fig. 4(e), 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 FedDRA and the FedAvg method, we plot the cosine-similarities of text and image embedding averaged across clients at each communication turn. As shown in Fig.4(f), FedAvg requires fewer communication rounds to converge, but results in fluctuations after certain communication turns. This aligns with findings in Sec.5.2, where local retraining a 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 a distributionally robust manner, and learns a stronger representation space.

Analysis on global constraint hyper-parameter $\mu$. As shown in Fig. 4(c), a larger $\mu$ encourages federated pre-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 $\mu$ from $1$ to $5$, the downstream performance consistently increases. However, a excessively large $\mu$ can decrease overall performance.

A larger uncertainty radius $\rho$ improves transferability in downstream tasks. Fig. 4(b) shows the downstream performance of models pre-trained with different uncertainty radii in the DRO process. As larger $\rho$ would bring higher performance in few-shot classification and segmentation tasks on out-of-domain datasets. We also observed that a smaller $\rho$ better supports cross-modal alignment learning, achieving better image-text retrieval performance on in-domain datasets, as shown in Table 8 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 on heterogeneity degree of client datasets. We changed the $\alpha$ 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 5, 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 in federated pre-training. As shown in Fig. 4(a), increasing the number of clients introduces greater diversity, which can enhance the downstream performance of the pre-trained model.