Source-linked AI summary
On the Compactness, Efficiency, and Representation of 3D Convolutional Networks: Brain Parcellation as a Pretext Task
Wenqi Li, Guotai Wang, Lucas Fidon, Sebastien Ourselin, M. Jorge Cardoso, Tom Vercauteren
TL;DR
The paper addresses the difficulty of designing efficient 3D CNNs for volumetric medical images. It introduces a high-resolution compact network using dilated convolutions and residual connections, validates it on 155-structure brain parcellation, and demonstrates dropout-based voxel-level uncertainty estimation. The network compares favourably with competing volumetric segmentation methods while using substantially fewer parameters.
Problem
Designing efficient 3D CNNs for volumetric medical images is challenging because voxel counts scale cubically with region size, increasing computational burden.
Method
The paper develops a high-resolution 3D segmentation network using dilated convolutions, residual connections, and dropout sampling for voxel-level uncertainty estimation.
Results
The network achieves competitive brain-parcellation performance on 155 neuroanatomical structures while using twenty times fewer parameters than state-of-the-art architectures.
Takeaways & Limitations
Brain parcellation serves as a pretext task, and the trained network potentially provides a starting point for transfer learning to other volumetric segmentation tasks.
Takeaways & Limitations
The network’s generalisation across brain MR scans from different scanning protocols and data centres remains to be extensively tested, and its uncertainty scores are not probabilities.
Abstract
from arXiv · showhide
Deep convolutional neural networks are powerful tools for learning visual representations from images. However, designing efficient deep architectures to analyse volumetric medical images remains challenging. This work investigates efficient and flexible elements of modern convolutional networks such as dilated convolution and residual connection. With these essential building blocks, we propose a high-resolution, compact convolutional network for volumetric image segmentation. To illustrate its efficiency of learning 3D representation from large-scale image data, the proposed network is validated with the challenging task of parcellating 155 neuroanatomical structures from brain MR images. Our experiments show that the proposed network architecture compares favourably with state-of-the-art volumetric segmentation networks while being an order of magnitude more compact. We consider the brain parcellation task as a pretext task for volumetric image segmentation; our trained network potentially provides a good starting point for transfer learning. Additionally, we show the feasibility of voxel-level uncertainty estimation using a sampling approximation through dropout.
1 Introduction
Volumetric medical images offer richer representations than 2D images but make 3D CNN design computationally challenging. The paper proposes a compact, high-resolution architecture that supports brain parcellation and voxel-level uncertainty estimation.
- Motivation: 3D medical image analysis is challenging because voxel counts scale cubically with region size, increasing pattern-learning complexity and computational burden.Volumetric representations are more informative than 2D ones, but their computational demands are substantially higher.
- Approach: The proposed network maintains high spatial resolution throughout its layers while using dilated convolutions and residual connections for volumetric segmentation.It is validated by parcellating 155 neuroanatomical structures from T1-weighted MR images.
- Results: With twenty times fewer parameters, the proposed architecture achieves competitive segmentation performance against state-of-the-art volumetric networks.The comparison concerns the 155-structure brain parcellation task.
- Representation learning: Brain parcellation is treated as a pretext task for learning volumetric representations that could support transfer learning to other segmentation tasks.The trained network is presented as a potential starting point for other volumetric image-recognition problems.
- Uncertainty estimation: The study demonstrates voxel-level uncertainty estimation by sampling the proposed dropout network at test time.The motivation is to indicate segmentation confidence and reliability.
2 On the elements of 3D convolutional networks
The paper combines small kernels, dilated convolutions, and residual connections to build compact, high-resolution 3D segmentation networks. These elements preserve spatial resolution, expand receptive-field diversity, and support voxel-level uncertainty estimation.
- Convolutions: 3 × 3 × 3 kernels reduce parameters while preserving receptive fields achievable with stacked convolutions.Stacking two 3 × 3 × 3 layers uses approximately 57% fewer parameters than a 5 × 5 × 5 kernel.
- Dilated convolutions: Dilated convolutions enlarge receptive fields while preserving spatial resolution and using the same parameter count as standard 3 × 3 × 3 convolutions.A dilation factor r provides a (2r + 1)^3-voxel receptive field; r = 1 recovers the standard convolution.
- Residual connections: Residual connections merge block inputs with nonlinear outputs, enabling direct information propagation through forward and back-propagation.The residual formulation is x_p+1 = x_p + F(x_p, w_p).
- Effective receptive field: 29 residual-network paths produce receptive fields ranging from 3 × 3 × 3 to 87 × 87 × 87 voxels in the proposed nine-block network.The receptive-field distribution follows a binomial distribution and contrasts with the smaller number of paths in DeepMedic and 3D U-net.
- Effective receptive field: Border distortions remain small despite the network’s large maximum receptive field, and test-time zero padding with output cropping is used to manage borders.The authors suggest residual connections reduce the effective receptive field relative to the maximum receptive field.
- Loss function: Dice loss is used instead of cross-entropy to address severe class imbalance in medical image segmentation.Cross-entropy can bias estimation toward majority classes because it treats training voxels equally.
- Uncertainty estimation: Dropout sampling approximates predictive distributions, allowing voxel-level segmentation uncertainty to be estimated from sample variance and voting disagreement.The extended network applies dropout at test time and aggregates multiple stochastic predictions.
3 The network architecture and its implementation
The proposed network is a 20-layer, end-to-end volumetric segmentation architecture that preserves spatial resolution while combining dilated convolutions and residual connections. It processes 96 × 96 × 96-voxel inputs, with preprocessing based on standardisation, augmentation, and Adam optimisation.
- Architecture: 20 convolutional layers process volumetric inputs, beginning with 3 × 3 × 3 convolutions for low-level features and later using dilated kernels for deeper representations.The subsequent kernels use dilation factors of two or four to encode mid- and high-level features.
- Architecture: Dilated convolutions gradually increase their dilation factor with depth to incorporate features at multiple scales while maintaining input spatial resolution.The architecture maps an image volume directly to voxel-level dense segmentation.
- Architecture: Residual connections group every two convolutional layers and directly merge block inputs with outputs through identity mappings.Each residual block uses pre-activation ordering with ReLU, batch normalisation, and convolutional layers.
- Implementation: The network is trained end-to-end on 96 × 96 × 96-voxel images, producing per-voxel classification scores through a final softmax layer.The architecture is designed for volumetric image segmentation.
- Implementation: Training preprocessing combines input standardisation with image- and subvolume-level augmentation, including rotations and spatial rescaling by factors from 0.9 to 1.1.Image-level augmentation also randomises the foreground threshold used during normalisation.
- Implementation: The networks are trained on two Nvidia K80 GPUs using averaged gradients from one input volume per GPU and the Adam optimisation method.Parameter initialisation follows He et al., while batch-normalisation scaling and shifting parameters start at 1 and 0.
4 Experiments and results
Experiments evaluate the compact network on highly granular brain parcellation, compare architectural and training variants with volumetric baselines, and assess border effects and dropout-based uncertainty estimation.
- Experimental setup: 543 healthy-control T1-weighted MR images were parcellated into 155 brain structures and 5 nonbrain tissues from the ADNI dataset.Volumes averaged about 182 × 244 × 246 voxels, with approximately 1.18mm × 1.05mm × 1.05mm voxel size.
- Experimental setup: The evaluation compared HC-default and variants without residual connections, with Dice or cross-entropy loss, and with dropout, alongside 3D U-net, V-net, and Deepmedic.The competing architectures used a 160-way softmax classifier in their final layer.
- Architecture and loss comparisons: Dice loss largely improved segmentation performance, addressing the severely unbalanced segmentation problem by implicitly re-weighting voxels.All networks were trained with Dice loss for a fair comparison after cross-entropy produced poor results.
- Architecture and loss comparisons: 2% DCS improvement followed adoption of dropout, while residual connections also improved mean DCS.The dropout variant achieved better segmentation results for all key structures.
- Architecture and loss comparisons: HC-default and HC-dropout outperformed competing methods in mean DCS despite using relatively few parameters.The lower DCS of V-net suggests batch normalisation was important for brain-parcellation training.
- Border effects: 17 voxels marked the approximate distorted border width in each dimension, with border effects not severely decreasing performance.The practical procedure padded volumes with 16 zeros per dimension and removed the same output border.
- Uncertainty estimation: Ten Monte Carlo dropout samples were sufficient for good segmentation, while uncertainty was relatively higher near boundaries between structures.Increasing the sample count had relatively small effects on DCS, and uncertainty thresholding enabled voxel-wise accuracy analysis.
- Uncertainty estimation: A typical 192×256×256-voxel volume required about 60 seconds for one prediction and slightly more than 10 minutes for ten Monte Carlo samples.The authors note that reusing pre-dropout features could reduce sampling runtime to near single-prediction cost.
5 Conclusion
The paper presents a compact high-resolution 3D network for brain-MR parcellation, with voxel-level uncertainty estimation and potential use for transfer learning. Its generalisation across scanning protocols and data centres remains future work, and its uncertainty scores are not yet calibrated probabilities.
- The proposed network uses dilated convolutions and residual connections to incorporate large volumetric context in high-resolution 3D segmentation.
- Monte Carlo dropout sampling demonstrates the feasibility of generating voxel-level uncertainty estimates for the brain parcellation network.
- Brain parcellation is treated as a pretext task, and the trained network potentially provides a starting point for transfer learning to other segmentation tasks.
- Generalisation across brain MR scans from varied scanning protocols and data centres remains to be tested extensively.
- The reported uncertainty estimates are not probabilities, and future work will investigate calibrating them into reliable probability estimates.