Representation Disentanglement for Multi-modal Brain MRI Analysis

We assume each subject in the training set has MRIs of $m$ modalities (sequences) and let $x_{i}\in\mathcal{X}_{i}$ denote the input image of the i-th modality. As shown in Fig. 1a, we aim to disentangle $x_{i}$ into an anatomical representation $s_{i}$ by an anatomical encoder $s_{i}=E_{i}^{A}(x_{i})$ and a modality representation $z_{i}$ by a modality encoder $z_{i}=E_{i}^{M}(x_{i})$. We assume that the anatomical representation $s_{i}$ encodes the morphological information of brain structures that is mostly impartial to the imaging modality, while $z_{i}$ provides image appearance information specific to a modality. The decoder $D$ then reconstructs $x_{i}$ from a pair of anatomical and modality representations. Prior works [2] have suggested that such disentangled representations can be learned by optimizing the self-reconstruction and cross-reconstruction losses; Given a pair of $s_{i}$ and $z_{j}$ derived from images of any two modalities, $D$ is supposed to synthesize an image that is similar to the input image $x_{j}\in\mathcal{X}_{j}$, whose synthesized domain corresponds to the j-th modality.

where $\widetilde{x_{ij}}=D(E_{i}^{A}(x_{i}),E_{j}^{M}(x_{j}))$. In addition to these reconstruction losses, another loss function commonly used for training image-to-image translation [14] is the latent consistency loss, which encourages the representations derived from raw inputs to be similar to the ones from the synthesized images.

where $\widetilde{z_{ji}}=E_{i}^{M}(D(E_{j}^{A}(x_{j}),E_{i}^{M}(x_{i})))$ is the modality representation derived from a synthesized image.

We resolve this problem by exploring the similarity relationships between representations. As the brain’s morphological shape is highly heterogeneous across subjects, we assume the anatomical representations $s$ from the same subject but different modalities should be more similar than those from the same modality but different subjects. Note, $s_{i}$ of the same subject are not necessary to be exactly the same, as multi-modal imaging is designed to capture distinct characteristics of brain anatomies. For instance, the brain tumor itself is more visible on T1-weighted MR with contrast (T1c) compared to T1 without contrast due to the injected contrast medium. On the other hand, modality representations $z$ from the same modality but different subjects should be more similar than those from the same subject but different modalities. We propose to model such relationships using a similarity loss inspired by the margin-based hinge loss [5].

where $p$ and $q$ correspond to a pair of subjects randomly sampled in a mini-batch, $\text{cos}(\cdot,\cdot)$ denotes the cosine distance between two vectors, and $f$ denotes a MaxPooling and flattening operation. Unlike the L₂-based similarity loss, Eq. (2) encourages the within-subject and across-subject distances to differ by the margins $\alpha_{s}$ and $\alpha_{z}$ and thereby avoids deriving identical representations.

As shown in Fig. 1b, after obtaining the disentangled representations, the anatomical representations from all available modalities of a subject are fused into one fixed-size encoding as the input for a downstream model $T$, which can be any state-of-the-art model for the downstream task. Note that the fusion function here should pool features of a various number of channels to a fixed number of channels. Let $s$ be the concatenation of anatomical representations from the available modalities $Concat(s_{i},...,s_{j})$. Then the fusion is the concatenation of several pooling functions: $Concat(MaxPool(s),MeanPool(s),MinPool(s)).$ With this fusion operation, one can use two strategies to train the downstream model $T$. We can either solely train $T$ based on the frozen $s$ derived by the self-supervised learning of the encoders (Section 3.1), or fine-tune the encoders jointly with the downstream task model $T$. Though the joint training can potentially result in representations that better suit the downstream task, we confine our analysis to the first strategy (fixing encoders) in this work to emphasize the impact of representation disentanglement.

ZeroDose The dataset comprised brain FDG-PET and two MR modalities (T1-weighted and T2 FLAIR) from 171 subjects with multiple diagnosis types including tumor, epilepsy, and dementia. The FLAIR and PET images were first registered to T1, and then all modalities were normalized to a standard template and resized to $192\times 160$ in the axial plane. Intensities in the brain region of each image were converted to z-scores. The top and bottom 20 slices in each image were omitted from analysis. Each 3 adjacent axial slices were converted to a 3-channel image as the input to the encoder models. Random flipping of brain hemispheres was used as augmentation during training. Five-fold cross-validation was conducted with 10% training cases used for validation. The downstream task was zero-dose PET reconstruction, i.e., to synthesize high quality FDG-PET from multi-modal MRIs. This is useful in practice as the injected radiotracer in current PET imaging protocols can lead to the risk of primary or secondary cancer in scanned subjects. Moreover, PET is more expensive than MRI and not offered in the majority of medical centers worldwide.

Implementation Details The anatomical encoder $E^{A}$ was a U-Net type model. Let C_k denote a Convolution-BatchNorm-ReLU block with $k$ filters ($4\times 4$ spatial filters with stride 2), and CD_k an Upsample-Convolution-BatchNorm-ReLU block. The architecture was designed as C₃₂-C₆₄-C₁₂₈-C₂₅₆-C₂₅₆-CD₂₅₆-CD₁₂₈-CD₆₄-CD₃₂. A convolution then mapped the resulting representations to 4 channels with softmax activation as the anatomical representation. The modality encoder $E^{M}$ consisted of 5 convolution layers of 3 x 3 filters and stride 2 with LeakyReLU of a 0.2 slope. Numbers of filters were 16-32-64-128-128. A fully connected layer mapped the resulting features to a 16-D representation. The decoder $D$ was based on SPADE [13] with the architecture used in [2]. The networks were trained for 50 epochs by the Adam optimizer with learning rate of $2\times 10^{-4}$ and weight decay of $10^{-5}$. The regularization rates were set to $\lambda_{c}=2.0$, $\lambda_{l}=0.1$, $\lambda_{s}=10.0$, $\lambda_{z}=2.0$. The margins in the similarity loss were set to $\alpha_{s}=\alpha_{z}=0.1$. For the downstream model $T$ in the zero-dose PET reconstruction, a U-Net based model with attention modules [12] was adopted. The downstream model for BraTS brain tumor segmentation was the BraTS 2018 challenge’s winner NVNet [11].

We first evaluated the methods on representation disentanglement and image cross reconstruction based on 5-fold cross-validation. We derived the anatomical and modality representations of the test subjects learned by the approaches. Fig. 2 visualizes the 4-channel anatomical representations of one subject from NCANDA. We observe that $s_{1}$ and $s_{2}$ (extracted from T1 and T2) learned by the two baselines (Conv+NA and Conv+Adv) were substantially different, indicating that they might still contain modality-specific information. On the other hand, our approach (Conv+Sim) produced visually more similar anatomical representations than the baselines. This indicates that the proposed similarity regularization can decouple the modality-specific appearance features from the structural information shared between modalities.

This result is also supported by the visualization of the learned representation spaces. As shown in Fig. 3a-c, we randomly selected 200 modality representations in the test set of the BraTS dataset and projected them into a 2D space by t-SNE. Only our approach clearly separated the representations of different modalities into 4 distinct clusters (Fig. 3c), which was in line with the regularization on $z$ in Eq. (2). The clustering with respect to modalities was not evident for the projections of the baseline approaches (Fig. 3a,b), indicating that complimentary information had leaked into the modalities representations. Moreover, the baseline approaches failed to disentangle T1 and T1Gd, two contrasts with high visual resemblance, as the red and blue dots were coupled in the representation space. Likewise, we visualized the space of anatomical representations in Fig. 3d-f. We randomly selected 4 subjects in the BraTS test set and projected the pooled anatomical representation $f(s_{i})$ of 4 consecutive slices into a 2D space. Now the 4 distinct clusters of our approach were defined with respect to subjects as opposed to modalities, and there was no apparent bias of modality in each subject’s cluster (Fig. 3f), indicating the representations solely encoded subject-specific yet modality-invariant information. The representation spaces learned by the two baselines contained both subject-specific anatomical and modality information (Fig. 3d,e); that is, although the projections could be separated by subjects (black circles), each subject-specific cluster could be further stratified by modalities (color of the markers).

The improved disentanglement also resulted in better cross-reconstruction. Fig. 4 shows the results of a test subject from the BraTS dataset. In each panel, $\widetilde{x_{ij}}$ is displayed on the $j^{th}$ row and $i^{th}$ column; Diagonal images correspond to self-reconstruction and off-diagonal ones are cross-reconstruction. The proposed Conv+Sim achieved the best visual quality (accurate structural details), especially the FLAIR reconstruction highlighted in red boxes, where the tumor area was more precisely reconstructed. This improvement was quantitatively supported by the higher similarity between ground-truth and synthesized images in terms of peak-signal-noise ratio (PSNR) and structural similarity index (SSIM) (Table 1). For each image of a specific modality, we synthesized it based on its own modality representation and the anatomical representation from another modality (For each image in the BraTS dataset, we computed the average metrics over all three cross reconstruction results). According to Table 1, Conv+NA achieved the lowest reconstruction quality for both SSIM and PSNR on all three datasets. The quality improved when adding adversarial training on the anatomical representations (Conv+Adv) but was still lower than the two implementations with the proposed similarity loss (except for T1 reconstruction in NCANDA). Of the two models with +Sim, CondConv+Sim recorded better performance on the ZeroDose and NCANDA datasets. This indicates that CondConv enabled more stable reconstruction results by coupling highly dependent modality-specific encoders. However, on the BraTS dataset, Conv+Sim achieved better cross-reconstruction for T2 and FLAIR. The reason could be that CondConv shared anatomical and modality encoders across modalities at the expense of model capacity, especially when more than two modalities were involved. It is also worth mentioning that all methods achieved higher performance on NCANDA because it was the only dataset of healthy controls. Taken all together, only our approach resulted in true disentanglement between anatomical and modality representations, which was not guaranteed by the baselines.

Deep learning models that rely on multi-modal input often suffer from the missing input problem. When a model is trained on data with complete inputs and tested on data with one modality missing, standard approaches either fill in all zero values (Standard+Zero) or use the average image of that modality over the entire cohort (Standard+Avg). Here, we demonstrate that an alternative solution is to train and test the model on the fusion of disentangled anatomical representations from all available modalities (Ours, CondConv+Sim). We show that this strategy can largely alleviate the impact of missing inputs.

In each run of the cross-validation, we first learned the disentangled representations and then trained the downstream models (based on the raw multi-modal images for the standard methods or fused anatomical representations for our proposed method) for zero-dose PET reconstruction. Then, the downstream model was tested on data with or without missing modalities. When all modalities were used for testing, the standard and proposed approaches achieved comparable reconstruction accuracy in terms of both PSNR and SSIM (N/A in Table 2), but we observe that our approach could generally result in more accurate tumor reconstruction (Fig. 5), which could not be reflected in the whole-brain similarity measure. The improvement of our approach became evident when one of the modalities was missing (Table 2). In particular, missing FLAIR induced larger performance drop for all approaches, as ZeroDose contained a large number of images with tumor, which was more pronounced in FLAIR than T1.

Next, we replicated this experiment for the downstream task of brain tumor segmentation on the BraTS dataset and measured the performance using the dice coefficient (DICE; Each cell in Table 2 records the average DICE across three categories: ET, ED, and NCR/NET). In line with the results of the ZeroDose experiment, the standard and proposed methods both obtained similar DICE scores, when complete inputs were used during testing. Standard+Zero recorded the lowest accuracy when missing any modality. Standard+Avg was less impacted when T1 or T2 was missing, but more impacted by T1Gd and FLAIR as the standard model relied mostly on those two modalities for localizing the tumor. The proposed method achieved the highest DICE score in all scenarios, among which the missing T2 recorded the largest drop on DICE. This might be because T2 had the most distinct appearance compared to other modalities, thus having the largest impact on the fused representation.