Multimodal Federated Learning with Missing Modality
via Prototype Mask and Contrast

Multimodal learning has attracted increasing attention from the research community. The model with dual-encoder architecture (Radford et al., 2021) uses separate encoders for each modality, with shallow modality interaction. On the contrary, the models with interactive-encoder architecture (Wang et al., 2023b; Kim et al., 2021) process input from different modalities and concentrate on modeling modality interactions via fuse tokens without modality-specific encoders.

In order to overcome the missing modality issue in multimodal learning, many methods have been developed. Ma et al. (Ma et al., 2021) propose the SMIL to train a feature reconstruction network using a meta-learning algorithm. Zhao et al. (Zhao et al., 2021) propose the MMIN to learn robust multimodal representations by training cascade residual autoencoders. Ma et al. (Ma et al., 2022) enhance the robustness of Transformer through multi-task learning and optimal fusion strategy search. Wang et al. (Wang et al., 2023a) propose the ShaSpec to generate features of missing modalities using a shared encoder. These methods have introduced additional branches and complex training algorithms to the model, which are unsuitable for FL.

Prototype refers to the centroid of the instances belonging to the identical class (Snell et al., 2017). Due to its scalability, prototypes are widely used to solve various problems in FL (Liao et al., 2023; Lyu et al., 2023; Liu et al., 2023; Li et al., 2023). Tan et al. (Tan et al., 2022) propose FedProto to conduct FL using prototypes rather than aggregation. Huang et al. (Huang et al., 2023) utilize prototypes to solve domain shift in FL. Dai et al. (Dai et al., 2023) utilize prototypes to alleviate the performance decline caused by heterogeneous data in FL. Some other FL frameworks (Xu et al., 2023; Chen et al., 2023a), while not explicitly mentioning prototypes, still leverage the concept of prototypes to address corresponding issues. In this paper, we primarily utilize prototypes to address modality missing issues.

Combining multimodal learning with federated learning is a novel research problem, with the key challenge being how to perform modality interactions when there are missing modalities for each client. Zhao et al. (Zhao et al., 2022) propose FedIoT to train autoencoders for each modality on every client and give more aggregation weight to multimodal clients. Chen et al. (Chen & Zhang, 2022) propose FedMSplit to construct a dynamic graph for adaptively selecting multimodal client models where some modalities may be missing. Yu et al. (Yu et al., 2023) propose CreamFL to conduct cross-modal interaction using inter-modal and intra-modal contrast. However, these studies only considered scenarios involving either unimodal or modality-complete multimodal clients. Feng et al. (Feng et al., 2023) propose FedMultimodal as the first attempt to consider scenarios where modalities of partial data may be missing on each client. They address this by masking the missing modalities using zero tensors. In this paper, we consider scenarios similar to FedMultimodal, with a focus on an efficient and versatile method for addressing modality missing.

We consider a federated learning setting including $N$ multimodal clients. Without loss of generality, we take images and text as instances of multimodal data. In Appendix A, we analyze how to generalize our framework to other multimodal tasks. We decouple the model architecture into a representation encoder $\mathcal{E}$, a deep fusion layer $\mathcal{F}$, and a task head $\mathcal{G}$. In each image-text pair, the image or text will be missing with a probability of $\rho$. In this case, all the images can be denoted as $\mathcal{I}_{n}=\{(i_{nk},y_{nk})\}_{k=1}^{|\mathcal{I}_{n}|}$, where $i_{nk}$ is the $k$-th image on the $n$-th client, $y_{nk}$ is its corresponding label. Similarly, all the text can be denoted as $\mathcal{T}_{n}=\{(t_{nk},y_{nk})\}_{k=1}^{|\mathcal{T}_{n}|}$, and all the image-text pairs can be denoted as $\mathcal{M}_{n}=\{(i_{nk},t_{nk},y_{nk})\}_{k=1}^{|\mathcal{M}_{n}|}$. Furthermore, each client develops a local model $(\mathcal{E}_{n},\mathcal{F}_{n},\mathcal{G}_{n}),n\in[1\dots N]$, with the same architecture as the global model.

As shown in Figure 2, PmcmFL constructs or updates the prototype library on the server (operation ①) at the beginning of each federated communication round. Subsequently, the prototype library and the global model are broadcast to all participants. Then, each participating client performs local training and then updates local prototypes. Using prototypes as masks of missing modalities, a task-calibrated training loss will supervise the training (operation ②). Additionally, to tackle the non-IID issue and the resulting client drift, a prototype-based contrastive loss is proposed (operation ④). Finally, all participating clients transmit their local models and local prototypes back to the server, after which the global model is updated by aggregating those local models. During inference, the prototype library also assists inference with missing modalities by transmitting prototypes to the global model (operation ③). The complete algorithm is in Appendix B.

For local training, We use Multiway Transformer as encoder $\mathcal{E}$. This encoder (Bao et al., 2022; Peng et al., 2022; Wang et al., 2023b) handles different missing patterns by switching among various encoding modes. When only image or text input exists, the multiway encoder functions as a modality-specific encoder, like ViT (Dosovitskiy et al., 2021) or BERT (Devlin et al., 2019). When inputting an image-text pair, the multiway encoder functions as an interactive encoder, enabling the representation of one modality to fuse information from the other modality. More details on the Multiway Transformer are in Appendix C.

Prototypes compact the semantics of data within the same class. Considering communication overhead in FL, low-dimension representations are used to construct the prototype library. As for the compacting strategy, we naively use the class centroids as the prototypes.

Selecting Hidden Representations.   For $(i_{nk},t_{nk})\in\mathcal{M}_{n}$, we denote $\mathbb{H}_{nk}=\mathcal{E}_{n}(i_{nk},t_{nk})\in\mathbb{R}^{m\times d}$ as the output of the multiway transformer encoder, where $m$ is the number of output tokens, and $d$ is the dimension of each token. The image CLS token $h_{nk}^{(i)}$ and the text CLS token $h_{nk}^{(t)}$ are extracted from $\mathbb{H}_{nk}$. We select the two CLS tokens along with the fused representation, i.e., $h_{nk}^{(f)}=\mathcal{F}_{n}(h_{nk}^{(i)},h_{nk}^{(t)})$, for constructing prototypes.

Construction and Update   In the prototype library, we construct and maintain three types of prototypes: image prototypes $\mathcal{P_{I}}$, text prototypes $\mathcal{P_{T}}$, and fusion prototypes $\mathcal{P_{F}}$. All prototypes are initialized with zero tensors or random tensors. Subsequently, after local training in each federated communication round, all participating clients construct local prototypes by computing class centroids. Finally, the global prototypes are aggregated from the local prototypes with numbers of client samples as weights. The construction process is formalized as follows:

where $C_{n}^{j}$ denotes the data of class $j$ on the $n$-th client, $P_{c}$ denotes the set of participating clients, and $\mathcal{D}$ denotes all data of participating clients. For $*\in[\mathcal{I,T,F}]$, $\mathcal{P_{*}}_{n}^{j}$ denotes the local prototypes and ${P_{*}}^{j}$ denotes the global prototypes.

When missing modalities occur, previous work (Feng et al., 2023) masks missing data with zero tensors. In PmcmFL, prototypes are employed as masks for missing modality representations, thereby incorporating global prior knowledge. Based on this, a task-calibrated training loss and a model-agnostic unimodal inference strategy are proposed.

Task-Calibrated Training Loss.   As shown in Figure 2, during local training with complete image-text pairs, the training loss is defined:

where $\mathcal{L}$ denotes the task loss function. For classification tasks, cross-entropy is commonly used. In the case of local training with text missing, a task-calibrated training loss is constructed using the text prototypes as masks:

where $\mathcal{P_{T}}^{j}$ has the same class as $y_{nk}$. Similarly, the task-calibrated training loss is defined when images are missing:

With supervision from the task-calibrated loss, the entire model is trained on the original multimodal task, even when there are missing modalities on clients. On the contrary, using zero masks will result in a task-drifted loss. For example, in the case where text is missing, the training loss with zero masks is denoted as $\mathcal{L}(\mathcal{G}_{n}(\mathcal{F}_{n}(h_{nk}^{(i)},\mathbf{0})),y_{nk})$. Essentially, this represents an unimodal training task that depicts the relationship between the unimodal representation $h_{nk}^{(i)}$ and the ground truth $y_{nk}$, which leads to inconsistent optimization directions for the model training.

Model-Agnostic Unimodal Inference.   During inference with missing modalities, prototypes are employed as masks to assist the global model. Unlike training, where labels are available, we must find corresponding cross-modal prototypes according to embedded representations. To this end, a matching function $m(\cdot)$ is established from hidden representations to prototypes. We use an example to describe the matching process of cross-modal prototypes. Without loss of generality, assume that the model only utilizes images for inference. First, the image $i_{nk}$ is input for computing the image representation $h_{nk}^{(i)}$. Sequentially, the image prototype with the same class $\mathcal{P_{I}}^{j}$ is determined using this matching function $\mathcal{P_{I}}^{j}=m(h_{nk}^{(i)})$. Following this, the corresponding text prototype $\mathcal{P_{T}}^{j}$ is determined according to class $j$. Finally, the image representation $h_{nk}^{(i)}$ and the text prototype $\mathcal{P_{T}}^{j}$ are input into the subsequent modules.

As for matching function $m(\cdot)$, we propose both model-free and model-based matching. Model-free matching involves matching the closest prototype for each representation using a distance metric, such as L1, L2 distance, or cosine similarity. And model-based matching involves training tiny models, treating the search of cross-modal prototypes as either a classification or a retrieval task. These tiny models are trained only in the final round of federated communication, incurring negligible overhead to the FL framework. For the classification task, we train a shallow MLP-Classifier using cross-entropy loss. For the retrieval task, inspired by DALLE$\cdot$2 (Ramesh et al., 2022), we sequentially use CLIP loss (Radford et al., 2021) and MSE loss to train an MLP-Prior to achieve transfer and retrieval. We attempt three aggregation strategies for tiny models from different clients: selecting the model with maximum client samples, model aggregation, and model ensemble.

Inspired by MixUp (Zhang et al., 2018), we propose ProtoMix to enrich the semantic information within prototypes used for inference. ProtoMix linearly combines the top-$k$ related prototypes with weights derived from applying the Softmax function to distance metrics, classification probabilities, or retrieval match scores.

Due to the non-IID issue, the learned representation distributions in latent space significantly differ among clients (i.e., client drift), leading to an uncoordinated model aggregation. To achieve model aggregation without conflicts, all clients must theoretically have similar local representation distributions. In PmcmFL, prototypes are used as learned representation targets to guide clients in learning fused representation distributions clustered by class. For this reason, we introduce a proximal term based on the unidirectional CLIP contrastive loss to regularize client training.

Sampling a batch $\mathcal{B}$ from the local dataset, the proximal term can be denoted as:

where $h_{s}$ is the $s$-th representation in batch, $\mathcal{P}_{s}$ is its corresponding prototype, and $\tau$ is a temperature factor. As the global model gradually converges, the prototypes converge step by step. This proximal term works by pulling representations closer to the prototypes of the same class and pushing representations away from prototypes of different classes so that the learned representations are around the corresponding global prototypes.

Since prototypes are the centroids of class representations, maintaining a certain distance (i.e., fine-grained semantic difference) from each representation, we use CLIP loss to constrain their similarities rather than relying on MSE to constrain their distances. Although the prototype-based proximal term can theoretically be constructed on image representations, text representations, and fused representations separately, we only utilize the last one due to the difficulty in coordinating multiple proximal terms. In this case, the total training loss for each client is the sum of task loss and contrastive loss of fused representations:

where $\gamma$ is the weighting hyper parameter for $L_{contr}^{f}$.

Dataset.   Following CreamFL (Yu et al., 2023), we utilize the challenging VQAv2 dataset (Goyal et al., 2017) to evaluate our PmcmFL. For the efficiency of the experiments, we construct a tiny VQA task, one-tenth the scale of the original VQA task. The tiny VQA is a classification task of 310 classes, with 64,000 training samples and 5,000 testing samples. The training data is distributed to 30 clients, employing the Dirichlet distribution ($\alpha=0.1$) for non-IID data partition (Hsu et al., 2019). Following suggestions from previous work (Feng et al., 2023), we set an equal training missing rate of $\rho_{train}=[0.1,0.2,0.3,0.4,0.5]$ for each modality. More details on the dataset are in Appendix D.

Model.   For the encoder $\mathcal{E}$, we utilize the same architecture as BEIT3-base (Wang et al., 2023b) along with its pre-train parameters. For the deep fusion layer $\mathcal{F}$, the image CLS token and the text CLS token are concatenated and projected through an MLP. For the task head $\mathcal{G}$, the fusion representation is projected to logits using another MLP. More details on the model architecture are in Appendix  E.

Baseline.   We compare our PmcmFL with the most relevant existing FL frameworks. Among them, 1) FedAvg+ignore, where training is conducted only on modality-complete data, ignoring data with modality missing; 2) FedMultimodal (Feng et al., 2023), Essentially an extension of FedAvg, using Zero Mask for missing modalities; 3) Fedmultimodal+random mask, where replaces Zero Mask with Gaussian Random Mask; 4) FedIoT+mask, a framework the same as FedAvg, where Mask is applied to missing modalities, using aggregation strategy is designed for multimodal FL in FedIoT (Zhao et al., 2022). 5) FedPAC+mask, a framework based on FedAvg and Mask for missing modalities, using prototype-based feature alignment (FA) in FedPAC (Xu et al., 2023). 6) FedHKD+mask, a framework based on FedAvg and Mask for missing modalities, using prototype-based hyper knowledge distillation (HKD) in FedHKD (Chen et al., 2023a). To demonstrate the effectiveness of our Prototype Contrast, we ablate it to obtain the variant PmmFL(PmcmFL without Contrast), and introduce the following two baselines: 7) PmmFL+FA and 8) PmmFL+HKD, which replace Prototype Contrast with Feature Alignment (FA) and Hyper Knowledge Distillation (HKD), respectively. We future compare with 9) CreamFL in simple modality missing scenarios. More details on baselines are in Appendix  F.

Implement.   To enable fair comparisons, our experiments remain consistent on method-agnostic hyperparameters, including federated communication rounds, local training epochs, client selection rate, learning rate, optimizer parameters, and batch size. For hyperparameters introduced by PmcmFL, to prevent them from overfitting to specific tasks, we only conducted a simple parameter search. Only the weight $\gamma$ is searched within $[5.0,1.0,0.5,0.1,0.01]$ to find the optimal value. More Implementation details for all experiments can be found in Appendix G.

Evaluation Metric.   We fix the random seed and utilize optimal accuracy in the last ten communication rounds as the final performance. Note that all evaluations are conducted with modality-complete testing samples, except in experiments on inference with missing modalities (Sec. 4.4). Appendix H provides accuracy under varying degrees of modality missing during the inference phase.

Table 1 displays the accuracy of the global model in all baselines and our PmcmFL. Our PmcmFL framework generally achieves noticeable performance improvement over all baselines in all modality missing rates.

Compared to ignoring incomplete data, the Mask method significantly improves the performance of federated training with 50% modality missing (a 2.448% boost), indicating that Mask is effective in handling severe modality missing.

Compared to FedAvg, FedIoT leads to a performance decline (a 0.796% drop on Acc@sum), which indicates that assigning more aggregation weight to multimodal clients no longer works in our task scenarios. Similarly, FedPAC leads to a performance decline (a 6.106% drop on Acc@sum), which is because the feature alignment (FA) employed by FedPAC, a strong regularization that narrows the distance between representations and prototypes through MSE, disrupts the semantic distinctiveness of representations.

Compared to all the baselines, our PmcmFL achieves superior performance (9.376-24.404% boost on Acc@sum), demonstrating the effectiveness of our Prototype Mask and Contrast. Furthermore, as the modality missing rate increases, the performance improvement brought about by our PmcmFL tends to be more significant (a 0.192% boost in 10% missing but a 3.666% boost in 50% missing).

Compared to PmcmFL’s variants, slightly superior performance (3.14-4.93% boost on Acc@sum) indicates that our Prototype Contrast is more generalizable than Feature Alignment and Hyper Knowledge Distillation. We further validate the better performance of PmcmFL compared with CreamFL in scenarios with a 100% missing rate of either modality. The experimental results are in Appendix I.

We study the effect of each PmcmFL component during training. Table 2 displays the results of ablation studies under various training modality missing rates.

The results show that Prototype Contrast and Mask contribute to performance improvement over the baselines, with gains of 7.430% and 11.168% on Acc@sum, respectively. Moreover, combining Prototype Contrast and Prototype Mask assist local training yields further performance boost.

Table 3 shows the testing accuracy of the global model in 100% text missing, demonstrating intuitively the performance of PmcmFL in handling inference with missing modalities. We use the global model trained with 30% modality missing, and its accuracy is 52.256% in modality-complete testing samples. For ProtoMix, top-1, top-10, top-20, and best prototype mix are exhibited.

We can obviously see that the inference accuracy of Zero Mask and Random Mask, serving as baselines, is only 0.996% and 3.662%, respectively, which indicates that the modality missing issues severely impair the model’s inference accuracy. In contrast, Prototype Mask achieves the highest accuracy of 27.502% (a 23.840% boost), demonstrating that our method helps restore modal inference accuracy. Moreover, the experiments also demonstrate the effectiveness of our ProtoMix strategy. More results and discussions can be found in Appendix  J.

We utilize T-SNE (Van der Maaten & Hinton, 2008) for visualization to qualitatively analyze the role of Prototypes Contrast in our FL framework.

Figure 3(a) visualizes fused representation distribution learned by different clients. We can observe that in FedAvg, some dots cluster by clients, indicating that each client forms its individual latent space. When local models are aggregated into a global model, individual latent spaces hinder the formation of a unified latent space. In contrast, in our PmcmFL, the heterogeneity of latent spaces among various clients is reduced due to prototypes acting as representation targets to guide local training.

Figure 3(b) visualizes the fused representation distribution learned by the global model. We randomly selected samples of three classes in the testing set. It can be observed that in FedAvg, dots from different classes are mixed, which is not conducive to subsequent classification tasks. In contrast, in our PmcmFL, dots from different classes cluster around their respective prototypes, which validates our starting point of using prototypes for constructing contrastive loss.

In our PmcmFL, the parameters transmitted during each federated communication round include model weights and the prototype library.

Figure 4 illustrates different FL frameworks’ communication overhead and communication efficiency. As shown, in our tiny VQA task, PmcmFL introduces only negligible additional communication overhead (+0.32%). In the original VQA task, PmcmFL exhibits even lower communication overhead than FedHKD (+3.21% v.s. +5.43%). Furthermore, we can observe that when transmitting the same amount of parameters, PmcmFL achieves optimal performance, demonstrating that our framework has the highest communication efficiency.

Figure 5 illustrates the performance of various FL Frameworks across different communication rounds, client selection rates, local training epochs, and learning rates. We conduct experiments with 50% modality missing while keeping all federated settings constant except the one being investigated. Four representative FL frameworks are compared: 1) FedAvg with ignoring modality-incomplete data; 2) FedAvg with Gaussian Random Mask; 3) FedHKD, the existing one with optimal performance; 4) our PmcmFL. The experimental results demonstrate that, under different FL settings, our PmcmFL consistently achieved optimal performance, showcasing robustness to FL settings.