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

Authors: sukwon  yun
inyoung  choi
jie  peng
yangfan  wu
jingxuan  bao
qiyiwen  zhang
jiayi  xin
qi  long
tianlong  chen
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 17:03:05

Abstract: Multimodal learning has gained increasing importance across various fields, offering the ability to integrate data from diverse sources such as images, text, and personalized records, which are frequently observed in medical domains. However, in scenarios where some modalities are missing, many existing frameworks struggle to accommodate arbitrary modality combinations, often relying heavily on a single modality or complete data. This oversight of potential modality combinations limits their applicability in real-world situations. To address this challenge, we propose Flex-MoE (Flexible Mixture-of-Experts), a new framework designed to flexibly incorporate arbitrary modality combinations while maintaining robustness to missing data. The core idea of Flex-MoE is to first address missing modalities using a new missing modality bank that integrates observed modality combinations with the corresponding missing ones. This is followed by a uniquely designed Sparse MoE framework. Specifically, Flex-MoE first trains experts using samples with all modalities to inject generalized knowledge through the generalized router ($\mathcal{G}$-Router). The $\mathcal{S}$-Router then specializes in handling fewer modality combinations by assigning the top-1 gate to the expert corresponding to the observed modality combination. We evaluate Flex-MoE on the ADNI dataset, which encompasses four modalities in the Alzheimer's Disease domain, as well as on the MIMIC-IV dataset. The results demonstrate the effectiveness of Flex-MoE, highlighting its ability to model arbitrary modality combinations in diverse missing modality scenarios. Code is available at: \url{https://github.com/UNITES-Lab/flex-moe}.

Main paper: 
 introduction: 
 In many fields, including healthcare, language, and vision, multimodal learning [6, 75, 37, 44] has emerged as a crucial approach for integrating data from multiple sources such as clinical records, imaging, and genetic data. Multimodal data enables more comprehensive analysis and decisionmaking, offering the potential for improved diagnosis and prediction in various applications [59, 33, 68] . However, a prominent challenge across these domains is the missing modality scenario [76, 60] , where not all modalities are consistently available for every instance due to diverse reasons such as individualized data collection protocols or the variable availability of certain modalities. As an representative example, in Alzheimer's Disease (AD) [4] , one of the most prevalent neurodegenerative disorders, handling this inherently multimodal data is crucial for accurate diagnosis (Figure 1 ). AD datasets often include a combination of clinical symptoms, imaging data [42] , and genetic profiles [48] . However, in real-world clinical settings, not all these modalities are readily available for each patient. Some data, such as clinical and imaging data, may be available from routine visits, whereas other data, such as genetic or biospecimen information, may require additional time to collect. This leads to incomplete datasets, which poses a challenge for existing models that tend to either rely heavily on single modalities or only utilize complete data, thereby missing the opportunity to leverage the full potential of multimodal learning (Figure 2 ).Single Modality and Complete Data Reliance. The reliance on single-modality data or complete data across many frameworks is a significant limitation in real-world scenarios, where missing data is the norm rather than the exception. As seen in Figure 2 , many current models either work with single-modality data or focus on the intersection of modalities, neglecting the potential contribution of partially available modalities. In healthcare, particularly for diseases such as AD, this often leads to missed opportunities in diagnosis and treatment due to the inability to fully exploit multimodal data when some modalities are missing.Oversight of Modality Combinations. Beyond the challenge of missing modalities, there is also the need to model the interactions between available modalities properly. Different combinations of modalities can provide complementary information, and each combination may hold unique significance for downstream tasks. For example, in AD diagnosis, combining biospecimen data and imaging data can reveal key insights: cerebrospinal fluid biomarkers may indicate early signs of AD [21] , while functional MRI can highlight cognitive impairments [42] . Hence, it is essential to develop models that not only handle missing modalities but also effectively utilize the available modality combinations.Given the general challenge of the missing modality scenario in multimodal learning, we propose a novel framework, Flex-MoE (Flexible Mixture-of-Experts), to flexibly incorporate arbitrary modality combinations while maintaining robustness to missing data. Flex-MoE first sort samples based on the available modalities and process them through modality-specific encoders. For missing modalities, we introduce a learnable missing modality bank, which provides learnable embeddings for missing modalities based on the observed ones. This approach ensures that the model can handle incomplete datasets effectively. Our framework also builds upon Sparse Mixture-of-Experts (SMoE) design, allowing us to generalize the expert knowledge from complete data (samples with all modalities) through the G-Router, followed by a specialized S-Router for handling fewer modality combinations. Each expert becomes specialized in handling different modality combinations, ensuring that the model can effectively process any combination of modalities. We demonstrate the effectiveness of Flex-MoE through comprehensive experiments on several real-world datasets, including the Alzheimer's Disease Neuroimaging Initiative (ADNI), which involves four key modalities for AD stage prediction, and the MIMIC-IV dataset. The results confirm the robustness of Flex-MoE in diverse missing modality scenarios.The contributions of this work can be summarized as follows:⋆ We introduce a flexible framework that effectively incorporates arbitrary modality combinations and addresses the missing modality scenario across various domains.⋆ Flex-MoE features a novel approach, including a missing modality bank and generalized and specialized expert training, which ensures robustness to missing modality scenario.⋆ Extensive experiments on real-world datasets, including ADNI and MIMIC-IV, showcase the consistent and robust performance of Flex-MoE in handling diverse modality combinations. 
related works: 
 Single Modality Approach In many fields, deep learning models often rely on single modality data for tasks such as classification [15, 13, 31, 71] , diagnosis [56, 72] , or prediction [12, 61, 34] . While effective in certain cases, these approaches fail to capture the potential synergies between different data sources, especially in contexts where multiple modalities are available. In the Alzheimer's Disease domain, many studies focus on specific modalities. For instance, image-based approaches include a VGG19 model [43] that diagnoses early-stage AD from MRI scans and a modified ResNet18 architecture [45] that predicts AD progression using fMRI data. Other studies focus on genomics, such as DLG [36] for classifying AD patients and SWAT-CNN [26] for discovering AD-associated genetic variants. In the biospecimen modality, a deep learning-assisted spectroscopy platform [29] diagnoses AD by analyzing blood-based amyloid-beta and metabolite biomarkers. Regarding clinical data, a deep learning model [5] outperforms earlier machine learning techniques in classifying AD patients. However, since AD data is inherently multimodal, methods based on a single modality are suboptimal, missing the potential to leverage interactions between different modalities.Multimodal Approach Across multiple fields, multimodal learning has become increasingly valuable for its ability to integrate and capture dynamics within and across different modalities, providing richer and more comprehensive representations of data. Approaches such as the Tensor Fusion Network [74] , Multimodal Transformer [58] , and Multimodal Adaptation Gate [53] highlight the effectiveness of combining multiple data sources. Recently, sparse mixture-of-experts-based methods, such as [44, 8, 19] , have been introduced to enhance modality interactions, though these methods are still relatively unexplored in the AD domain due to the complexity of handling various modality combinations. In AD research, some works have emerged to leverage multimodal data, such as [46] and [59] , which integrated a deep learning framework that combines imaging, genetic, and clinical data, achieving superior AD staging accuracy. Another multimodal model [33] , incorporating longitudinal and cross-sectional data, provided more accurate AD predictions. While multimodal AD studies have shown significant progress, the challenge of missing modalities, especially in the context of how to effectively cope with modality combinations, remains largely underexplored.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We discuss limitations and future directions of the work in the Conclusion section.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
preliminaries and notations: 
 Why Sparse Mixture-of-Experts? Given a multimodal nature, we choose to utilize Sparse Mixtureof-Experts (SMoE) [55] due to its computational efficiency and its ability to handle multimodal data by effectively alleviating the gradient conflict optimization issue between modalities [50] . To briefly introduce SMoE, Traditional Mixture-of-Experts (MoE) models [23, 28, 9, 70] evolved by incorporating sparsity into their structure, optimizing computational efficiency and model performance. SMoE selectively activates only the most relevant experts for a given task, reducing overhead and improving scalability. This innovation is particularly beneficial in handling complex, high-dimensional datasets across diverse applications. It has been widely used in vision [54, 40, 16, 2, 18, 62, 69, 1, 49] and language processing [35, 30, 78, 77, 79, 25] fields, dynamically assigning different parts of the network to specific tasks [41, 3, 20, 11] or data modalities [32, 44] . Research has shown its effectiveness in classification tasks in digital number recognition [20] and medical signal processing [3] . In this work, we further explore the use of SMoE to model arbitrary modality combinations and address the missing modality scenario.Notation. Formally, the SMoE consists of multiple experts, denoted as f 1 , f 2 , . . . , f |E| , where |E| represents the total number of experts, and a router, R, which is responsible for the routing mechanism and sparsely selects the top-k experts. For a given embedding or token x, the router R engages the top-k experts based on the highest scores obtained from softmax function with learnable gating function, g(•) (usually one or two layer MLP), and output R(x) i , where i denotes the expert index. This process can be described as follows:y = |E| i=1 R(x) i • f i (x), R(x) = Top-K(softmax(g(x)), k), TopK(v, k) = v, if v is in the top k, 0, otherwise.(1)In this section, we present our novel algorithm, Flex-MoE, specifically designed to flexibly address the challenge of missing modalities in the multimodal domain. We start by sorting the samples based on their number of observed modalities. Following a modality-specific encoder, we supplement the embeddings for missing parts via missing modality bank completion. This effectively manages missing modalities by learning embedding banks that capture the information specific to observed modality combinations. Next, a Transformer coupled with an SMoE layer is employed. We introduce an expert generalization and specialization step to optimize modality utilization by fully leveraging samples with complete modalities and obtaining modality combination-specific knowledge through samples with fewer modalities. A comprehensive illustration of Flex-MoE is provided in Figure 3 . Throughout the details in the following section, while our work is exemplified through the AD domain for predicting AD stages using four representative modalities-image, clinical, biospecimen, and genetic-it is important to note that Flex-MoE can be generalized to any other multimodal domain. Given samples with diverse modality combinations, we first sort the samples based on their number of available modalities in descending order, and then pass through the modality-specific encoder. (b) Each encoder is only trained with their available samples. For the missing embeddings, we introduce a missing modality bank containing learnable embeddings given the observed modality combination with their corresponding missing modality index. Equipped with this embedding, Flex-MoE passes through the Transformer where the FFN layer is replaced with a Sparse MoE layer. Here, (c) full modality samples take charge of training generalized experts in a balanced manner via G-router, then (d) the remaining few modality combinations further specialize the expert knowledge with S-Router, which fixes the top-1 gate as the corresponding observed modality combination expert. In this figure, top-2 selection of experts is illustrated as an example.Given a set of samples with their own modalities, it is straightforward to pass them through modalityspecific encoders, such as a 3D-CNN for MRI images. However, we are dealing with a missing modality scenario in multimodal data, where specific modalities are often missing based on their observed modality combinations. Thus, it is common to use padded or imputed inputs for the corresponding missing modalities in a multimodal setting. This approach becomes troublesome when considering interactions between modalities. For instance, given a batch with samples, some samples might have image, genetic, and clinical modalities, while others might be missing image and genetic modalities. In such cases, the image encoder and genetic encoder would take zero-padded or hypothesized imputed inputs derived from the missing samples, which are synthetic and of lower quality than the observed ones. This negatively affects the training of modality-specific encoders. Additionally, heavy imputation for each modality in a multimodal setting increases time complexity, which is not desirable in clinical settings.Given this situation, we propose training each encoder solely with observed samples to fully leverage the potential of the encoder by using only observed inputs. Unlike existing approaches, our design principle considers modality combinations to ensure robust and flexible training and effective handling of missing modalities. As illustrated in Figure 1 , each patient has diverse symptoms and personalized treatments, leading to variations in available modalities. For example, patients might lack image and genomic modalities (i.e., possessing only biospecimen and clinical modalities) due to various reasons such as patient conditions, resource limitations, or specific clinical settings [22, 51] . Imputing these missing modalities must be handled within this context, rather than applying a global learnable representative embedding for each modality without considering the observed environment.Therefore, we propose a learnable missing modality bank. Given the number of modality combinations without fully observed scenarios, the total cases would be, given a modality set M = {I, C, B, G},|M|-1 m=1 |M| m = 2 |M| -1.The resulting missing modality bank can be defined as B ∈ R 2 |M| -1×|M| . By using this bank, the concatenated embedding of all modalities for patient i,h i = [e I i , e C i , e B i , e G i ]would be represented as follows:EQUATIONwhere e m i ∈ R d denotes the embedding from the corresponding modality m for patient i, and d is the hidden dimension. Here, 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., B M\m,m = B {I,G,B},C ). By doing so, the encoder for each modality can be trained without encountering non-observed, incomplete input features. Then, we move on to the Transformer layer, where the FFN layer is replaced by an SMoE layer, a common approach in the SMoE domain [24, 38, 10] .While adopting the SMoE backbone, it is important to note that our environment differs from concurrent SMoE studies, especially in terms of multimodal learning with missing modalities. In this context, choosing the most relevant tokens is challenging for experts, since the significance, i.e., the quality of input information, varies with the missing modalities. This significantly motivates us to take a distinct approach from concurrent SMoE studies, where the input token is derived from fully observed scenarios. To address the unique challenges of the missing modality scenario, we propose a modality combination-specific MoE design. Specifically, we assign expert indices based on all possible modality combinations. For example, 'IGCB' is assigned as 0, 'IGC' as 1, ..., up to 'B' as 14. The remaining experts act as buffers, allowing the Router to select the most relevant top-k experts and activate them automatically. This approach leaves room for flexibility and maintains the initial intuition of the MoE design.Generalization It now becomes clear why the samples used for training Flex-MoE are sorted in descending order. Inspired by curriculum learning [7, 63] , where easy samples are presented first and more challenging samples appear later, we regard the level of difficulty as the number of missing modalities. We first train our SMoE layer with easy samples, where all modalities are fully observed. Using this intersection as a gold standard, we initially train all the experts in the MoE model. The procedure essentially follows the vanilla SMoE design as described in Equation 1, but with one key difference: the input tokens consist only of inputs where all modalities are fully observed. Hence, we refer to this router as the Generalized Router, G-Router. This approach leverages the completeness of information in these samples, which should be fully utilized before specializing the experts in their respective areas. To ensure balanced activation of the experts initially, which will later specialize, we incorporate the load and importance balancing loss [55] , which will later be exemplified in Equation 4.Specialization Once the experts are initially trained using fully observed samples, we aim to specialize each expert, which is the key advantage of the MoE design. We leverage the remaining samples, which encompass diverse modality combination configurations. Each modality combination requires its own specialized expertise. For instance, samples with Image, Biospecimen, and Genetic data will have a corresponding expert index activated through the top-1 gating mechanism to fully utilize the specialized knowledge of that expert (i.e., expert 'IBG' in Figure 3 ). To effectively specialize the modality combination-specific experts, we propose a Specialized Router design, S-Router, which encompasses following technical novelties. First, to facilitate targeted expert selection when an input token is provided, we avoid manually replacing the selected routing policy with our preferred choice in a post-hoc manner, which would stop the continuous gradient flow. Instead, we innovatively introduce a cross-entropy loss between the top-1 expert selection and the targeted expert indices for each token by the S-Router. Formally, this can be described as follows:EQUATIONwhere MC(x j ) denotes the modality combination index of a given token x j in a total of n tokens. max(S-Router(x j )) denotes the maximum prediction probability of the corresponding activated expert index, which corresponds to the probability of the top-1 expert index. By comparing these two, the router is trained to activate the corresponding expert index for a given input token with a certain modality combination. Thus, the specialized knowledge inherent in specific modality combinations is naturally contained within the target expert.Moreover, when calculating load and importance balancing loss [55] , we specifically compute the loss for the remaining top-k-1 experts, as the top-1 selection is manually handled and thus considered biased rather than balanced. We aim for the selection of the remaining k-1 experts to occur in a balanced manner, allowing interaction with other related modality combinations. Formarlly, it can be expressed as follows:EQUATIONwhereEQUATION, σ(x) is the standard deviation of x, µ(x) is the mean of x, g ie is the gate value for sample i with expert index e as discussed in Equation 1, and δ(• > 0) is an indicator function that is 1 when the inner value is greater than 0. E \ e top-1 denotes the set of expert indices excluding the top-1 selected expert index. This ensures that the resulting MoE model not only retains global knowledge but also incorporates specialized expert knowledge tailored to specific modality combinations. During the inference phase, the specified expert index for a particular modality combination can be activated alongside other experts, enabling predictions to consider both the specific modality combination and intersections with other modalities.Finally, the output embeddings of the Sparse MoE layer pass through a 1-layer MLP prediction head to predict one of the three stages of AD, i.e., Dementia, CN, or MCI. To further facilitate a curriculum-learning approach, we first use warm-up epochs with sorted samples, followed by shuffled samples for the remaining epochs. This strategy enhances generalizability across diverse input samples, enabling better handling of variability during the inference phase. [66] ). Before integrating all modalities, to address the initial missing data within each modality, we applied simple mean imputation [39] for each column. For more detailed data table with preprocessing steps for each modality, please refer to Appendix A.1.We use the Medical Information Mart for Intensive Care IV (MIMIC-IV) database [27] , which contains de-identified health data for patients who were admitted to either the emergency department or stayed in critical care units of the Beth Israel Deaconess Medical Center in Boston, Massachusetts24. MIMIC-IV excludes patients under 18 years of age. We take a subset of the MIMIC-IV data, where each patient has at least more than 1 visit in the dataset as this subset corresponds to patients who likely have more serious health conditions. For each datapoint, we extract ICD-9 codes, clinical text, and labs and vital values. Using this data, we perform binary classification on one-year mortality, which foresees whether or not this patient will pass away in a year. We drop visits that occur at the same time as the patient's death. In order to align the experimental setup with the ADNI data, which does not contain temporal data, we take the last visit for each patient.Baselines We compare the performance of Flex-MoE with various state-of-the-art baselines across modality-specific, e.g., image or genetic, and multimodal approaches. For the image-only modality, we first experimented with 3D MRI scans by utilizing a 3D CNN [17] and an architecture that combines 3D CNN and 3D CLSTM [68] . To decrease computational complexity, we also extracted 2D slices from the 3D volumes. For 2D MRI scans, we implemented the VGG architecture with pre-trained weights and applied layer-wise transfer learning [43] , as well as a modified ResNet-18 network [45] . For the genetic-only approach, we employed a ResNet-34 based architecture to handle the high-dimensional genetic data [36] . In ADNI dataset, we further implemented domain speicfic baselines, such as auto-encoder and 3D CNN-based architecture that incorporates imaging, genetic, and clinical data [59] , and a GRU-based architecture that considers imaging, genetic, clinical, and biospecimen data [33] . Moreover, we include ShaeSpec [60] , which utilizes a spectral attention mechanism to emphasize important features across modalities, and mmFormer [76] , which is based on transformer-based multimodal fusion with an attention mechanism. For multimodal approaches in both ADNI and MIMIC-IV, we incorporate the recent FuseMOE [19] model, which directly integrates multimodal data through a mixture of experts strategy, as the most straightforward baseline. Additionally, we compare the following methods: MulT [57] , which captures cross-modal interactions through cross-attention mechanisms; MAG [52] , which fuses multimodal features by mapping them to an adaptation vector; TF [73] , which combines multimodal embedding sub-networks and a tensor fusion layer; and LIMoE [44] , which addresses training stability in multimodal learning using entropy regularization based on contrastive learning.Experimental Settings. To ensure a fair comparison with other baselines, we used the best hyperparameter settings provided in the original papers. If not available, we tuned the learning rate in 1e-3, 1e-4, 1e-5, the hidden dimension in 64, 128, 256, and the batch size in 8, 16. For our proposed method, we searched the number of experts in 16, 32, and Top-k in 2, 3, 4. We set the coefficient of the sum of additional losses (importance and load balancing) combined with our cross-entropy loss to 0.01, scaling it within the task classification loss. For the dataset split, we chose 70% for training, with the remaining 30% split evenly between validation and test sets (15% each). It is important to note that, to share the same inference space, where single and multimodal baselines should both be able to predict, we opted to choose the intersection as the test and validation sets. This means that during the training phase, the dataset can be incomplete. For the multi-modal baselines, if they had the ability to impute or interact with other modalities, we leveraged their methods. Otherwise, we used zero-padding to facilitate batch-wise training. For single-modal and multi-modal baselines that solely work on the intersection region, we filtered that data and used it during training. All experiments were conducted using NVIDIA A100 GPUs. Each experiment was run three times with different seeds to ensure reproducibility, and the results were averaged. The optimal hyperparameter settings for Flex-MoE can be found in Appendix A. Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
primary results: 
 In Table 1 and Table 2 , we provide a comprehensive comparison of Flex-MoE with various multimodal baselines. We have the following observations: (1) Overall, Flex-MoE performs effectively in diverse multimodal settings, fully harnessing its potential as more modalities become available. This is supported by the large margin of improvement (7.6% and 11.07% over the best performing baselines, MAG and the most recent work FuseMoE, respectively, in full modality settings in Table 1 ).(2) Although the recently proposed FuseMoE [19] suggested its ability to handle missing scenarios, the lack of effective modality combination creates a bottleneck in such AD domain, even performing worse when a smaller number of modalities is used (FuseMoE performs better with three modalities than with full modalities), which is not optimal given the diverse missing modality scenarios.(3) Despite its specific characteristics in the AD domain [33, 59] , the reliance on intersection data and the lack of consideration for how missing modalities relate to observed modality combinations have been overlooked. (4) Overall, the performance gain derived from Flex-MoE can be attributed to its unique ability to cope with diverse modality combinations through a missing modality bank, and its capability to fully harness the knowledge of samples via a generalization followed by a specialization step for experts. For additional results on different metrics, such as Macro-F1 and AUC, please refer to Appendix A.3.To validate the effectiveness of the two essential modules of Flex-MoE-the missing modality bank and the unique SMoE design-under a missing modality scenario, we evaluate them followingly.First, to evaluate the effectiveness of the missing modality bank introduced in Figure 3 (b), we assess whether it captures relevant embedding information given an observed modality combination. Specifically, we validate this by examining the inter-relationship between modalities, focusing on the consistency of how the missing modality bank handles missing information. In Figure 4 , we measure the cosine similarity between observed modality combinations. The key observation is that (1) with more overlapping modality combination, it tends to share more similar embedding information. This is evident in the left side of Figure 4 , where the full modality scenario (ICBG) shows higher similarity with ICB and CBG (0.56 and 0.58, respectively) compared to IC and CB (0.46 and 0.45). The clinical modality (C) is most similar across combinations, which aligns with the dataset characteristic that clinical data is present in all input combinations, as shown in Figure 2 . On the right side of Figure 4 , the similarity between missing modalities is shown. When (2) modality G is missing, it is more similar to the cases where C and B are missing compared to I, suggesting that certain missing modalities share more commonality in how they are handled by the model. This insight underscores the importance of careful consideration when modeling missing modalities and demonstrates how the missing modality bank effectively captures necessary embedding information in terms of modality combination. Furthermore, in Figure 5 , we show the activation ratio of input modality combinations across each expert index, i.e., possible modality combinations. We observe two key findings: (1) Thanks to expert generalization using full modality samples (BCGI), generalized knowledge is distributed across all experts. This allows each expert to leverage commonly shared knowledge while activating the most relevant inputs for the downstream task.(2) Expert specialization enables each expert to acquire specialized knowledge. For instance, in the case of the BCG expert, the two most activated input tokens were BCGI and its corresponding token, BCG. Similarly, for the BCI and CI experts, they not only possess general knowledge from BCGI but also hold their own specialized knowledge from BCI and CI, respectively, to effectively handle these inputs. This demonstrates that when a certain modality combination is provided, the top-1 expert is successfully selected, allowing it to supplement its specialized and necessary information. Overall, these routing experiments demonstrate that Flex-MoE contains both globally generalized and locally expertspecific knowledge, achieved by leveraging samples with both full and fewer modalities.Ablation Study. In this section, we investigate the crucial components that contribute most positively to the performance gain of Flex-MoE. From Table 3 , we observe that (1) when both expert specialization and generalization are absent, the performance drop is most severe. Additionally, (2) the performance decline in the embedding bank negatively affects overall performance, indicating that the missing modality bank combined with expert generalization and specialization is crucial for handling missing modality scenarios. Furthermore, (3) the sorting based on descending order appear most effective as the expert generalization occurs first withi the full modality samples.Sensitivity Study. In Figure 6 , we also varied the hyperparameters used in this study. We examined the number of experts, number of SMoE layers and top-k selection. We found that (1) employing many experts does not always guraantee a higher performance compared to its increase in complexity, showing using 16 experts apeear to be a suitable choice to equip fine-grained specialized knowledge.(2) Using sing a single layer of the SMoE was most effective, as stacking more layers or adding a Transformer block caused an overload to parameter learning. Additionally, (3) compared to the commonly used top-2 gating network in concurrent SMoE studies, we found that top-4 selection was the most effective. This is because manually assigning the top-1 expert index to the target modality combination leaves more room for better harmonization with the SMoE design.Complexity Study. In Table 4 , we further verify the benefits of utilizing the SMoE design in terms of mean time per iteration, GFLOPs, and the number of parameters compared to the baselines. We observed the following: (1) Compared to recent baseline FuseMoE, Flex-MoE achieves notable efficiency gains (e.g., 22.74%, 1.15%, and 89.17% gain in mean time, GFLOPs, and # of parameters, respectively.) while delivering higher performance.(2) Although TF appears to be a lightweight design in the I, G and I, G, C settings, it trades off computational efficiency with significantly lower performance compared to Flex-MoE.(3) Notably, as the number of modalities increases, existing models tend to become more complex in terms of GFLOPs and the number of parameters to manage the additional complexity. However, Flex-MoE remains robust and efficient, maintaining higher performance due to its effective use of sparsely activated experts, brought by the SMoE framework.  
conclusion: 
 While multimodal learning brings new opportunities and challenges across various domains, including medical fields, existing approaches struggle to handle arbitrary modality combinations, especially in missing modality scenarios, often relying on single modalities or complete datasets. In this work, we propose a flexible multimodal learning framework, Flex-MoE, capable of managing arbitrary subsets of available modalities. By carefully considering modality combination, it leverages a learnable embedding bank to capture missing modality information and utilizes a unique SMoE design to enhance expert generalization and specialization. Extensive experiments on the representative ADNI and MIMIC-IV datasets validate its effectiveness in handling diverse modality combinations. Future work includes extending the framework to explore the scaling laws of available modalities, which in turn presents numerous modality combinations, offering significant room for further improvement. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our work does not use any data or models that have a high risk of misuse such as pretrained language models, image generators, or scraped datasets.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We use the ADNI dataset and provide the proper citation. All materials related to the paper will adhere to the copyright policies of NeurIPS.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained? Answer: [NA] Justification: IRB approval was not necessary for this project. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: The proposed algorithm has the potential to significantly improve early diagnosis and treatment outcomes for patients, reducing the burden on healthcare systems. However, its effectiveness can be limited by the availability of comprehensive and high-quality patient data, and there may be challenges in integrating this tool into existing clinical workflows. Justification: We clearly state the challenges of the problem at hand and the contributions that we make in the abstract and introduction.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: This paper addresses multimodal learning with missing modalities with Flex-MoE, which is a sparse mixture of experts method.  The approach is tested for Alzheimer’s disease classification using the ADNI dataset, and results show SOTA performance.  The reviewers identified several strengths, such as: 1) idea is clear, 2) paper is well written, 3) research topic is relevant, 4) motivation is well explained, and 5) the ablation study is comprehensive. The reviews also identified a few issues, which include: 1) missing more comprehensive experiments in other tasks, 2) missing discussion on complexity, 3) equation (2) needs clarification, and 4) missing comparison with other methods (e.g., mmformer, shaspec). The rebuttal addressed all points above and the final scores of the paper are 5,7,7 (average score of 6.33, with an average confidence of 3.67). I recommend this paper to be accepted.

Reviews: review: paper_summary: paper_summary: This paper introduces Flex-MoE, a novel multimodal learning framework for Alzheimer's Disease that handles missing modalities using a Sparse Mixture-of-Experts design and demonstrates its efficacy on the ADNI dataset.
strength_weakness: strength_weakness: 1. The paper emphasizes that the Sparse Mixture-of-Experts (SMoE) selectively activates only the most relevant experts to improve scalability, but does not provide experimental results demonstrating computational efficiency or scalability.

2. The interpretation of \( B_{M \setminus m, m} \) in Equation (2) is unclear; it should be better explained in details. Whether it involves a learnable module that outputs missing modality features based on available modality features, or via other method?

3. The approach to handling missing modalities differs from the method in the paper "Multi-modal Learning with Missing Modality via Shared-Specific Feature Modelling," which uses averaging for missing modalities. A comparison with this approach would be beneficial. The paper provides code https://github.com/billhhh/ShaSpec.

4. The paper overlooks some existing missing modality baselines, such as mmFormer and ShaSpec.

Some minor weaknesses:
1. The abstract states that "few recent studies attempt to integrate multiple modalities," but there are actually many methods for handling missing modalities (as noted in point 4 above), which are not compared experimentally in this paper.

2. The font size in Table 1 is too small and needs to be redrawn for better readability.

3. In line 180, the bold formatting for "It" is incorrect.

4. In Table 2(b), the value 61.08 in the first row is not the best result, which should be corrected.

5. Table 2 should ideally use a visualization approach similar to that in the ShaSpec paper to represent different modality combinations more clearly.
questions: questions: The most biggest question for me is pointed in weaknesses Q2: The interpretation of \( B_{M \setminus m, m} \) in Equation (2) is unclear; it should be better explained in details. Whether it involves a learnable module that outputs missing modality features based on available modality features, or via other method?
limitations: limitations: Please refer to the weaknesses section
paper_summary: paper_summary: The paper presents a multimodal learning framework, Flex-MoE (Flexible Mixture-of-Experts), designed to integrate diverse modalities in Alzheimer's Disease (AD) research using a Sparse Mixture-of-Experts design. Flex-MoE sorts samples based on the number of available modalities and processes them through modality-specific encoders. The framework trains experts with full modality samples and uses an S-Router to adapt the knowledge for fewer modality combinations.
strength_weakness: strength_weakness: 1. The model's specificity to Alzheimer's disease is unclear, as it appears to be a general approach for handling missing modalities.
2. The experimental setup lacks clarity. Details on how labels (e.g., dementia, CN, MCI) were decided, the prediction timeline, and the handling of patients at different stages are not provided. Additionally, the use of clinical data and whether it is time series data is not explained.
3. The evaluation is unconvincing. The imbalance rate of the datasets is not discussed. Only accuracy and F1 score are used for evaluation, but additional metrics like AUC could provide a more comprehensive understanding of the model’s performance. The process for selecting the threshold for classification to obtain the F1 score is also unclear.
4. The computational resources required for training and inference using Flex-MoE are not specified, especially given the need to train the model for various modality combinations.
5. No external dataset is included to validate the model's performance.
questions: questions: 1. Why is the model specifically designed for Alzheimer's disease, and how does it differ from a general model for missing modalities?
2. How were the labels (dementia, CN, MCI) decided? What is the prediction timeline, and were patients at different stages treated as separate samples? What time period of data was used, and was the clinical data time series data?
3. What is the imbalance rate of the datasets? Including metrics like AUC can provide a more complete understanding of the model’s performance. How was the threshold chosen for determining the classification boundary to obtain the F1 score?
4. What specific computational resources are required for training and inference using Flex-MoE?
5. How would the model perform on an additional external dataset?
limitations: limitations: 1. The evaluation and practical application of the model are challenging to assess. The necessity of training different models for each modality combination is not well justified.
2. The lack of external dataset validation limits the generalizability of the model's performance.
paper_summary: paper_summary: The paper aims to deal with the task of AD patient classification with missing modalities, and proposes a framework, Flex-MoE.
Flex-MoE also utilizes a learnable missing modality combination bank to complete the missing embeddings of missing modalities.
The training strategy and the gating mechanism of Flex-MoE are also well-designed to handle this task.
The experiments verify that Flex-MoE is able to achieve SoTA performance in this task on the shown modality combinations.
strength_weakness: strength_weakness: - The expression from Line 159 to Line 167 is a little bit confusing. More explanation (or even illustration) is needed.
- Is the usage of the missing modality bank just a simple look-up operation? I am concerned that it seems to be naive and there may be a more smart approach to achieve this.
- There is not $R_{(x_j)}$ of Line 199 in your equations. You need to correct this.
- Why the S-router can be trained to activate the corresponding expert index by Eq. (3)?
- It would be better if you could use some equations to assist to express the content of Line  204-215.
- This paper also leaks sufficient discussion of some related works, such as other tasks of biomedical images [1,2], in Sec.2 or appendix.
- The modality combinations of the experiments in Table 2 is few. It would be better to provide the experiments based on more combinations to showcase the performance of Flex-MoE.

[1] Zhou, D., Gu, C., Xu, J., Liu, F., Wang, Q., Chen, G., & Heng, P. A. (2023). RepMode: learning to re-parameterize diverse experts for subcellular structure prediction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 3312-3322).

[2] Jiang, Y., & Shen, Y. (2024). M$^4$oE: A Foundation Model for Medical Multimodal Image Segmentation with Mixture of Experts. arXiv preprint arXiv:2405.09446.
questions: questions: Please see the section of weaknesses. 
I may raise or lower my score after the rebuttal.
limitations: limitations: The authors have claimed the limitations and broader impacts of the proposed method.

Comments: title: Adjusted the score to 7
Thank you and we will ensure inclusion of discussions
Appreciation and further clarification on cognitive assessment in ADNI dataset
Response from Reviewer 6yys
Appreciation and kind request for your consideration of score adjustment
Responses after rebuttal
Thank you for your consideration
The authors have addressed my concerns well
Please check the authors' rebuttal
Please check the authors' rebuttal
Please check the authors' rebuttal
Eager for Your Feedback on Our Rebuttal
Eager for Your Feedback on Our Rebuttal
Eager for Your Feedback on Our Rebuttal
Please read the rebuttal to check if the authors addressed your concerns
Please read the rebuttal to check if the authors addressed your concerns
Please read the rebuttal to check if the authors addressed your concerns
comment: I'd like to adjust my score to 7 according to the efforts of rebuttal by authors
Dear Reviewer **6yys**,  

Thank you for your response.  
As noted by the reviewer, we will ensure that the following points are included in our final version:

- 1. **More detailed explanations** regarding lines 159-167 and the equations in lines 204-215.
- 2. **Additional modality combination** results, expanding beyond those currently presented ($\mathcal{I}$, $\mathcal{I}+\mathcal{G}$, $\mathcal{I}+\mathcal{G}+\mathcal{C}$, $\mathcal{I}+\mathcal{G}+\mathcal{C}+\mathcal{B}$).
- 3. Inclusion of **more baseline** discussions and results, specifically for ShaSpec and mmFormer.
- 4. Incorporation of **additional datasets**, such as the MIMIC dataset.
- 5. Analysis of **computational efficiency**, including training time, GFLOPs, and parameters.

Once again, thank you for all your insightful comments to improve our paper. 


Best,  
Authors
Dear Reviewer **qdnn**,

Thank you for your response. We appreciate that our rebuttal was well received, and we thank you for your willingness to increase the score.

In our above rebuttal response [link](https://openreview.net/forum?id=ihEHCbqZEx&noteId=nYk2KzEUPG) under **[W 2, 3 & Q 2: Clarity in Experimental Setup]**, we provided the specific file referenced in this study, where the labels (e.g., Dementia, CN, MCI) for each patient were annotated according to the ADNI dataset. We also presented the statistical distribution and imbalance ratio.

However, we realize that there may still be some concerns regarding the details of the **cognitive assessments** used in our model. These assessments are crucial in determining the diagnostic categories such as Cognitively Normal (CN), Mild Cognitive Impairment (MCI), and Dementia. Below, based on ADNI documentation [1], we provide a more detailed explanation of each label including the specific cognitive tests involved and how they contribute to the classification process.

1. **Cognitively Normal (CN) (Interchangeable with Cognitively Unimpaired (CU))**:
    - **Meaning**: Participants in this group show no signs of significant cognitive decline at the time of their evaluation.
   - **Key Cognitive Assessments**: To be classified as CN/CU, individuals must score 0 on the Clinical Dementia Rating (CDR) and show no memory impairment (memory box score of 0). Their performance on the Mini-Mental State Exam (MMSE) and the Wechsler Logical Memory II sub-scale must be within the normal range, adjusted for education level.
   - **Significance**: These tests ensure that participants classified as CN/CU exhibit cognitive functioning consistent with normal aging, without evidence of dementia or mild cognitive impairment.

2. **Mild Cognitive Impairment (MCI)**:
    - **Meaning**: MCI is an intermediate stage between normal aging and dementia, where participants show noticeable cognitive impairments that do not yet meet the criteria for dementia.
   - **Key Cognitive Assessments**: Individuals are labeled as MCI if they have a CDR global score of 0.5 and a memory box score of at least 0.5. The MMSE and Wechsler Logical Memory II sub-scale are also used, with specific cutoffs to diagnose MCI based on educational levels. These cognitive tests reveal impairments in memory and other cognitive domains that are more pronounced than in normal aging but not severe enough to warrant a dementia diagnosis.
   - **Significance**: These assessments are crucial for identifying individuals at risk for developing dementia, making them a key focus for early interventions.

3. **Dementia (Dementia/AD)**:
   - **Meaning**: Participants in this group meet the clinical criteria for dementia, typically due to Alzheimer's Disease, characterized by significant cognitive and functional impairment.
   - **Key Cognitive Assessments**: A CDR global score of 0.5 or 1, combined with impaired scores on the MMSE and Wechsler Logical Memory II sub-scale, typically categorizes a participant as having dementia. Additionally, the classification process involves screening to exclude those whose symptoms suggest non-Alzheimer's forms of dementia, such as Frontotemporal Dementia.
    - **Significance**: These cognitive assessments are vital for studying the progression of Alzheimer's Disease and for developing treatments targeted at this advanced stage.

For further details regarding the protocols and ADNI study information (we used ADNI 3 in this study), you may kindly refer to the official ADNI documentation [2].

In summary, to provide a clearer understanding for readers, we will incorporate this discussion into our final version.

Once again, thank you for all your comments to improve our paper.

Best,  
Authors

[1] https://adni.loni.usc.edu/data-samples/adni-data/study-cohort-information/  
[2] https://adni.loni.usc.edu/help-faqs/adni-documentation/
Dear authors,

Please also make sure that you will include all explanations and discussions for these concerns in your final paper or appendix, which can decently improve the quality and clearness of this work.

Best, 

Reviewer 6yys
Thank you for the authors' response. I still have some concerns about the specific cognitive assessments included in the model. I’m willing to adjust my scores to 5 and my confidence to 3.
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
Dear Reviewer **6yyS**,

We are pleased to hear that our rebuttal has addressed your concerns. We will ensure that those additional experiments are included in the final version.

Best wishes,  
Authors
Dear authors and AC,

Sorry for my late reply.
The authors have addressed my concerns, with additional experiments, explanations, and discussion for some unclear expression in the paper.
I would like to raise my score to 7, if the authors can add these contents to the main paper or the appendix.
Thanks for your time and effort.

Best,
Reviewer 6yys
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.
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.
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 **6yyS**,

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](https://openreview.net/attachment?id=KkGeR9PVFe&name=pdf)

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

Best,  
Authors
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](https://openreview.net/attachment?id=KkGeR9PVFe&name=pdf)

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

Best,  
Authors
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](https://openreview.net/attachment?id=KkGeR9PVFe&name=pdf)

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

Best,  
Authors
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 qdnn, 

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

Best regards
Your AC.
Dear Reviewer BwQM, 

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

Best regards
Your AC.