Flex-MoE: Modeling Arbitrary Modality Combination via the Flexible Mixture-of-Experts

We thank reviewer 6yyS for mentioning our work as practical, addressing missing modalities is worth attention to the AD community and rationale and decent techniques. For the concerns, we provide responses below:

[W 1, 3: Expression from Line 159 - 167, 199]
The main idea of the missing modality bank completion is to supplement missing modalities from a predefined bank, ensuring robust data integration with observed ones. For example, if a patient lacks clinical data but has imaging, biospecimen, and genetic data, the observed modalities pass through their specific encoders. The missing clinical embedding is supplemented from the missing modality bank, indexed by the observed modalities (e.g., {Imaging, Biospecimen, Genetic}, Clinical). This approach prevents reliance on incomplete or naively imputed data, as encoders only process observed modalities. Once embeddings are standardized in dimension, they proceed to the Sparse MoE layer within the Transformer architecture. For line 199, it should be $\text{max}(\mathcal{S}\text{-Router}(x_j))$. To improve readability, we will include such details in the final version.

[W 2: Usage of Missing Modality Bank]
As the reviewer mentioned, this can be seen as a look-up operation, but the way we construct and leverage this look-up table is what makes our work novel and distinctive. The rationale behind designing the missing modality bank is to differentiate the context of patient groups based on their unique observed modality combinations and corresponding missing modalities. In the AD domain, the missing modality problem provides a unique context for each patient’s diagnosis. For example, patients with biospecimen and image data but missing clinical and genetic data exhibit unique characteristics [1]. Motivated by this, we designed the missing modality bank with learnable embeddings indexed by observed modality combinations and their corresponding missing modality. This allows existing encoders to process only the observed features without being distracted by missing samples, equipping the model to handle various observed modality contexts flexibly. Empirically, as shown in Figure 4 of the manuscript, the clinical modality (C) shows more significant similarities with the full modality combination when predicting AD diagnosis compared to other modalities (I, G, B). Additionally, patients missing clinical (C) and genetic (G) modalities exhibit more similarities compared to other missing modality combinations, reinforcing our claim that missing modality combinations provide unique diagnostic contexts in AD.

[1] https://pubmed.ncbi.nlm.nih.gov/24360540/

[W 4: Equation to assist line 204 - 215]
Here, we aim to utilize load balancing loss [1] especially targeted for the remaining experts (i.e., experts except top-1 selected expert) to ensure experts for each modality combination to be balanced activated. Formally it can be expressed as,

1. Load Balancing Loss
   $\mathcal{L}_{\text{balance}} = \text{CV}^2(\sum{_j^N} \text{importance}{_j}) + \text{CV}^2(\sum{_j^N} \text{load}{_j})$

• where $N$ is the number of samples

2. Coefficient of Variation Squared (CV²)
   $\text{CV}^2(x) = \left( \frac{\sigma(x)}{\mu(x)} \right)^2$

• where $\sigma(x)$ is the standard deviation of $x$ and $\mu(x)$ is the mean of $x$.

3. Importance and Load
   $\text{importance}{_e} = \sum{_j^N} g _ {je}, \forall e \in E \setminus \{e _ {\text{top-1}}\}$

• where $e$ is expert index and $g{_ie}$ is the gate value for sample $j$ with expert $e$
  $\text{load}{_e} = \sum{_j^N} \delta(g _ {je} > 0), \forall e \in E \setminus \{e _ {\text{top-1}}\}$
• where $\delta(g _ {je} > 0)$ is an indicator function that is 1 if the gate value $g _ {ie}$ is greater than 0, indicating that the expert $e$ is selected for sample $j$.

For the prediction head, we can simply denote it as: $\mathbf{Y} = \mathbf{Z} \cdot \mathbf{W}$ where $\mathbf{Y} \in \mathbb{R}^{N \times |\mathcal{C}|}$ denote the predicted probability for each sample on each class, $\mathbf{Z} \in \mathbb{R}^{N \times (D*|\mathcal{M}|)}$ denote the concatenated embedding of each modality, and $\mathbf{W} \in\mathbb{R}^{(D*|\mathcal{M}|) \times |\mathcal{C}|}$ denote the weight matrix that transforms the concatenated embedding into final prediction space.

For the final version, we will incorporate above equations to provide better understanding.

[1] https://arxiv.org/abs/1701.06538

[W 5: Discussion of related works]
To briefly summarize, RepMode predicts 3D fluorescent images of subcellular structures from 3D transmitted-light images, addressing challenges of partial labeling and multi-scale variations. M4oE uses modality-specific experts and a gating mechanism to enhance medical image segmentation across various modalities. However, RepMode focuses on subcellular structure prediction, and M4oE focuses on medical image segmentation, differing from our work, which aims to accurately classify the stage of AD disease, particularly in the multimodal AD domain. While these works leverage image modalities, they lack the ability to cope with various other modalities, such as genetic and clinical. Moreover, although M4oE uses modality-specific experts, it does not consider the characteristic of modality combination information in practical missing modality scenarios, which is our main focus.

[W 6: More modality combination experiments]
To achieve greater generalizability, besides more modality combinations, we added the MIMIC dataset [1], the AUC score as an additional metric, and two more baselines (ShaSpec and mmFormer) during the rebuttal period. The results can be found here: PDF.

We still observe Flex-MoE outperforms existing baselines across various modality combinations.

[1] https://physionet.org/content/mimiciv/3.0/

Dear Reviewer 6yyS,

Can you have a look at the rebuttal and see if your concerns have been addressed?

Best regards Your AC.

Dear Reviewer 6yyS,

Please don't forget to read the authors' rebuttal to reach a final decision about this paper in the next day or so. If you have any further questions to the authors, that would be a great chance to reach out to them.

Best regards Your AC.

We thank reviewer qdnn for the comprehensive review of our paper. Specifically, we appreciate the reviewer's recognition of our key contribution in addressing missing modalities in AD as crucial and highly relevant. For the concerns raised, we provide the following responses:

[W 1 & Q 1: model’s specificity to AD]
As the reviewer mentioned, our approach can be generalized, similar to common phenomena in various biomedical papers (e.g., Transformer in 3D MRI [1], Graph Neural Network in the single-cell domain [2,3]). However, we emphasize the importance of a careful application that considers domain-specific characteristics, especially in response to recent improvements in the ML field. In this paper, we focused on the AD domain for two main reasons:

• Underexplored Characteristic: Despite the comprehensive multi-modal nature of AD (e.g., clinical, biospecimen, image, genetic) [3], it remains underexplored compared to domains with fewer modalities (e.g., image, text) [4, 5]. Empirically, as shown in Table 2 of the manuscript, the direct application of multimodal ML, e.g., LIMoE [5], underperforms compared to Flex-MoE in handling 3 or more modalities and in overall performance. This supports our motivation to address the diverse multimodal nature prevalent in the AD domain and the specific missing modality combinations.

• Unique Characteristics of Missing Modalities: In the AD domain, the missing modality problem provides a unique context for each patient’s diagnosis. For instance, patients with biospecimen and image modalities but missing clinical and genetic data present unique characteristics [6]. Empirically, as shown in Figure 4 of the manuscript, clinical modality (C) shares more significant similarities with the full modality combination when predicting AD diagnosis compared to other modalities (I, G, B). Additionally, patients missing clinical (C) and genetic (G) modalities show more similarities compared to other missing modality combinations, reinforcing our claim that missing modality combinations provide unique diagnostic contexts in AD.

In summary, our study aims to advance clinical diagnosis in the AD domain by effectively integrating multimodal data, especially in the growing era of AI4Science (Biology, Medicine, Health).

[1] https://www.nature.com/articles/s41598-024-59578-3
[2] https://www.nature.com/articles/s41467-021-22197-x
[3] https://www.nature.com/articles/s41467-023-36559-0
[4] https://arxiv.org/abs/2309.15857
[5] https://arxiv.org/abs/2206.02770
[6] https://pubmed.ncbi.nlm.nih.gov/24360540/

[W 2, 3 & Q 2: Clarity in Experimental Setup]
To validate the effectiveness of Flex-MoE, we focused on the gold-standard ADNI dataset. Following existing studies [1,2], we performed an AD stage prediction task to predict specific labels for Dementia, CN, and MCI. Specifically, we used the diagnostic summary file 'DXSUM_PDXCONV_22Apr2024.csv' and the 'DIAGNOSIS' column, which underwent clinical and neuropsychological evaluation as reported by ADNI [3]. The label statistics in our study were: CN (1030 patients, 43.3%), MCI (860 patients, 36.1%), Dementia (490 patients, 20.6%). Despite concerns about imbalance, we observed a relatively balanced ratio of 490/1030 = 0.47. The F1-score was calculated using the scikit-learn package [4], which automatically computes the score without manual selection of threshold.

Regarding timeline information, timestamp data is not readily available across genetic, biospecimen, and clinical modalities. However, for the image modality, we selected the most recent image to ensure relevance and mapped the patient ID with other modalities. Thus, our method remains static, but dynamic analysis with time information is a promising direction for future work.

We will add this information to the final version for clarity. Detailed statistics and preprocessing steps for each modality are provided in Section 4.1 and Appendix A of the manuscript, addressing the reviewer's concerns about the clinical modality.

[1] https://www.nature.com/articles/s41598-020-74399-w
[2] https://www.nature.com/articles/s41598-018-37769-z
[3] https://adni.loni.usc.edu/wp-content/uploads/2008/07/inst_about_data.pdf
[4] https://scikit-learn.org/stable/modules/generated/sklearn.metrics.f1_score.html

[W 4 & Q 4: Computational Resources]
We appreciate the reviewer's concerns about computational resources. We used NVIDIA A100 GPUs for training and inference but highlight the efficiency of adopting the Sparse MoE design, where only the selected top-k experts are utilized. We performed a computational comparison experiment in terms of Mean Time, FLOPs, and the number of activated parameters, which can be found here: PDF.
The result demonstrate the efficiency of the SMoE design, showcasing the potential for various real-world datasets.

[W 3, 5 & Q 4, 5 & L 1, 2: More Comprehensive Experiment]
To achieve greater generalizability, we added the MIMIC dataset [1], the AUC as an additional metric, and two more baselines (ShaSpec and mmFormer) during the rebuttal period. We tested all possible modality combinations, and the results can be found here: PDF.
Our experiments show that Flex-MoE outperforms existing baselines across various metrics and modality combinations, thanks to its careful consideration of modality combinations. Additionally, to justify the necessity of training different experts, Figure 5 shows that each expert index contains both global context knowledge and specialized modality combination knowledge, providing flexibility in responding to various modality combinations. Figure 4 further illustrates the interpretability of how different observed and missing modality combinations share similarities in predicting AD status.

[1] https://physionet.org/content/mimiciv/3.0/

Dear Reviewer qdnn,

Can you have a look at the rebuttal and see if your concerns have been addressed?

Best regards Your AC.

Dear Reviewer qdnn,

We sincerely thank you for dedicating your time to review our work and for your constructive feedback. As the deadline for the discussion period approaches, we are eager to engage further and understand if our responses address your concerns satisfactorily.

For quick access to additional experiments, including more baseline results, datasets, and computational efficiency, you can find them here: PDF

We would greatly appreciate it if you could kindly review our response. Thank you for your consideration.

Best,
Authors

Dear Reviewer qdnn,

Please don't forget to read the authors' rebuttal to reach a final decision about this paper in the next day or so. If you have any further questions to the authors, that would be a great chance to reach out to them.

Best regards Your AC.

We thank reviewer BwQM for their careful review and for mentioning that Flex-MoE is straightforward, effectively addresses missing modalities in AD, and easily understandable. For the remaining concerns, we provide details below:

[W 1: Computational efficiency and scalability]
Throughout the paper, we explained the rationale behind the SMoE design principle for handling missing modality scenarios by assigning modality combinations to each expert index. As the reviewer mentioned, utilizing the SMoE design indeed benefits computational efficiency and scalability. To further verify its effectiveness, we compared the training time, GFLOPs, and number of parameters used in training at this link: PDF.

We observe that Flex-MoE not only performs best in the I, G, C, B modality configuration, as shown in Table 2 of the manuscript, but also achieves efficiency in terms of training time and GFLOPs. Additionally, it demonstrates scalability with fewer parameters compared to existing baselines, including the state-of-the-art model, FuseMoE [1].

[1] https://arxiv.org/abs/2402.03226

[W 2 & Q 1: Interpretation of $\mathbf{B} _ {\mathcal{M} \setminus m}$ in Equation (2)]
The main idea of the missing modality bank completion is to supplement missing modalities from a learnable embedding bank. If a sample contains missing modality data, we retrieve the missing information (i.e., embedding) from this bank, which includes all possible modality combinations and their respective missing modalities. Otherwise, for the observed input, the data directly passes through the modality-specific encoder.

For example, if a patient lacks clinical data (i.e., $m = \mathcal{C}$ where $m$ denotes the missing modality, thus $\mathcal{M} \setminus m = \{ \mathcal{I}, \mathcal{B}, \mathcal{G} \}$) but has imaging, biospecimen, and genetic data, the observed modalities pass through their respective encoders. The missing clinical embedding is supplemented from the missing modality bank, indexed by the observed modalities (i.e., $\mathbf{B} _ {\mathcal{M} \setminus m} = \mathbf{B} _ {\{\mathcal{I}, \mathcal{B}, \mathcal{G}\}}$). This approach ensures that encoders only process observed features, preventing them from being influenced by incomplete or naively imputed data. As this missing modality bank is learnable throughout the iterations, it eventually possesses the knowledge relevant to the downstream task, such as AD stage prediction.

Empirically, as shown in Figure 4 of the manuscript, the clinical modality (C) shows more significant similarities with the full modality combination when predicting AD diagnosis compared to other modalities (I, G, B). Additionally, patients missing clinical (C) and genetic (G) modalities exhibit more similarities compared to other missing modality combinations, reinforcing our claim that missing modality combinations provide unique diagnostic contexts in AD.

[W 3 & W 4 & M-W 5: Comparison with other baselines in ShaSpec style]
We appreciate highlighting these two relevant baselines in our study. To briefly summarize each paper: ShaSpec employs a strategy using auxiliary tasks based on distribution alignment and domain classification, along with a residual feature fusion procedure, to learn shared and specific features for handling the missing modality issue. During the rebuttal period, we employed ShaSpec, and the performance comparison on the ADNI dataset with additional metric, AUC can be found here: PDF.
Interestingly, while ShaSpec has a specific module to handle missing modalities, it falls short of Flex-MoE. We believe this is due to forcing distribution alignment without carefully considering the context of each modality combination. For instance, when the clinical modality is missing, ShaSpec leverages other modalities to generate the clinical feature, but this may overlook the unique contextual meaning of the missing clinical data. In contrast, Flex-MoE retrieves a learnable embedding from the missing modality bank to supplement the clinical embedding without forcing observed modalities to contribute, allowing the clinical embedding more flexibility to learn downstream task-specific information.

Moreover, mmFormer is an end-to-end framework initially designed for different MRI modalities, consisting of hybrid modality-specific encoders, a modality-correlated encoder, and a convolutional decoder. This approach differs from ours, which considers images as one modality and extends the perspective to genetic, biospecimen, and clinical data. When implementing mmFormer for our task, we followed its proposed architecture for processing image data and replaced their encoder with our modality-specific encoders used in Flex-MoE. The results can be found here: PDF.

In summary, mmFormer lacks the ability to handle diverse modalities beyond images, which have different contexts, making it significantly less effective than Flex-MoE. Additionally, by supplementing missing modality embeddings as global modality-invariant features, mmFormer fails to handle the specific context of modality combinations and corresponding missing modalities. This justifies the necessity of incorporating modality combination information while handling missing modalities, as proposed in Flex-MoE.

[M-W 1,2,3: Improvements]
Thank you for the constructive feedback. We will incorporate the discussed baselines, change the font size in Table 1 for better readability, and carefully revise the bold font.

[M-W 4: Best result in Table 2 (b)]
We have checked the bold formatting in Table 2 (b) and found no errors. Flex-MoE with 61.08 performs best in that row. Could we kindly ask reviewer BwQM to re-clarify the point? We are ready to provide more details and refine accordingly.

Dear Reviewer BwQM,

Can you have a look at the rebuttal and see if your concerns have been addressed?

Best regards Your AC.

Dear Reviewer BwQM,

We sincerely thank you for dedicating your time to review our work and for your constructive feedback. As the deadline for the discussion period approaches, we are eager to engage further and understand if our responses address your concerns satisfactorily.

For quick access to additional experiments, including more baseline results, datasets, and computational efficiency, you can find them here: PDF

We would greatly appreciate it if you could kindly review our response. Thank you for your consideration.

Best,
Authors

Dear Reviewer BwQM,

Please don't forget to read the authors' rebuttal to reach a final decision about this paper in the next day or so. If you have any further questions to the authors, that would be a great chance to reach out to them.

Best regards Your AC.

Dear Reviewer BwQM,

Thank you once again for your thoughtful feedback and for confirming that all of your concerns have been addressed. We greatly appreciate your time and effort in reviewing our work.

As you noted, we will include all modifications and results to our final version. We would be deeply grateful if you could kindly consider adjusting the score to reflect an improved positive evaluation of our submission. Your support at this stage would be of great help to us.

Thank you once again for all your support.

Best,
Authors

The rebuttal of authors addresses all of my questions. They provide more baseline results, datasets, and computational efficiency. So I have no more questions related to the paper. W.r.t. the Best result in Table 2 (b), thank you for the clarification and I've double checked it again. It should be correct. Kindly add the new results in the newer version of the paper. Thanks for the effort!

Best