Scalable Neural Architecture Search for 3D Medical Image Segmentation

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

Keywords: AutoML, Neural Architecture Search, Medical Image Segmentation

TL;DR: Scalable Neural Architecture Search for 3D Medical Image Segmentation

Abstract: In this paper, a neural architecture search (NAS) framework is formulated for 3D medical image segmentation, to automatically optimize a neural architecture from a large design space. For this, a novel NAS framework is proposed to produce the structure of each layer including neural connectivities and operation types in both of the encoder and decoder of a target 3D U-Net. In the proposed NAS framework, having a sufficiently large search space is important in generating an improved network architecture, however optimizing over such a large space is difficult due to the extremely large memory usage and the long run-time originated from high-resolution 3D medical images. Therefore, a novel stochastic sampling algorithm based on the continuous relaxation on the discrete architecture parameters is also proposed for scalable joint optimization of both of the architecture parameters and the neural operation parameters. This makes it possible to maintain a large search space with small computational cost as well as to obtain an unbiased architecture by reducing the discrepancy between the training-time and test-time architectures. On the 3D medical image segmentation tasks with a benchmark dataset, an automatically designed 3D U-Net by the proposed NAS framework outperforms the previous human-designed 3D U-Net as well as the randomly designed 3D U-Net, and moreover this optimized architecture is more compact and also well suited to be transferred for similar but different tasks.

Code Of Conduct: I have read and accept the code of conduct.