Scalable Neural Architecture Search for 3D Medical Image Segmentation

Sungwoong Kim,Ildoo Kim,Sungbin Lim,Woonhyuk Baek,Chiheon Kim,Hyungjoo Cho,Boogeon Yoon,Taesup Kim

DOI: https://doi.org/10.1007/978-3-030-32248-9_25

2019-06-14

Abstract:In this paper, a neural architecture search (NAS) framework is proposed for 3D medical image segmentation, to automatically optimize a neural architecture from a large design space. Our NAS framework searches the structure of each layer including neural connectivities and operation types in both of the encoder and decoder. Since optimizing over a large discrete architecture space is difficult due to high-resolution 3D medical images, a novel stochastic sampling algorithm based on a continuous relaxation is also proposed for scalable gradient based optimization. On the 3D medical image segmentation tasks with a benchmark dataset, an automatically designed architecture by the proposed NAS framework outperforms the human-designed 3D U-Net, and moreover this optimized architecture is well suited to be transferred for different tasks.

Machine Learning,Computer Vision and Pattern Recognition,Image and Video Processing

What problem does this paper attempt to address?

The main problem that this paper attempts to solve is for 3D medical image segmentation tasks, automatically optimizing the neural network architecture to improve performance and achieve rapid transfer to related tasks. Specifically:

1. **Limitations of Manual Design**: When traditional deep neural networks such as 3D U - Net are used to process 3D medical image segmentation, their performance depends on manual trial - and - error to determine the network architecture, training hyper - parameters, and pre - /post - processing steps. This method is not only time - consuming but also has limitations in terms of performance improvement and rapid transfer.

2. **Application of Automated Machine Learning (AutoML)**: Although the general deep - learning field has achieved rapid development through techniques such as Neural Architecture Search (NAS) and Neural Optimizer Search, in medical imaging tasks, only a few studies have used simple hyper - parameter optimization methods.

3. **Challenges of Large - Scale Discrete - Space Optimization**: Due to the high resolution of 3D medical images, optimizing the large - scale discrete architecture space is very difficult. This requires not only extremely high memory usage but also long - running times.

To solve the above problems, the author proposes a new scalable neural architecture search framework, which is especially suitable for 3D medical image segmentation tasks. This framework can automatically optimize the neural network architecture from a large - scale design space and, by introducing techniques such as continuous relaxation and Gumbel - softmax sampling, achieves scalable gradient optimization, thereby overcoming the optimization problems brought about by high - resolution 3D images.

In summary, the core objective of this paper is to develop an efficient automated method for designing a neural network architecture more suitable for 3D medical image segmentation tasks while ensuring that this architecture can be well transferred between different tasks.