PaperlayerView exact revision on arXiv

Source-linked AI summary

One-pass Multi-task Networks with Cross-task Guided Attention for Brain Tumor Segmentation

Chenhong Zhou, Changxing Ding, Xinchao Wang, Zhentai Lu, Dacheng Tao

arXiv:1906.01796v2cs.CVcs.AIcs.LGeess.IV

TL;DR

Manual brain-tumor segmentation from MRI is subjective and time-consuming, while model cascades add complexity and overlook task correlations. OM-Net unifies interconnected segmentation tasks in one pass, using shared and task-specific learning, cross-task guidance, and post-processing; experiments show advantages over MC and existing state-of-the-art methods.

  • Problem

    Manual brain-tumor segmentation from MRI is subjective and time-consuming, motivating automatic and robust segmentation tools.

  • Method

    OM-Net integrates interconnected segmentation tasks into one model with shared and task-specific parameters, curriculum-based data transfer, cross-task guided attention, and post-processing.

  • Results

    OM-Net outperforms MC with one-third of its model parameters, while experiments on three datasets demonstrate advantages over existing state-of-the-art brain-tumor segmentation methods.

  • Takeaways & Limitations

    One-pass computation and cross-task sharing reduce model size and system complexity while exploiting correlations among segmentation tasks.

Abstract

from arXiv · show

Class imbalance has emerged as one of the major challenges for medical image segmentation. The model cascade (MC) strategy 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. To handle these flaws, we propose 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. 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. 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 and design a cross-task guided attention (CGA) module which can adaptively recalibrate channel-wise feature responses based on the category-specific statistics. Finally, a simple yet effective post-processing method is introduced to refine the segmentation results. 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. The code is publicly available at https://github.com/chenhong-zhou/OM-Net.

I. INTRODUCTION

The paper addresses class imbalance and complexity in brain tumor segmentation by replacing sequential model cascades with OM-Net, a one-pass multi-task model that exploits task relationships.

  • Manual brain-tumor segmentation from MRI is subjective and time-consuming, motivating automatic and robust tools.
  • Severe class imbalance between healthy, tumor, and intra-tumoral classes degrades model optimization by making healthy tissue dominant during training.
  • OM-Net integrates three interconnected segmentation tasks into one model with one-pass prediction instead of three sequential networks.
  • Curriculum learning gradually adds tasks in increasing difficulty, while online data transfer shares training material across hierarchically related tasks.
  • Cross-task Guided Attention uses preceding-task predictions to learn category-specific channel statistics, unlike self-attention methods that assign one weight per channel.
  • Across BraTS 2015, 2017, and 2018, OM-Net outperforms MC with one-third of its model parameters, and CGA further improves performance.

A. Brain Tumor Segmentation

The paper reviews cascade-based brain tumor segmentation and attention mechanisms, then motivates OM-Net and CGA as ways to reduce cascade complexity and provide category-specific channel attention.

  • Model cascades address class imbalance through coarse-to-fine segmentation but are cumbersome and ignore correlations among tasks.
  • OM-Net uses shared parameters and training data to integrate the three cascade tasks while outperforming MC with one-third as many model parameters.
  • Existing attention methods include hard regional, soft spatial, and channel attention, with channel attention recalibrating feature responses across channels.
  • SE assigns each channel a single weight from average responses across all voxels, regardless of voxel category, which may be suboptimal for segmentation.
  • CGA uses OM-Net’s task structure to provide cross-task guidance for category-specific channel attention.
  • The method combines multi-task OM-Net, curriculum-based training, CGA, and post-processing to refine segmentation results.
  • The three tasks represent hierarchical complete-tumor, tumor-core, and enhancing-tumor regions, with NCR and NET merged into one class.

B. One-pass Multi-task Network

OM-Net integrates three cascade tasks through shared computation, curriculum learning, and online training-data transfer, producing simultaneous predictions during inference.

  • OM-Net combines the three MC tasks in one network while preserving each task’s original training data.
  • Curriculum learning introduces tasks progressively in increasing difficulty to improve convergence quality.
  • The first task trains tumor-versus-normal-tissue knowledge until its loss curve flattens before later tasks are added.
  • When the second task is added, task data are concatenated by batch, shared features are split, and task-specific parameters are optimized.
  • Training patches from easier tasks are transferred to harder tasks when they satisfy hierarchical sampling conditions.
  • During inference, concatenation, feature slicing, and data transfer are removed, allowing OM-Net to produce all three task predictions simultaneously.

C. Cross-task Guided Attention

CGA uses predictions from preceding tasks to provide category-specific guidance for channel attention. OM-Net obtains these coarse predictions through cross-task forward flow during training, without extra computation during testing.

  • Cross-task Guided Attention: CGA addresses the limitations of global average pooling by computing channel statistics in category-specific regions.The regions are estimated from predictions shared between tasks.
  • Cross-task Guided Attention: Feature-2 and Feature-3 pass through earlier tasks to generate coarse predictions for the second and third tasks, respectively.These operations are called cross-task forward flow and introduce negligible training computation with no testing overhead.
  • Cross-task Guided Attention: The CGA module contains category-specific channel importance and complementary segmentation blocks.The architecture applies CGA modules to task outputs and uses SoftmaxWithLoss for each task.

1) CSCI Block:

The CSCI block uses preceding-task probabilities and feature maps to estimate channel importance separately for tumor and non-tumor categories. This category-specific weighting is finer-grained than the single global value produced by an SE block.

  • 1) CSCI Block:: CSCI computes tumor and non-tumor probabilities from the preceding task’s prediction tensor.Tumor classes are grouped in Ctumor, while normal tissue and background form Cnon-tumor.
  • 1) CSCI Block:: Feature maps and category probabilities are reshaped, multiplied, and L1-normalized to estimate channel-importance vectors.The tumor vector mt and corresponding non-tumor vector encode category-specific channel importance.
  • 1) CSCI Block:: CGA evaluates channel importance for a specific category, whereas SE compresses global channel information into one value.This gives CGA finer category-specific statistics for recalibration.

2) CompSeg Block:

CompSeg uses tumor- and non-tumor-specific channel recalibration to produce complementary score maps, then combines them with probability-based spatial weighting for final prediction.

  • 2) CompSeg Block:: The CompSeg block recalibrates each feature channel separately with tumor and non-tumor importance vectors.The resulting maps Ut and Un highlight important channels and suppress less important ones for their respective categories.
  • 2) CompSeg Block:: Two complementary pathways focus separately on tumor and non-tumor voxels before classification.Each pathway feeds a 1 × 1 × 1 convolutional layer to produce its own score map.
  • 2) CompSeg Block:: The tumor and non-tumor score maps are merged by probability-weighted averaging and passed through another 1 × 1 × 1 convolution.This produces the ultimate prediction for the current task.
  • 2) CompSeg Block:: CGA reuses Pt and Pn first to derive channel importance and second as soft spatial masks for score-map fusion.The same category probabilities therefore guide both channel recalibration and spatial merging.
  • 2) CompSeg Block:: In CGA-core, cross-task guidance comes from the second OM-Net task, and Pt and Pn represent core and non-core tumor probabilities.The CGA-core and CGA-tumor structures otherwise have only minor differences.

D. Post-processing

The post-processing pipeline removes small isolated tumor clusters and applies K-means under specified volume conditions to correct likely edema misclassification. The second step improves tumor-core Dice without changing complete- or enhancing-tumor results.

  • D. Post-processing: The proposed post-processing method uses fewer steps than the method that inspired it and replaces hand-defined intensity thresholds with K-means clustering.It is presented as simple and flexible for refining network predictions.
  • D. Post-processing: The method removes isolated predicted tumor clusters smaller than τV OL = min(2000, 0.1 × Vmax).Vmax is the volume of the largest predicted 3D connected tumor area, and removing false positives can slightly improve complete-tumor Dice.
  • D. Post-processing: K-means is applied when vole/volt < 0.1, vole(n)/volt(n) < 0.05, and vole(n) < 1000.The clustering uses MRI intensity values to classify segmented edema voxels in the connected component.
  • D. Post-processing: The second post-processing step significantly improves tumor-core Dice while leaving complete-tumor and enhancing-tumor segmentation unchanged.It changes only voxels predicted as edema.
  • D. Post-processing: Experiments evaluate the approaches on BraTS 2018, BraTS 2017, and BraTS 2015 datasets using online evaluation for hidden test or validation sets.The datasets contain co-registered, interpolated, skull-stripped MRI images with four modalities.

B. Evaluation Metrics

The evaluation follows official dataset metrics, using Dice score as the main measure alongside PPV, Sensitivity, and Hausdorff distance. Training uses normalized MRI modalities, task-specific patch counts, and SoftmaxWithLoss.

  • Dice score, PPV, Sensitivity, and Hausdorff distance are the official evaluation metrics used across the datasets.Dice, PPV, and Sensitivity measure voxel-wise overlap, while Hausdorff distance measures surface-distance discrepancy.
  • Dice score is adopted as the main evaluation metric because it is consistently used across BraTS challenges.
  • MRI voxel intensities are normalized within the brain area to zero mean and unit variance for each modality.
  • The first, second, and third tasks use approximately 400,000, 400,000, and 200,000 training patches, respectively.

V. EXPERIMENTAL RESULTS AND DISCUSSION

Experiments evaluate the proposed components on BraTS 2018, showing that OM-Net improves over increasingly parameterized cascades while CGA and complementary fusion further enhance segmentation.

  • Increasing the model count in MC steadily improves Dice scores but also increases parameter count, storage consumption, and system complexity.
  • OM-Net achieves better segmentation than MC3 despite using only one-third of MC3’s parameters, especially for tumor-core and enhancing-tumor Dice scores.
  • 3) Effectiveness of Cross-task Guided Attention: OM-Net + CGA outperforms both OM-Net and OM-Net + SE, with a maximum 2.28% tumor-core Dice improvement over OM-Net on the local validation subset.On larger datasets, CGA consistently improves enhancing-tumor Dice, whereas OM-Net and OM-Net + SE show no clear performance difference.
  • 3) Effectiveness of Cross-task Guided Attention: Feature maps with high category-specific statistics respond strongly to tumor or non-tumor regions, while low-statistic channels are suppressed by CGA.These observations support the CSCI block’s generation of category-specific channel dependence.
  • 3) Effectiveness of Cross-task Guided Attention: OM-Net + CGA significantly outperforms CGA− because soft spatial masks fuse complementary prediction results in the CompSeg block.
  • 3) Effectiveness of Cross-task Guided Attention: OM-Net + CGA outperforms both CGAt and CGAn, indicating that both CompSeg pathways exploit complementary information beneficially.

4) Effectiveness of Post-processing:

The post-processing scheme refines OM-Net predictions, improving complete-tumor and tumor-core segmentation, while BraTS 2015 results show additional gains over MC3 and DeepMedic.

  • 4) Effectiveness of Post-processing:: Post-processing slightly improves complete-tumor Dice and increases tumor-core Dice by 2.6% for OM-Net.The reported gains are attributed to removing false positives in the first step and applying a second refinement step.
  • 4) Effectiveness of Post-processing:: Post-processing consistently improves complete-tumor and tumor-core performance when applied to OM-Net + CGA.
  • 4) Effectiveness of Post-processing:: The proposed methods steadily improve qualitative brain-tumor segmentation quality across comparisons with MC3, OM-Net, OM-Net + CGA, and OM-Net + CGAp.
  • B. Performance Comparison on BraTS 2015 Testing Set: OM-Net exceeds MC3 by 1% Dice on both tumor core and enhancing tumor in the BraTS 2015 testing set.
  • B. Performance Comparison on BraTS 2015 Testing Set: OM-Net + CGA adds 1% enhancing-tumor Dice, while post-processing adds 1% complete-tumor and 4% tumor-core Dice over OM-Net + CGA.
  • B. Performance Comparison on BraTS 2015 Testing Set: OM-Net + CGAp outperforms DeepMedic by 2%, 8%, and 2% Dice on complete tumor, tumor core, and enhancing tumor, respectively.
  • B. Performance Comparison on BraTS 2015 Testing Set: OM-Net + CGAp ranks first on the BraTS 2015 online leaderboard at the time of submission.

C. Performance Comparison on BraTS 2017 Validation Set

On BraTS 2017 validation, OM-Net improves on MC3 while using one model, and CGA, post-processing, and ensembling provide further gains. The resulting OM-Net variants achieve competitive leaderboard performance with lower memory requirements than approaches using large patches.

  • Model comparison: OM-Net outperforms MC3, particularly on enhancing-tumor Dice scores, while requiring only one-third of MC’s model parameters.The comparison evaluates models on the online BraTS 2017 validation set.
  • CGA ablation: CGA raises OM-Net Dice by 1.88% on tumor core and 2.09% on enhancing tumor.These gains are reported relative to OM-Net and exceed the performance of the SE variant.
  • Post-processing: OM-Net + CGAp further improves tumor-core Dice, supporting the effectiveness of the proposed post-processing method.Qualitative comparisons also report steadily improved segmentation quality across MC3, OM-Net, CGA, and CGAp.
  • Ensemble: The ten-model OM-Net + CGA⋆ ensemble consistently obtains higher Dice scores than OM-Net + CGA and outperforms all other methods on the online leaderboard.The ensemble trains each model on nine of ten folds and selects the best snapshot on the held-out fold.
  • Efficiency: OM-Net uses 32×32×16-voxel patches, making it more memory-efficient than a competing approach using 160 × 192 × 128-voxel patches and 32GB GPUs.The paper links the smaller patches to training and deployment on low-cost GPU devices.
  • Challenge result: The proposed techniques achieved joint third place among 64 teams in the BraTS 2018 Challenge.This result is reported alongside evaluations on BraTS 2015, 2017, and 2018.

VI. CONCLUSION

The paper concludes that OM-Net addresses class imbalance with one-pass coarse-to-fine segmentation while reducing model size and exploiting task correlations. Experiments across three datasets report advantages over MC and existing state-of-the-art methods.

  • Conclusion: OM-Net performs coarse-to-fine brain-tumor segmentation with one-pass computation, unlike the multi-model MC framework.The model is designed to handle class imbalance while simplifying the prediction system.
  • Conclusion: OM-Net reduces model size and system complexity while exploiting task correlations through shared parameters, training data, and prediction results.These mechanisms connect the separate segmentation tasks within one model.
  • Conclusion: CGA uses cross-task guidance to learn category-specific channel attention and significantly outperforms the SE block.The module is presented as a segmentation-oriented alternative to existing self-attention approaches.
  • Conclusion: A post-processing method refines OM-Net segmentation results to achieve better accuracy.The paper reports this method as an additional component of the proposed system.
Loading 1906.01796v2…