Cross-Modality Deep Feature Learning for Brain Tumor Segmentation

Brain tumor segmentation is a hot topic in the medical image analysis and machine learning community. It has received great attention in the past few years. Early efforts in this filed designed hand-crafted features and adopted the classic machine learning models to predict the brain tumor areas. Due to the rapid development of the deep learning technique, the recent brain tumor segmentation approaches mainly apply the deep features and classifiers from the DCNN models. Based on the type of the convolutional operation used in the DCNN models, we briefly divide the existing methods into two groups, i.e., the 2D CNN-based methods and 3D CNN-based methods. The 2D CNN-based methods [4, 5, 6] apply the 2D convolutional operations and split the 3D volume samples into 2D slices or 2D patches. While the 3D CNN-based methods [7, 8, 9] apply the 3D convolutional operations, which can take the whole 3D volume samples or the extracted sub 3D patches as the network input.

Although these deep learning-based methods can already obtain much powerful feature representation when compared to the early classical methods that are based on the hand-crafted features, they did not make full use of the multi-modality data in the feature learning process, which limits the effectiveness of the learned feature representations and the final segmentation results. Realizing this issue, Fidon et al. [10] proposed a multi-modal convolutional network for brain tumor segmentation, where nested network structure was designed to explicitly leverage deep features within or across modalities. Different from our approach, they did not formulate the across modality transition process and did not employ the mask guidance scheme in the feature fusion process.

Multi-modality feature learning is gaining more and more attention in the recent years as the multi-modality data can provide richer information for sensing the physical world. Existing works have applied multi-modality feature learning in many computer vision-based tasks such as 3D shape recognition and retrieval [11], survival prediction [12], RGB-D object recognition [13] and person re-identification [14]. Among these methods, Bu et al. [11] built a multi-modality fusion head to fuse the deep features learnt by a CNN network branch and a Deep Belief Network (DBN) branch. To integrate multiple modalities and eliminate view variations, Yao et al. [12] designed a deep correlational learning module for learning informative features on the pathological data and the molecular data. In [15], Wang et al. proposed a large-margin multi-modal deep learning framework to discover the most discriminative features for each modality and harness the complementary relationship between different modalities.

Although the multi-modality feature learning technique has been applied in many computer vision tasks, it is still a under-studied issue in the research field of medical image understanding, especially for the task of brain tumor segmentation. To this end, this paper makes an early effort to build a cross-modality deep feature learning framework for brain tumor segmentation. The cross-modality feature transition (CMFT) process and the cross-modality feature fusion (CMFF) process designed in this work are also novel to the existing multi-modality feature learning methods.

As shown in the left part of Fig. 2, given modality $A$ and modality $B$, we adopt the generative adversarial learning strategy to facilitate the knowledge transition across the different modality data, which in turn captures the informative patterns from each modality data. To be specific, for each modality data, we build a generative network, i.e., the generator $G$, and a discriminative network, i.e., the discriminator $D$, to formulate the feature transition process. For achieving this goal, we apply the CycleGAN learning scheme [16] to learn the transition $G^{A}_{B}:A\rightarrow B$ and $G^{B}_{A}:B\rightarrow A$ so that

where A and B indicates the “real” input sample from the modality $A$ and modality $B$, respectively. Compared with other generative adversarial learning schemes, the cycle consistency-based learning scheme adopted by the CycleGAN model has the following advantages for learning representative features: Firstly, it learns the transition $G^{A}_{B}:A\rightarrow B$ and $G^{B}_{A}:B\rightarrow A$ simultaneously, thus facilitating a better exploration of the relationship between the two modality data and maintaining the content of each modality data. Secondly, during the training process, it does not necessarily require matched modality data which might be hard to obtain in practical applications.

Besides the generators, there are also two discriminators $D_{A}$ and $D_{B}$, where $D_{A}$ distinguish the “fake” $A$-modality data generated by $G^{B}_{A}(\textbf{B})$ from the “real” $A$-modality data while $D_{B}$ distinguish the “fake” $B$-modality data generated by $G^{A}_{B}(\textbf{A})$ from the “real” $B$-modality data. During the generative adversarial learning process, we adopt the adversarial loss to match the distribution of the generated fake data to the distribution of the “real” data. To this end, the adversarial loss is defined as:

where $\mathbb{E}_{M}[\tau]$ indicates the expectation of $\tau$ for all the samples from modality M.

In addition, we also follow [16] to apply the cycle consistency loss to constrain the modality transition function $G^{A}_{B}$ and $G^{B}_{A}$ from random permution in the target modality domain. To enforce the modality transition function $G^{A}_{B}$ and $G^{B}_{A}$ to be cycle consistent, we encourage $G^{A}_{B}$ to transit the generated “fake” $A$-modality data $G^{B}_{A}(\textbf{B})$ back to the “real” $B$-modality data, and similarly encourage $G^{B}_{A}$ to transit the generated “fake” $B$-modality data $G^{A}_{B}(\textbf{A})$ back to the “real” $A$-modality data. To this end, the cycle consistency loss is defined as:

By considering both the adversarial loss and the cycle consistency loss, the full learning object function of the cross-modality feature transition process becomes:

where

$\lambda$ is a hyper-parameter to weigh the adversarial loss and the cycle consistency loss.

Network Architecture: When designing the generator, we adopt a U-net architecture due to its effectiveness in both image-to-image translation [17] and brain tumor segmentation [6, 4, 9]. Considering the training samples are in form of 3D volumes, we adopt 3D convolutions in the network layers, thus obtaining the 3D U-net architecture. The concrete network architecture is shown in Fig. 3. For the discriminator, we follow the existing work [16] to construct it by using several convolutional layers to obtain the classification results. The concrete network architecture of the discriminator is shown in Table 1.

To implement the cross-modality feature fusion process, we establish a novel cross-modality feature fusion network for brain tumor segmentation. Equipped with the newly designed fusion branch which uses the single-modality prediction results to guide the feature fusion process, the proposed network can not only transfer the features learned from the feature transition process conveniently but also learn powerful fusion features for segmenting the desired brain tumor areas.

Given the input data from modality $A$ and $B$, the cross-modality feature fusion network contains two single-modality feature learning branches $S_{A}$ and $S_{B}$ and a cross-modality feature fusion branch $S_{F}$ for segmenting the desired brain tumor areas. Specifically, the single-modality feature learning branch $S_{A}$ takes the $A$-modality data as the input and learns representative features to predict the segmentation masks of the brain tumor areas $S_{A}(\textbf{A})$ as the output. Similarly, the single-modality feature learning branch $S_{B}$ takes the $B$-modality data as the input and learns representative features to predict the segmentation masks of the brain tumor areas $S_{B}(\textbf{B})$ as the output. The cross-modality fusion branch takes the deep features as well as the predicted segmentation masks of the two single-modality feature learning branches as input to learn more powerful fusion features to generate the final segmentation masks of the bairn tumor areas $S_{F}(\textbf{A},\textbf{B})$. To learn the cross-modality feature fusion network, we introduce the following object function:

To prevent the model from being heavily affected by the unbalance among different types of tumor areas, we follow [18] to calculate $\mathcal{L}_{seg}(S_{A})$, $\mathcal{L}_{seg}(S_{B})$, and $\mathcal{L}_{seg}(S_{F})$ by the Dice Similarity Coefficient (DSC). Thus, for $\mathcal{L}_{seg}(S_{A})$, we have

where Y is the ground-truth annotation for the desired brain tumor areas. It goes the same for $\mathcal{L}_{seg}(S_{B})$ and $\mathcal{L}_{seg}(S_{F})$.

Network Architecture: Although the single-modality feature learning branch does not necessarily be the same with the generator used in cross-modality feature transition, the more network layers shared by these two networks, the richer features can be conveniently transferred from the feature transition process to the feature fusion process. To this end, we adopt a quite similar network architecture to the generator $G^{A}_{B}$ (or $G^{B}_{A}$) to build the single-modality feature learning network branches $S_{A}$ and $S_{B}$ (see the right part of Fig. 2). Compared to the generator, the only difference is the number of kernels set to the last convolutional layer. As shwon in Fig. 3, the last convolutional layer of the single-modality feature learning network branch uses four convolutional kernels, while the generator only uses one convolutional kernel in the last convolutional layer. As can be seen, designing the single-modality feature learning network branch in this way could share the most network layers with the generator and thus can take full advantage of the features learned from the cross-modality feature transition process.

For fusing the knowledge from each modality data, we propose a novel cross-modality feature fusion branch. As shown in Fig. 4, the proposed cross-modality feature fusion branch contains several convolutional layers to fuse deep features from different layers of the two single-modality feature learning network branches. The convolutional layers are then followed by a mask-guided attention block to learn more powerful fusion features for brain tumor segmentation. Different from the conventional attention modules, such as [19, 20], the attention masks in our mask-guided attention block are the segmentation masks predicted by the single-modality feature learning branches rather than those inferred from the deep feature maps from previous network layers. In other words, the attention masks in the conventional attention network blocks/modules are used to guide the network learning on its own network branch. They are learned in a bottom-up manner. In contrast, the attention masks in this work are used to guide the network learning on a different network branch and they are learned in a top-down manner.

As the data used in the investigated brain tumor segmentation task usually have four modalities, i.e., the T1, T1-c, T2, and FLAIR modality (see Fig. 1), we also explore effective extension strategies to enable the aforementioned cross-modality deep feature learning framework to work on the modality quaternions. An naive extension is to adopt six cycGAN models, i.e, $\{G^{A}_{B},G^{B}_{A}\}$, $\{G^{A}_{C},G^{C}_{A}\}$, $\{G^{A}_{D},G^{D}_{A}\}$, $\{G^{B}_{C},G^{C}_{B}\}$, $\{G^{B}_{D},G^{D}_{B}\}$, $\{G^{C}_{D},G^{D}_{C}\}$, to learn the transition functions between each modality data and fuse the four single-modality feature learning branches in the cross-modality feature fusion network. Although this strategy can also learn rich feature representations from both the cross-modality feature transition and cross-modality feature fusion processes, it requires too large computational cost to implement in practice.

To this end, we propose a simple yet effective way to implement the learning framework on modality quaternions. Instead of transiting knowledge between each modality data, we implement the transition process between each modality pair. That is to say, the transition process is extended to transit knowledge from a modality pair to another modality pair. In this work, we use the T1 and T1-c modalities to form a modality pair while T2 and FLAIR modalities to form another modality pair. In this way, the information within each modality pair tends to be consistent while the information from different modality pairs tends to be distinct and complementary (see Fig. 5), which enables the cross-modality feature transition process to learn rich patterns. Based on this strategy, we implement the proposed approach on modality quaternions by simply converting the input data of the generators and discriminators in the CMFT process and the input data of the feature learning branch in the CMFF process to be the concatenation of two modality data, while other parts of the learning framework remain unchanged (see Fig. 6).

As described in previous sections, the proposed learning framework contains two processes, i.e., the CMFT process and the CMFF process. In fact, these two processes can also be considered as two learning phases of a unified DCNN model. Specifically, imaging that we have a DCNN model contains two generators, two discriminators and a fusion network branch, our approach trains the two generators and the two discriminators in the first learning phase and then trains the two generators (with a modified prediction layer and loss function) together with the fusion network branch in the second learning phase. From this point of view, our proposed deep learning framework can be seen as a unified end-to-end learning model with two-phase training strategy.

Besides the two-phase training strategy, we can actually learn CMFT and CMFF simultaneously, where both Eq. 5 and Eq. 7 would be introduced to form the new objective function of each training sample. However, by simultaneously learning the two generators, the two discriminators and the fusion network branch, this strategy has too much memory costs especially when exploring the 3D volume data like in this task. Thus, we choose to adopt the two-phase training strategy to implement our approach.

In the BraTS 2017 and BraTS 2018 benchmark datasets, there are four modalities, i.e., T1, T1-c, T2, and FLAIR, for each patient. The BraTS 2017 benchmark has two sub-sets: a training set, which contains 285 subjects, and a validation set containing 46 subjects with hidden ground truth. The BraTS 2018 benchmark contains the same number of subjects in its training set but has 66 subjects in the validation set with hidden ground truth. When implementing the experiments on each of the benchmarks, we use the training set to train the brain tumor segmentation models while use the validation set to test the segmentation performance. We adopted the official metrics that are used by the online evaluation system of BraTS for quantitative evaluation. They are the Dice score, Sensitivity, Specificity, and the $95th$ percentile of the Hausdorff Distance (HD95).

Before training, each of the input modality data was normalized to have zero mean and unit variance. We randomly sampled patches of size $128\times 128\times 128$ within the brain tumor area as the inputs of both the cross-modality feature transition model and the cross-modality feature fusion model. As a trade-off between performance and memory consumption, the base number of filters in the U-Net was designed to be 16, which was increased to twice after each down-sampling layer. The Adam optimizer with an initial learning rate of $10^{-4}$ was applied to optimize the objective function, where $\lambda$ was set to be 10. When training the cross-modality feature fusion network, the pre-trained parameters of the transition mappings $G^{A}_{B}$ and $G^{B}_{A}$ were transferred to the $S_{A}$ and $S_{B}$ for further fine-tuning. The $S_{A}$ and $S_{B}$ took the same input modality data as the $G^{A}_{B}$ and $G^{B}_{A}$. The parameters of the cross-modality fusion branch are randomly initialized. We used the Adam optimizer with an initial learning rate of $10^{-4}$ and a batch size of 1 to train this network branch. All of the network branches were implemented in Pytorch on a NVIDIA GTX 1080TI GPU. It totally takes 18 hours and 57 minutes to complete the training process and the test speed is 3.2 seconds per subject.

In this subsection, we evaluate the proposed approach on the BraTS 2017 benchmark. We first analyze the effect of the main network branches of the proposed learning model by conducting the experiments on the following baseline models. The first two baseline models train the single-modality-pair feature learning branches $S_{A}$ and $S_{B}$ with the input modality data {T1,T1c} and {T2,FLAIR}, respectively. The third baseline model “$S_{A}+S_{B}$” fuses the prediction of $S_{A}$ and $S_{B}$ by directly computing the average of the obtained segmentation maps. Then, we compare our approach with the baseline models “Ours w/o MG” and “Ours w CA” which adopt the proposed cross-modality feature fusion branch but without using the mask-guided attention block or directly using the conventional attention block [21]. All the aforementioned baseline models are fine-tuned based on the network parameters obtained from the cross-modality feature transition process. The experimental results are reported in top rows of Table 2.

By comparing $S_{A}$, $S_{B}$ and our approach, we can observe that simply using a single-modality-pair feature learning branch only obtains poor performance due to the inadequate modality information. The performance of $S_{A}+S_{B}$ is better than $S_{B}$ but worse than $S_{A}$, which might be caused by the large performance gap between $S_{B}$ and $S_{A}$. By comparing “Ours w/o MG”, $S_{A}+S_{B}$, and “Ours” we can observe that using the proposed feature fusion branch can significantly improve the feature learning capacity of our approach and using the mask-guided attention block can further improve the segmentation accuracy. Notice that when using the conventional attention block, the network works better for the ET area but worse for the TC and WT areas, making the average performance of “Ours w CA” less effective than “Ours w/o MG” and “Ours”.

In addition, we also conducted the ablation study by implementing three baseline models which directly train the CMFF network to obtain the segmentation results without the CMFT process. The first baseline “Ours_random” used the random values to initialize the CMFF network, while the second baseline “Ours_voc” used the parameters pre-trained on the PASCAL VOC segmentation dataset [33]⁷⁷7PASCAL VOC segmentation dataset is a large-scale image set that consists of RGB images and the corresponding segmentation annotation. to initialize the CMFF network. To facilitate the parameter transferring between the 2D image data and 3D volume data, we first trained a 2D-Unet on the PASCAL VOC segmentation dataset and then extended its convolution kernels to 3D convolution kernels as in [34]. For the third baseline “Ours_self”, we replaced the proposed CMFT process by a self reconstruction-based feature learning process that learns patterns by reconstructing the input data.

The experimental results are reported in bottom rows of Table 2. From the comparison results, we can observe that 1) due to the inadequate of medical imaging data, directly training the DCNN models with random parameter initialization is not able to achieve satisfying learning performance; 2) while using the large-scale RGB image data (together with the segmentation annotation) still cannot solve this problem because of the large domain gap; and 3) the proposed cross-modality feature transition process can learn informative features from the medical imaging data without using any human annotation, which also works better than the self reconstruction-based learning strategy.

Next, we compare the proposed approach with several state-of-the-art methods, which include six ensemble methods and five single prediction methods. The ensemble methods integrate multiple deep brain tumor segmentation models that are trained from different views or different training sub-sets to obtain the predicted segmentation masks for each test data, while the single prediction methods only apply one deep model to fulfill the brain tumor segmentation task. Thus, the ensemble methods can usually obtain better performance but with higher complexity both in computational cost and time consumption. The quantitative results are reported in Table 3. From Table 3, we can observe that as a single prediction method⁸⁸8Although our model has a cross-modality feature transition process and a cross-modality feature fusion process, the cross-modality feature transition process only learns features and does not predict segmentation results. In other words, our segmentation results are predicted by the cross-modality feature fusion process only rather than the combination of the segmentation results obtained by both processes. Thus, our approach is considered as a single prediction method rather than an ensemble method., our proposed approach outperforms all the state-of-the-art single prediction methods both in terms of Dice score and Hausdorff95. More encouragingly, our approach can also obtain better performance than most (nine out of ten) ensemble methods. Thus, the comparison results in Table 3 demonstrate the effectiveness of the proposed approach.

On the larger-scale BraTS 2018 benchmark, we first compare the proposed approach with five baseline models, including “$S_{A}$”, “$S_{B}$”, “$S_{A}+S_{B}$”, “Ours w/o MG”, and “Ours w CA” to analyze the effect of the main network branches designed in our learning framework. The experimental results are reported in top rows of Table 4. Being consistent with the results on the BraTS 2017 benchmark, there is obvious performance gap between “$S_{A}$” and “$S_{B}$” and the straightforward fusion strategy “$S_{A}+S_{B}$” can only obtain performance better than “$S_{B}$” but worse than “$S_{A}$”. Compared to “$S_{A}+S_{B}$”, our approach obtains 4.6% performance gain (in terms of the Dice score), which demonstrates that the feature fusion branch proposed by our approach plays an important role in fusing informative features and predicting accurate tumor areas. Notice that “Ours w CA” obtains better performance than “Ours w/o MG” on this dataset. But its performance is still worse than “Ours”. Some examples of the comparison results on the BraTS 2018 validation set are shown in Fig.7. For better understanding the segmentation results, we also shown examples of our approach on the BraTS 2018 training set with the corresponding ground-truth annotations (see Fig.8). Besides, we also study the failure cases in Fig. 9, from which we can observe that the main challenges to our approach are the LGG cases when the ground-truth tumor areas are with absent ET area, discontinuous tumor regions, or ragged tumor contours.

To evaluate the effectiveness of the proposed CMFT process, we also compare our approach with the “Ours_random”, “Ours_voc”, and “Ours_self” baselines. The experimental results are reported in bottom rows of Table 4, from which we can observe obvious performance gain when compare our approach to the aforementioned baseline methods. Some examples of the comparison results are shown in Fig. 7, which can better illustrate the advantage of our approach. In addition, to verify the effectiveness of our strategy to build the modality pairs as described in Sec. 3.3, we further implement a baseline model which constructs the modality pair A by using the T1 modality and FLAIR modality and modality pair B by using the T2 modality and T1-c modality. Based on our experiment, this baseline obtains 0.822 Dice score, 0.844 sensitivity, and 5.789 Hausdorff Distance on the BraTS 2018 dataset. The comparison between this baseline and the proposed approach demonstrates the effectiveness of our approach in making the information contained within each modality pair relatively consistent and the information contained across the different modality pairs relatively distinct and complementary.

Finally, we compare the proposed approach with other state-of-the-art methods on the BraTS 2018 benchmark, which include three ensemble models [35, 36, 37] and three single prediction models [38, 39, 40]. It is worth mentioning that as different works adopt various ways to obtain their ensemble models and the concrete processes for obtaining the ensemble model are not clear to us, it is hard to implement an ensemble model that could compare with the existing ensemble models fairly. However, from the experimental results reported in Table 5, we can observe that our single-prediction model has already achieved better performance when compared to the ensemble models of [37, 41]. When compared to the state-of-the-art single prediction models, our approach also obtains the outperforming performance both in terms of Dice score and Hausdorff95. Thus, we believe the above experiments have already demonstrated the effectiveness of the proposed approach.