Source-linked AI summary
3D MRI brain tumor segmentation using autoencoder regularization
Andriy Myronenko
TL;DR
Manual 3D brain-tumor delineation is time-consuming and can be inaccurate, motivating automated, reproducible segmentation for analysis and monitoring. The paper presents a multimodal 3D MRI encoder-decoder with a variational autoencoder regularization branch, and the approach won first place in the BraTS 2018 challenge.
Problem
Manual brain-tumor delineation requires anatomical expertise, is expensive and time-consuming, and can introduce human error, motivating automated reproducible segmentation for tumor analysis and monitoring.
Method
A volumetric encoder-decoder CNN segments three tumor subregions from multimodal 3D MRIs while a variational autoencoder branch reconstructs input images to regularize the shared network during training.
Results
The approach won 1st place in the BraTS 2018 challenge; an ensemble of 10 models improved validation performance by 1%.
Takeaways & Limitations
At inference, only the main segmentation encoder-decoder is used, without the additional reconstruction branch.
Takeaways & Limitations
GPU memory limits constrained the training configuration: using a larger batch size required smaller image crops and resulted in worse performance.
Abstract
from arXiv · showhide
Automated segmentation of brain tumors from 3D magnetic resonance images (MRIs) is necessary for the diagnosis, monitoring, and treatment planning of the disease. Manual delineation practices require anatomical knowledge, are expensive, time consuming and can be inaccurate due to human error. Here, we describe a semantic segmentation network for tumor subregion segmentation from 3D MRIs based on encoder-decoder architecture. Due to a limited training dataset size, a variational auto-encoder branch is added to reconstruct the input image itself in order to regularize the shared decoder and impose additional constraints on its layers. The current approach won 1st place in the BraTS 2018 challenge.
1 Introduction
Brain tumors, particularly gliomas, require multimodal MRI for diagnosis and treatment planning, while automated 3D segmentation can reduce physician workload and improve reproducibility. The proposed approach addresses volumetric tumor-subregion segmentation with an encoder-decoder network augmented by a VAE branch that regularizes shared representations.
- Clinical context: Multimodal 3D MRI supports brain-tumor diagnosis, monitoring, and surgery planning by emphasizing different tissue properties and tumor regions.The modalities include T1, contrast-enhanced T1, T2, and FLAIR.
- Motivation: Manual tumor delineation is time-consuming, costly, dependent on anatomical expertise, and susceptible to human error.
- Motivation: Automated 3D segmentation can save physicians time and provide accurate, reproducible outputs for tumor analysis and monitoring.
- Benchmark: The BraTS 2018 training set contained 285 cases from 19 institutions, with four aligned MRI modalities and annotations for nested whole-tumor, tumor-core, and enhancing-tumor regions.The dataset included 210 high-grade glioma and 75 low-grade glioma cases.
- Approach: The proposed encoder-decoder segments multimodal 3D MRIs while a VAE branch reconstructs the input during training to regularize the shared encoder.At inference, only the main segmentation encoder-decoder is used.
2 Related work
Prior BraTS systems used ensembles, cascaded segmentation, modified U-nets, dilated convolutions, and uncertainty-aware losses. This work instead uses a large crop with batch size one, directly predicts all three nested regions, and adds a training-only reconstruction branch.
- Prior approaches: BraTS 2017 high-performing systems included EMMA, an ensemble of independently trained architectures, and a cascaded approach using anisotropic convolutions.
- Prior approaches: BraTS 2018 competitors explored generic U-net modifications, additional institutional data, DenseNet blocks with dilated convolutions, uncertainty-aware loss, and multi-scale ensembles.
- This work: This work uses a 160x192x128 crop with batch size 1, directly predicts all three nested tumor subregions, and avoids additional training data.
- This work: The architecture adds a training-only VAE reconstruction branch to regularize the shared encoder, alongside direct sigmoid outputs for the three tumor subregions.
3 Methods
The method combines an encoder-decoder CNN for 3D tumor subregion segmentation with a VAE branch that reconstructs inputs to regularize learned features during training.
- VAE regularization: A VAE branch reconstructs the original image from a 256-dimensional latent representation, providing additional encoder guidance and regularization for limited training data.The latent representation contains 128 dimensions for the mean and 128 for the standard deviation.
- Network architecture: The encoder uses ResNet blocks with Group Normalization, progressively downsizing spatial dimensions while increasing feature size.Group Normalization is used because training uses a batch size of 1.
- Network architecture: The decoder upsamples features, adds encoder outputs from matching spatial levels, and produces three sigmoid tumor-subregion channels.The VAE decoder omits these inter-level skip connections and reconstructs the input image from a sampled latent representation.
- Loss function: The training objective combines soft Dice segmentation loss, weighted L2 image-reconstruction loss, and weighted KL divergence loss.The reconstruction and KL terms each receive an empirically selected weight of 0.1.
- Optimization: The model is optimized with Adam from an initial learning rate of 1e-4, a 300-epoch schedule, batch size 1, and random image ordering.L2 kernel regularization uses weight 1e-5, while spatial dropout uses rate 0.2 after the initial encoder convolution.
- Data preparation: Training uses normalized, randomly shifted and scaled input channels plus random axis mirror flips for augmentation.Normalization is based on nonzero voxels, and each of the three axes is mirrored with probability 0.5.
4 Results
The proposed model was evaluated on BraTS 2018 validation and testing datasets using segmentation metrics, test-time augmentation, and model ensembling. It achieved second-best validation Dice performance and won the testing challenge.
- Evaluation setup: The model was evaluated on 66 validation cases and 191 testing cases using Dice, sensitivity, specificity, and Hausdorff distance.The datasets had unknown glioma grade and segmentation labels, and results were submitted to the BraTS 2018 server.
- Validation results: The network received second-best Dice accuracy for all three segmentation labels: ET, WT, and TC.This result refers to the initial short-paper submission on July 13, 2018.
- Evaluation setup: The 10-model ensemble improved performance by 1%, whereas test-time augmentation produced only a marginal improvement.Test-time augmentation averaged predictions from eight mirror-flipped input volumes.
- Testing results: The single testing submission won first place in the BraTS 2018 challenge.The testing results were reported in Table 3 using mean Dice and Hausdorff measurements for ET, WT, and TC.
- Efficiency: Training took 9 minutes per epoch on one V100 GPU, while inference took 0.4 seconds for a single model on one V100 GPU.Training for 300 epochs took two days on one GPU; using eight interconnected V100 GPUs reduced training to six hours, a 7.8x speedup.
5 Discussion and Conclusion
The study concludes that the proposed 3D multimodal MRI segmentation network won BraTS 2018, while several architectural and post-processing alternatives did not improve performance.
- Conclusion: The proposed semantic segmentation network for multimodal 3D MRI brain-tumor segmentation won the BraTS 2018 challenge.
- Architecture experiments: Increasing network width consistently improved results, whereas increasing network depth further did not improve performance.
- Architecture experiments: Larger-batch BatchNorm training performed worse because GPU memory limits required smaller image crops.
- Post-processing and augmentation: More sophisticated data augmentation and CRF-based post-processing produced no additional overall improvement.CRF helped some images but worsened segmentation results for others.