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

TL;DR

Flex-MoE introduces a novel framework that effectively models arbitrary modality combinations in multimodal learning, especially when some modalities are missing, by leveraging a flexible mixture-of-experts approach with specialized routers.

Contribution

The paper proposes Flex-MoE, a new framework that handles arbitrary modality combinations and missing data using a generalized router and a sparse MoE structure, improving robustness and flexibility.

Findings

Successfully models arbitrary modality combinations in diverse missing data scenarios.

Demonstrates effectiveness on ADNI and MIMIC-IV datasets with multiple modalities.

Achieves robust performance despite missing modalities.

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 https://github.com/UNITES-Lab/flex-moe.