Completed Feature Disentanglement Learning for Multimodal MRIs Analysis

Feature representation learning is a crucial aspect in the field of MML, which contains three types of methods. Some approaches [4, 3] have focused on the first type of method, which extracted specific features from each modality and subsequently fused them with obtained embeddings. Braman et al. [4] designed the Multimodal Orthogonalization (MMO) loss function to obtain the maximum specific representation for each of radiology, pathology, genomic and clinical data. Several methods [5, 6] have concentrated on the second type of method, which captured modality-shared features from multiple modalities. Ning et al. [6] built a bi-directional mapping between original space and shared space of multimodal to effectively obtained multi-modal shared representation. However, the fist two type of methods have primarily emphasized either modality-specific or modality-shared features, thus failing to learn a comprehensive representation of multimodal data. The third type of method, FD, has proven to be effective in separating multimodal information into meaningful components and has been successfully applied in various applications [7, 9, 8, 10]. Hu et al. [7] proposed a disentangled-multimodal adversarial autoencoder (DMM-AAE) model that employed a VAE to disentangle multimodal MRIs information into modal-common features and modal-specific features. However, this method only addressed the two-modal fusion scenario. Cheng et al. [9] extended this approach to multimodal fusion scenario. It is worth noting that both of these methods cannot be trained end-to-end due to the reliance on hand-crafted features as model inputs. Hazarika et al. [8] decoupled multimodal information into modality-invariant and modality-specific features using Central Moment Discrepancy metric, orthogonality constraints and reconstruction loss. Li et al. [10] proposed decoupled multimodal distillation (DMD), which first separated the representation of each modality into modality-irrelevant space and modality-exclusive space. Then, a graph distillation unit was employed to each space for dynamic enhancing the features of each modality.

The aforementioned FD methods have a common drawback that can result in incomplete feature representation learning in the case of three or more modalities, as depicted in Fig. 1. In contrast, the proposed CFD strategy addresses this limitation by decoupling multimodal information into modality-shared features, modality-specific features, and modality-partial-shared features, thereby enabling comprehensive feature representation learning.

The fusion strategy is another crucial aspect of MML. Several approaches have involved concatenating features extracted from different modalities [20, 3] or representation spaces [7, 9, 10], such as modality-shared and modality-specific features. However, concatenation fusion has not effectively utilized the correlations between multiple modalities.

In recent years, there has been an increasing focus on exploring the correlations among multiple modalities to obtain effective features. Some methods achieve multimodal fusion by assigning weights or probabilities to each modality. [15, 16, 17] explored the uncertainties of different modalities to obtain reliable multimodal fusion information. Choi et al. [21] proposed EmbraceNet which performed multimodal representation fusion based on a probabilistic approach. Zhou et al. [22] introduced a canonical correlation analysis (CCA)-based method named ADCCA to exploit the correlation between multiple modalities and integrate the complementary information from these modalities. Zhuang et al. [23] proposed a global-guided fusion method which consider both global and local correlations of multiple modalities.

With the proven ability of attention mechanisms to enhance feature representation and explore complex correlation between multiple modalities, many attention-based multimodal fusion methods have emerged [24, 18, 25, 19, 26]. Zhang et al. [18] proposed a modality-aware mutual learning (MAML) framework that weighted the multimodal features using an attention-based modality-aware (MA) module. Zhu et al. [25] captured complementary information from multimodal data using self-attention and cross-modal attention, and further designed a triple network to obtain more discriminative information. Xing et al. [19] developed the NestedFormer framework, which included the Nested Modality-aware Feature Aggregation (NMaFA) module to explore long-range correlations within and between modalities for effective and comprehensive information learning. However, these attention-based methods cannot explicitly reveal the contribution of each decoupled features during fusion process.

Mixture-of-Experts (MoE) [27] employs multiple experts to extract distinct representation spaces from the input and generated corresponding weights using a gating network. MoE have the ability to dynamically capture the mixture information from multiple experts. Several studies [28, 29] have extended MoE to handle multi-input scenarios, where each expert processes a specific input. These approaches leverage the dynamic nature of MoE. However, these methods concatenate all inputs to generate weights in the gating network without thoroughly considering the relationships between different inputs, which can limit the effectiveness of the fusion process. In contrast, we introduce a gating network to capture the local-global relationships between the decoupled features.

The proposed CFDL framework, illustrated in Fig. 3, comprises three parts: a) feature extraction from multimodal MRIs, b) completed feature disentanglement for feature decoupling, and c) dynamic MoE fusion for dynamically integrating the decoupled features. The framework employs the same type of backbone to extract latent features $x_{j}$ from each modality. To capture a comprehensive representation of the multimodal data, the latent features are decoupled into modality-shared features, modality-specific features, and modality-partial-shared features using the proposed CFD strategy. The decoupled features are then integrated via the DMF module. Within the DMF module, each decoupled feature is paired with a specific expert, and a gating network named LinG_GN generates weights for multiple experts. The fused features are obtained by aggregating the weighted features from the experts. In the following, we take three-modal condition as examples to illustrate the CFD strategy and DMF module.

Inspired by previous FD methods [7, 8, 9], we first decouple the extracted latent features into modality-shared features and modality-specific features. We employ a shared-encoder $\bm{E^{c}}$ to decouple modality-shared features and three private-encoders $\bm{E^{p}}_{j}$ to decouple modality-specific feature for each modality. Three modality-shared features can be formulated as follows:

and three modality-specific features can be obtained with:

The final modality-shared feature $F$ is the mean of all modality-shared features, given by

We further consider modality-partial-shared features between pair-wise modalities. As a result, three groups of modality-partial-shared features are decoupled, with each group consisting of two features. The two features in the same group ($G_{jk}^{j}$, $G_{jk}^{k}$) are decoupled with the same partial-shared-encoder named $\bm{E^{s}}_{jk}$, i.e.,

Specifically, $G_{jk}^{j}$ represents the modality-partial-shared feature between the $j$-th modality and the $k$-th modality, which is decoupled from the $j$-th modality. The final modality-partial-shared feature $G_{jk}$ can be calculated by averaging the modality-partial-shared features in each group,

In the end, we obtain three modality-shared features ($F_{1}$, $F_{2}$, $F_{3}$), three modality-specific features ($P_{1}$, $P_{2}$, $P_{3}$), and three groups of modality-partial-shared features including $\{G_{12}^{1},G_{12}^{2}\}$, $\{G_{13}^{1},G_{13}^{3}\}$, $\{G_{23}^{2},G_{23}^{3}\}$. Furthermore, we get 7 final decoupled features, which is denoted as a set $\mathcal{S}$, where $\mathcal{S}=\{F,P_{1},P_{2},P_{3},G_{12},G_{13},G_{23}\}$.

To enhance the completeness of the decoupled representation, we use the following three constraints:

• 1)

  Modality-shared features should exhibit high similarity to one another.

• 2)

  Modality-partial-shared features within each group should have maximum similarity.

• 3)

  The final decoupled features should exhibit maximum dissimilarity from one another.

To ensure the effective decoupling of modality-shared features and modality-partial-shared features, we employ the mean squared error (MSE) loss as a constraint. The MSE loss measures the discrepancy between two features, and we aim to increase similarity between two features by minimizing this loss. The losses for modality-shared features $\mathcal{L}_{sim}^{c}$ and modality-partial-shared features $\mathcal{L}_{sim}^{s}$ are expressed as:

We denote $\mathcal{L}_{sim}$ as the sum of $\mathcal{L}_{sim}^{c}$ and $\mathcal{L}_{sim}^{s}$.

To enhance the decoupling of modality-specific features and minimize redundancy among all final decoupled features, we incorporate cosine similarity as a constraint for better optimization. Our objective is to increase the dissimilarity between the final decoupled features by reducing the cosine similarity between each pair of these features. The loss for all final decoupled features $L_{diff}$ is calculated by:

where $CS(\cdot,\cdot)$ represents the cosine similarity function.

To ensure dynamic fusion of the final decoupled features, we introduce the DMF module based on the MoE architecture, which is shown in Fig. 3 (c). In DMF module, each final decoupled feature $\mathcal{S}_{j}\in\mathbb{R}^{dim}$ is associated with a specific expert $\bm{E^{x}}_{j}$, which is implemented as a fully-connected layer. We introduce LinG_GN to dynamically generate weights for these experts, taking into account the relationships of these final decoupled features. The LinG_GN operates with two inputs to capture a comprehensive understanding of these features. Firstly, we concatenate all the final decoupled features as one input for the LinG_GN. To integrate this concatenated feature, we map the concatenated feature into a unified representation utilizing a fully-connected layer $FC$. This unified representation captures the collective information from all the final decoupled features and is treated as the global feature $g\in\mathbb{R}^{dim}$,

where $Cat(\cdot,...,\cdot)$ means column concatenation operation and $ReLU$ is the activation function. Secondly, we stack all the final decoupled features as the second input for the LinG_GN. This stacked features retain the individual information from each final decoupled feature and are considered as the local features, denoted as $\textbf{O}=[{\mathcal{S}_{1}}^{T},...,{\mathcal{S}_{7}}^{T}]$, where $\textbf{O}\in\mathbb{R}^{7\times dim}$. By obtaining both the global feature and the local features, we can explore the importance of each local feature within the context of the global feature. This exploration allows us to determine the weight $\omega$ of local features in contributing to the fused representation. The weight $\omega$ is calculated by:

where $\times$ represents matrix multiplication operation, $\omega\in\mathbb{R}^{7}$. Each element in $\omega$ represents the weight for the corresponding expert network $\bm{E^{x}}_{j}$. The fused feature $F_{fused}$ is obtained by concatenating the linearly weighted experts,

For the final prediction $\hat{y}$, we utilize a Multi-Layer Perceptron ($MLP$) as the classifier. The cross-entropy (CE) loss is employed as the supervision for prediction. The classification loss $\mathcal{L}_{cls}$ is defined as:

The final loss $\mathcal{L}$ can be defined as the weighted sum of aforementioned losses,

where $\alpha$ and $\beta$ are balance factors.

We utilize the 3D ResNet18 [30] as the backbone for feature extraction, and the parameters of these backbones are not shared. The dimension of the each extracted latent feature $x_{j}$ is 512. The shared-encoder $\bm{E^{c}}$, each private-encoders $\bm{E^{p}}_{j}$, each partial-shared-encoders $\bm{E^{s}}_{jk}$, $FC$ and each expert $\bm{E^{x}}_{j}$ are all implemented as one fully-connected layer with $dim$ neurons. The $MLP$ consists of two fully-connected layers with $dim$ neurons and one output layer with $num\_cls$ neurons, where $num\_cls$ represents the number of classes. Each fully-connected layer in $MLP$ is followed by a ReLU layer and a Dropout layer. We empirically set $dim$ as 32.

We utilize three multimodal MRI datasets, which consist of two public datasets and one private dataset, to verify the effectiveness of the proposed framework. Some cases from these three datasets are shown in Fig 4.

We also verify the effectiveness of the proposed CFD strategy and DMF module. Ablation studies are conducted on both adopted datasets, and the results are summarized in Table 3 and Table 4, respectively. In the tables, the models are termed as baseline1, baseline2, …, baseline6 from the top row to the bottom row, with baseline6 representing the proposed framework. The ablation studies consider three factors: $dis\_ps$, $MoE$ and $LinG$. The $dis\_ps$ factor, which related to CFD strategy, determines whether to decouple modality-partial-shared features during the feature decoupling process. There are two factors related to the DMF module, namely $MoE$ and $LinG$. The $MoE$ factor represents whether to adopt MoE for the feature fusion, with “$\times$” indicating fusion with concatenation operation. The $LinG$ factor denotes whether to utilize proposed LinG_GN to generate weights in MoE, with “$\times$” representing the use of concatenation of decoupled features as the input of the gating network. The concatenation operation for inputs in gating network is the common setting in MoE-based MML methods [42, 29].

Specifically, the order of the three modalities is PD-axial, T1-coronal and T2-sagittal on the MRNet dataset, and T1C, Flair-C, ADC on the MEN dataset.

To enhance the interpretability of our framework, Grad-CAM [44] was employed to visualize the activation maps of the network for each modality in the MEN dataset. As shown in Fig 7, the model pays more attention to the tumor-edema junction across different modalities, with variations in the specific areas of focus. For the T1C modality, the model predominantly focuses on the tumor and its surrounding region. In the FLAIR-C modality, the model highlights the tumor and edema regions. Conversely, in the ADC modality, the model primarily attends to the edema region. Moreover, based on the feature weights in the DMF module (Fig. 6 (c)), the characteristics of the tumor-edema junction ($F$, the final modality-shared feature), the edema ($G_{23}$, the final modality-partial-shared feature between FLAIR-C and ADC) and the tumor-surrounding region ($P1$, the modality-specific feature of T1C) contribute significantly to the final task.

To validate the effectiveness of our framework, we visualized the feature space distribution using the Manifold Discovery and Analysis (MDA) [45] algorithm on the MEN dataset, randomly selecting balanced samples across categories for optimal visualization. Fig. 8 demonstrates that categorical separation becomes increasingly distinct in deeper network blocks, with B2-C2 through B5-C2 showing progressive improvement in class discrimination. The temporal evolution of feature distributions across training epochs (Fig. 9 (a-c)) reveals increasingly pronounced categorical boundaries. Furthermore, we examined the model’s robustness by analyzing feature distributions under Gaussian noise perturbation of inputs, as shown in Fig. 9 (d). The results showed that the feature distributions maintained clear categorical separation, demonstrating the method’s resilience to input variations.

In order to assess the generalization of the proposed framework on a wider range of modalities, we extend the proposed framework to MGMT promoter status prediction dataset (BraTS 2021) in the four-modal case, and the concept map is shown in Fig. 10. In this case, the complexity increases as we needed to decouple not only the pair-wise modality-partial-shared features but also the triplet-wise modality-partial-shared features. This leads to a total of 32 decoupled features and 15 final decoupled features.

The comparison results on BraTS 2021 dataset are listed in Table 5. The seven SOTA methods obtain comparable results. Our proposed framework achieves the best results in five metrics, including ACC ($0.6137$, $0.0121$ better than the 2nd), G-Mean ($0.6089$, $0.0114$ better than the 2nd), Ba_ACC ($0.6123$, $0.0127$ better than the 2nd), AUPRC ($0.5934$, $0.0085$ better than the 2nd), AUC ($0.6177$, $0.0173$ better than the 2nd). The statistical test results obtained on the BraTS 2021 dataset are similar to those of the MRNet and MEN datasets.

However, it is important to acknowledge that as the number of modalities increases, the relationships between them become more intricate, posing challenges for the CFD strategy. Additionally, incorporating more modalities increases the number of network parameters, further complicating network optimization. However, in clinical studies, commonly used multimodal MRI datasets typically contain two to four modalities. Our proposed framework has been designed to achieve effective performance under these conditions.

As shown in Fig. 11, we compare the number of parameters and GFLOPs between the proposed method and the comparison methods⁴⁴4We compare only with methods that utilize the same backbone as the proposed approach. For clarity, MAML is not shown in Fig. 11, as its parameter and GFLOPs values are significantly larger than those of other methods. Specifically, in the three-modality case, MAML has 202.61M parameters and 165.63 GFLOPs, while in the four-modality case, these values increase to 270.09M parameters and 167.89 GFLOPs.. In the three-modal case (left part of Fig. 11), using MRNet as an example, the proposed method achieves the best performance with the fewest parameters and the lowest GFLOPs. In the four-modal case (right part of Fig. 11), while EmbraceNet and CCML have the fewest parameters and GFLOPs, our method achieves a significant performance improvement with a comparable number of parameters and relatively low GFLOPs. Although the number of parameters increases for all methods from the three-modal case to the four-modal case, our proposed method achieves superior performance with only a modest increase in parameters.