FANCL: Feature-Guided Attention Network with Curriculum Learning for Brain Metastases Segmentation

TL;DR

A novel model called feature-guided attention network with curriculum learning (FANCL) is proposed, based on CNNs, which can effectively compensate for the loss of high-level feature from small tumors with the information of large tumors.

Abstract

Accurate segmentation of brain metastases (BMs) in MR image is crucial for the diagnosis and follow-up of patients. Methods based on deep convolutional neural networks (CNNs) have achieved high segmentation performance. However, due to the loss of critical feature information caused by convolutional and pooling operations, CNNs still face great challenges in small BMs segmentation. Besides, BMs are irregular and easily confused with healthy tissues, which makes it difficult for the model to effectively learn tumor structure during training. To address these issues, this paper proposes a novel model called feature-guided attention network with curriculum learning (FANCL). Based on CNNs, FANCL utilizes the input image and its feature to establish the intrinsic connections between metastases of different sizes, which can effectively compensate for the loss of high-level feature from small tumors with the information of large tumors. Furthermore, FANCL applies the voxel-level curriculum learning strategy to help the model gradually learn the structure and details of BMs. And baseline models of varying depths are employed as curriculum-mining networks for organizing the curriculum progression. The evaluation results on the BraTS-METS 2023 dataset indicate that FANCL significantly improves the segmentation performance, confirming the effectiveness of our method.

Figures (11)

• Figure 1: Illustration of visual examples of BMs. Here, red, green, and blue represent the nonenhancing tumor, peritumoral edema and enhancing tumor, respectively. Small tumors are indicated by yellow arrows. The high heterogeneity in size, shape and density of metastases, along with the complex anatomical structure of the brain, makes the segmentation task extremely difficult.
• Figure 2: Visualization of the original image and its corresponding convolutional feature map. Here, Conv is a $3\times 3\times 3$ convolutional layer with the stride of 2. (a)$240\times 240$ 2D slice of the original image. (b)$80\times 80$ feature map slice of the original image generated by Conv. (c) The larger tumor in the original image. (d) The smaller tumor in the original image. (e) The feature map of the larger tumor. (f) The feature map of the smaller tumor. It is evident that there is a significant difference in tumor size between (c) and (d), but a small difference between (d) and (e). This indicates that the large tumors in the feature maps and small tumors in the original images have more comparable sizes. Meanwhile, (f) shows that the feature map of small tumors obtained through convolutional operation is very small, making it difficult to provide sufficient information.
• Figure 3: Architecture diagram of FANCL.
• Figure 4: The framework of our methods. (a) The overview of our model, where the baseline architecture is nnU-Net. The feature-guided attention mechanism obtains FGA from the input image and feature map of baseline, and adjusts the initial prediction via FGA to obtain the modulated result $\mathrm{R}_t$. We also calculate the golden FGA (GFGA) through the ground truth (GT) for supervision. (b) Model training pipeline under the CL strategy. In the $t$-th stage, model is supervised by $\mathrm{C}_t$ obtained in (c) and GT simultaneously. When $t\geq 2$, the model continues training from the previous stage. (c) Curriculum-mining strategy, which extracts different feature information in the input image through networks of different depths, and then forms curricula $\mathrm{C}_t$ for different stages. Our model in (a) and the curriculum-mining networks in (c) are trained independently. Besides, in testing phase, the prediction can be obtained by simply inputting the image into our model.
• Figure 5: Feature-guided attention mechanism. a: The calculation process of FGA. b: Use the FGA obtained in a to adjust the initial prediction of baseline. c: The generation process of GFGA.