[OM-Net] One-pass Multi-task Networks with Cross-task Guided Attention for Brain Tumor Segmentation

0) Motivation, Object and Related works:

Motivation:

• Class imbalance has emerged as one of the major challenges for medical image segmentation. The model cascade (MC) strategy, a popular scheme, significantly alleviates the class imbalance issue via running a set of individual deep models for coarse-to-fine segmentation. Despite its outstanding performance, however, this method leads to undesired system complexity and also ignores the correlation among the models. 

Objectives:

• To handle these flaws in the MC approach, we propose in this paper a light-weight deep model, i.e., the One-pass Multi-task Network (OM-Net) to solve class imbalance better than MC does, while requiring only one-pass computation for brain tumor segmentation. 

  1. First, OM-Net integrates the separate segmentation tasks into one deep model, which consists of shared parameters to learn joint features, as well as task-specific parameters to learn discriminative features. 

  2. Second, to more effectively optimize OM-Net, we take advantage of the correlation among tasks to design both an online training data transfer strategy and a curriculum learning-based training strategy. Third, we further propose sharing prediction results between tasks, which enables us to design a cross-task guided attention (CGA) module. By following the guidance of the prediction results provided by the previous task, CGA can adaptively recalibrate channel-wise feature responses based on the category-specific statistics. 

  3. Finally, a simple yet effective post-processing method is introduced to refine the segmentation results of the proposed attention network. Extensive experiments are conducted to demonstrate the effectiveness of the proposed techniques. Most impressively, we achieve state-of-the-art performance on the BraTS 2015 testing set and BraTS 2017 online validation set. Using these proposed approaches, we also won joint third place in the BraTS 2018 challenge among 64 participating teams. 

1) Model:

• Firstly, we split multi-class brain tumor segmentation into three segmentation tasks: 

  1. Coarse segmentation to detect complete tumor. 

  2. Refined segmentation for complete tumor and its intra-tumoral classes. 

  3. Precise segmentation for enhancing tumor.

• For model cascade framework, each task is implemented by an independent basic network. Our basic network used in model cascade framework is a 3D variant of the Fusionnet. See picture below.

• OM-Net integrates three tasks with their respective training data being the same as those in model cascade framework. Its architecture is illustrated in the picture below. Each task owns its task-specific parameters and the shared backbone model is aimed to learn the correlation among the tasks. Specifically, the shared backbone model refers to the network layers outlined by the yellow dashed line.

• Cross-task guided attention (CGA) module is further proposed to share prediction results between tasks of OM-Net, which can make use of the prediction results generated by the previous task to produce finer category-specific statistics. Thus, CGA can adaptively recalibrate channel-wise feature responses based on these category-specific statistics, so as to boost the segmentation performance for the current task.

• 
  

References:

1. C. Zhou, C. Ding, Z. Lu, X. Wang, and D. Tao, “One-pass Multi-task Convolutional Neural Networks for Efficient Brain Tumor Segmentation,” in Proc. Int. Conf. Med. Image Comput. Comput.-Assist. Intervent, 2018, pp. 637–645.

2. C. Zhou, C. Ding, X. Wang, Z. Lu, and D. Tao, “One-pass Multi-task Networks with Cross-task Guided Attention for Brain Tumor Segmentation,” IEEE Trans. Image Process., doi: 10.1109/TIP.2020.2973510, 2020.