Multi-modal Learning with Missing Modality in Predicting Axillary Lymph Node Metastasis
^††thanks: This study was partially supported by the National Natural Science Foundation of China (Grant no. 92270108), Zhejiang Province Natural Science Foundation of China (Grant no. XHD23F0201).

Missing modality of clinical data in the test time is considered in the paper. The dataset is divided into a training set and test set: $\mathcal{D}=\{\mathcal{D}^{d},\mathcal{D}^{v}\}$. We consider the training set $\mathcal{D}^{d}=\{(x^{w}_{i},x^{c}_{i},y_{i})\}_{i=0}^{n-1}$ as a modality-complete dataset, where $x^{w}_{i}$ and $x^{c}_{i}$ represent two different modalities (whole slide images (WSI) and clinical information) of the $i$-th sample, $y_{i}$ is the corresponding label and $n$ is the total number of the samples in the training set. The test set $\mathcal{D}^{v}=\{(x^{w}_{0},x^{c}_{0},y_{0}),(x^{w}_{1},y_{1}),...\}$ is a modality-incomplete dataset. There exist samples in the data set $\mathcal{D}^{v}$ that do not contain the clinical information. In this paper, we aim to make full use of the multi-modal information in training set to improve the model performance and flexibly deal with the modality-missing problem in the test set.

Specifically, a WSI $x^{w}_{i}$ is divided into $t$ small patches which feed into an encoder. In the multi-modal branch, deep features $\{f_{i,j}^{wm}\}_{j=0}^{t-1}$ from the encoder are aggregated to the fused feature $f_{i}^{wm}$ by simple attention [22].

where $\mathcal{H}(\cdot)$ is a non-linear projection function whose parameters are learnable. The output of $\mathcal{H}(\cdot)$ is a 1D vector with length $t$. The $j$th element corresponds to the patch feature $f_{i,j}^{wm}$ during summation. We combine $f_{i}^{wm}$ and the mapped feature $f_{i}^{c}$ of clinical data to calculate the final classification loss $\mathcal{L}_{mul}^{c}$.

where $\mathcal{G}_{mul}(\cdot)$ is the final classifier in the multi-modal branch and $[\cdot,\cdot]$ is a feature splicing operation.

In order to avoid the feature learning of WSI being affected by clinical information, we transfer the knowledge from the single-modal branch to the multi-modal branch. We define the intermediate output $\{f_{k}^{mul}\}_{k=0}^{l-1}$ of the deep layers in the multi-modal branch and $\{f_{k}^{sgl}\}_{k=0}^{l-1}$ in the single-modal branch. $l$ is the number of deep feature layers in the network. In this paper, we choose the WSI features from the final layer $f_{l-1}^{mul}$ and $f_{l-1}^{sgl}$ (the output of the attention module $f_{i}^{wm}$ and $f_{i}^{ws}$ in Fig. 2), which exhibit the most robust semantics of WSIs. The knowledge distillation for the WSI features can be represented as follows:

where $\mathcal{D}$ is a distance function that measures the gap of features between the single-modal and multi-modal branches. We choose the mean square error as the distance measure function. $\mathcal{M}_{mul}$ and $\mathcal{M}_{sgl}$ are the projection modules that can transfer the intermediate output feature to the target representation. The total loss function for the learning of the multi-modal branch is

where $\lambda_{m}$ is a hyper-parameter to weigh different items. We utilize the classification loss $\mathcal{L}_{mul}^{c}$ and the distillation loss $\mathcal{L}_{mul}^{f}$ to update the multi-modal branch simultaneously.

Following patch fusion steps described in Multi-modal Branch Learning, we convert deep features $\{f_{i,j}^{ws}\}_{j=0}^{t-1}$ to $f_{i}^{ws}$. We employ a learnable prompt [23] $x_{i}^{p}$ to signal the single-model branch when missing modality and memorize the missing information of clinical data $x_{i}^{c}$. We map $x_{i}^{p}$ to a feature $f_{i}^{p}$ by a non-linear function. The dimension of $f_{i}^{p}$ is as same as that of the feature $f_{i}^{c}$ in the multi-modal branch. We then combine the WSI feature $f_{i}^{ws}$ and the prompt feature $f_{i}^{p}$. Afterward, the knowledge of clinical data is transferred from the multi-modal branch to the single-modal branch based on the distillation loss:

We also apply mean square error for $\mathcal{D}$. $KL(\cdot,\cdot)$ [24] is the KL divergence function for the predicted confidence. $\mathcal{G}_{sgl}$ and $\mathcal{G}_{mul}$ are two final classifiers for single-modal branch and multi-modal branch, respectively. The loss function $\mathcal{L}_{sgl}^{f}$ is only used for the learning of the prompt as shown in Fig.2.

Similarly, there is also a classification loss $\mathcal{L}_{sgl}^{c}$ for the learning of WSIs in the single-modal branch. Consequently, the total loss function for the training of a single-modal branch is presented as follows:

The loss function is used for the update of the single-modal branch while the multi-modal branch is frozen. During testing, the BD framework can tackle modality complete or incomplete inputs in a unified manner by turning off or on the single branch.

The experimental dataset is from a grand challenge named Early Breast Cancer Core-Needle Biopsy WSI (BCNB) [10]. Paired multi-modal data containing WSIs and clinical information is provided by the dataset. All WSIs are hematoxylin and eosin stained and the clinical data consists the information of age, tumor size, ER, PR and HER2. We use the information to predict the metastatic status ($N_{0}$ and $N_{+}$) of axillary lymph nodes. Since it is a binary classification task, we use the metrics Area Under Curve (AUC) and F1-scores (F1) to validate the proposed method. F1 represents the averaged results in the prediction of metastatic status.

We randomly split the dataset into a training and test set with 80% and 20%. A subset with 20% is separated from the training set for validation. We assume that the training set is complete with paired modalities. While the clinical data in the test set can be missed at a random rate. In the training process of our method, stochastic gradient descent with a momentum of 0.3, a weight decay rate of $1\times 10^{-3}$ serves as the optimizer. The learning rate is initialized at $1\times 10^{-4}$. The hyper-parameters $\tau$, $\lambda_{s}$ and $\lambda_{m}$ are set to 1.2, 0.5 and 0.6, respectively. We initialize the learnable prompt with a length of 50. Early stopping is used to avoid overfitting by monitoring the F1 scores in the training set. The code is implemented based on python3 and pytorch-1.9 and all experiments are conducted using NVIDIA A100 GPUs.

All non-linear and linear projection modules are composed of fully connected layers and the ReLU non-linear activation function. The function $\mathcal{H}(\cdot)$ within the attention module consists of two hidden layers with corresponding activation functions. We employ two layers with hidden sizes of 100 and 50 to map the learnable prompt $x_{i}^{p}$ to $f_{i}^{p}$. Both $\mathcal{M}_{mul}$ and $\mathcal{M}_{sgl}$ are fully connected layers used to map features to a dimension of 64.

We compared our proposed approach with representative methods (AE [25], Ensemble [26], Filling, SMIL [16]) in dealing with the missing modality problem for multi-modal learning, The mechanism of first three intuitive methods are as shown in Fig. 3.

• •

  Filling is the method that aims to fill the missing clinical data with zero vectors. The model structure is based on the model LNMP [10]. It is the same as LNMP when the modalities are complete during test.

• •

  AE is designed to generate the missed deep features of clinical data automatically. This model is trained with two stages. First, we train an LNMP model with the modality-complete training set. Then, an auto-encoder is trained to generate missed features, the input and output of which are features of the WSIs and clinical data, respectively.

• •

  Ensemble is the model that has two individual networks. One is the WSI recognition network, whose output is the predicted probability. The other one is the classification network for clinical data. We get the final prediction result by fusing the probabilities from the two networks. We only use the first network if there is no input of clinical data.

The results of comparisons are shown in Table I. Our method achieves the best performance (bold) regarding F1 scores and AUC. Compared to others, AE yields relatively worse performance. This method often requires a large amount of paired training data, therefore, is difficult to effectively generate the accurate features of clinical data for the prediction. The direct filling method meets the requirement of test flexibility but does not provide valuable information about the missing clinical data. Thus, performance decreases greatly with the increase of the missing ratio and even becomes worse than that of the method with only images. The integration of two separate networks in the Ensemble method has better performance both on the complete modality and incomplete modality among the three intuitive approaches. However, the two networks are independent, and the complete modality in the training set is underutilized. For instance, a test sample only with the modality of WSI is not helped by the clinical data in the training set. Our method is also inspired by this finding and further improves the performance based on Ensemble. SMIL can also be regarded as a generative model. Differently, it is trained end-to-end, but the shared encoder may be perturbed by clinical data.

Ablation study on distillation directions. We split the BD framework into two parts: the single-branch learning from the multi-branch ($\mathcal{M}\rightarrow\mathcal{S}$) and the multi-branch learning from the single-branch ($\mathcal{S}\rightarrow\mathcal{M}$). Then, we design the ablation study to verify the effectiveness of each part. We test in two situations with the missing ratio of 0% (complete modality) and 80%. F1-scores of the model with or without each part are presented in Table II. We regard the independent two branches without distillation as the baseline in the ablation study. Under the test case of complete modality, the performance remains the same after adding the part $\mathcal{M}\rightarrow\mathcal{S}$ due to turning off the single branch when testing. But there is a substantial improvement after appending the part $\mathcal{S}\rightarrow\mathcal{M}$. In the case of missing modality, it is exactly the other way around. The part $\mathcal{M}\rightarrow\mathcal{S}$ is crucial to the performance of the model, while $\mathcal{S}\rightarrow\mathcal{M}$ has little effect on it. Thus, $\mathcal{M}\rightarrow\mathcal{S}$ and $\mathcal{S}\rightarrow\mathcal{M}$ are necessary for the incomplete and complete modality respectively.

Ablation study on the initial length of learnable prompt. We opt for the scenario where 100% of clinical data is missing for comparison. As shown in the subfigure (c) of Fig.4, the model performs best when the initialization length is 50 (the green line). Too long initialization of the prompt may result in memory redundant information of missing modality. We believe that shorter initializations might convey less information, yet the prompt can still serve as a reminder to the model regarding the absence of the modality. The performance will not drop significantly.

Ablation study on $\lambda_{m}$ and $\lambda_{s}$. Missing 100% clinical data is considered in this experiment. We first fix $\lambda_{s}$ and vary $\lambda_{m}$ to record model performance. Then we choose the best $\lambda_{m}$ and change $\lambda_{s}$ to various values. As shown in TABLE III, larger values of $\lambda_{m}$ lead to better performance, illustrating the necessity of distilling WSI features from the single-modal branch. It may overwhelm the classification loss $\mathcal{L}_{sgl}^{c}$ as $\lambda_{s}$ increases, resulting in performance degradation. From the values of the saved epoch, the model converges faster when $\lambda_{s}$ is larger.

For further study, we analyze the deep features of WSI before and after the addition of clinical information. We train a single-modal model with only WSIs and a multi-modal model with completely paired data, including WSIs and clinical data. Then, the deep features of WSI from the two models are collected. We perform feature dimensionality reduction based on t-SNE[27] and visualize these features on the two-dimensional plane as shown in Fig. 5. (a): The WSI features are extracted by the model trained with only WSIs. (b): The WSI features are from the intermediate output of a trained multi-modal model, in which the deep features of clinical data are expanded by 20 times compared to the original dimension. (c): The WSI features are also collected from the multi-modal model, where the deep features of clinical data are expanded by 40 times.

We find that the WSI features from the single-modal model are more aggregated and divisible. And there is a sign that we may get worse WSI features as the feature dimension of clinical data increases. Thus, we conclude that the addition of clinical information might affect the representation learning of WSIs. Inspired by this finding, we keep the part that the multi-modal branch learns from the single-modal branch ($\mathcal{S}\rightarrow\mathcal{M}$).

To validate the efficacy of our model, we consider the scenario of WSI absence. We employ the learnable prompt to alert the model of the presence of missing modality and memorize the information of WSIs from the multi-modal branch. The image encoder is removed from the single-modal branch, and the prompt is directly mapped to the deep feature. The subfigure (a) and (b) of Fig.4 illustrate that our model far outperforms the base model ($Filling$), and it consistently outperforms the single-modal model at various missing rates. This demonstrates the effectiveness of our model no matter which modality is missing.