EMONAS-Net: Efficient multiobjective neural architecture search using surrogate-assisted evolutionary algorithm for 3D medical image segmentation

Highlights

• An efficient multiobjective neural architecture search framework is presented.

• The search space considers the micro and macro structure of the architecture.

• The proposed surrogate-assisted evolutionary algorithm improves convergence.

• The framework was tested on three 3D medical image segmentation challenges.

• The proposed framework efficiently identifies accurate and smaller architectures.

Abstract

Deep learning plays a critical role in medical image segmentation. Nevertheless, manually designing a neural network for a specific segmentation problem is a very difficult and time-consuming task due to the massive hyperparameter search space, long training time and large volumetric data. Therefore, most designed networks are highly complex, task specific and over-parametrized. Recently, multiobjective neural architecture search (NAS) methods have been proposed to automate the design of accurate and efficient segmentation architectures. However, they only search for either the micro- or macro-structure of the architecture, do not use the information produced during the optimization process to increase the efficiency of the search, or do not consider the volumetric nature of medical images. In this work, we present EMONAS-Net, an Efficient MultiObjective NAS framework for 3D medical image segmentation that optimizes both the segmentation accuracy and size of the network. EMONAS-Net has two key components, a novel search space that considers the configuration of the micro- and macro-structure of the architecture and a Surrogate-assisted Multiobjective Evolutionary based Algorithm (SaMEA algorithm) that efficiently searches for the best hyperparameter values. The SaMEA algorithm uses the information collected during the initial generations of the evolutionary process to identify the most promising subproblems and select the best performing hyperparameter values during mutation to improve the convergence speed. Furthermore, a Random Forest surrogate model is incorporated to accelerate the fitness evaluation of the candidate architectures. EMONAS-Net is tested on the tasks of prostate segmentation from the MICCAI PROMISE12 challenge, hippocampus segmentation from the Medical Segmentation Decathlon challenge, and cardiac segmentation from the MICCAI ACDC challenge. In all the benchmarks, the proposed framework finds architectures that perform better or comparable with competing state-of-the-art NAS methods while being considerably smaller and reducing the architecture search time by more than 50%.

Introduction

Convolutional Neural Networks (CNNs) have shown to be very successful in solving various computer vision and medical imaging problems such as image classification, detection and segmentation. However, a crucial factor in their success is designing the appropriate architecture for the particular application and dataset. Most of the currently employed architectures have been manually designed to solve the specific task, resulting in highly complex and over-parametrized networks. Moreover, manually designing an architecture is a very difficult and time-consuming process due to the massive number of hyperparameters and high computational cost of training a network.

Recently, methods that automate the search of the neural architecture have been developed known as Neural Architecture Search (NAS). NAS can be modeled as an optimization problem, where the aim is to find the architectural parameters that maximize the CNNs performance on the test data. Different optimization approaches have been presented such as evolutionary algorithms [[1], [2], [3]], reinforcement learning [[4], [5], [6]], and Bayesian optimization [7,8]. Unfortunately, NAS algorithms are usually time-consuming and computationally expensive, requiring even thousands of GPU days to obtain state-of-the-art architectures [1,2,4]. To overcome this challenge, recent works have proposed embedding the hyperparameter search space from a discrete to a continuous space [9,10], reusing learned weights [11,12], using a one-shot architecture search [13], or applying a surrogate function to estimate an architecture's validation loss [14,15]. Nevertheless, these techniques were developed for NAS methods that only optimize the model's accuracy and do not consider additional objective functions or constraints that might be imposed by the particular application.

Despite the rapid development and success of NAS in natural image processing tasks, only a few approaches have been presented for 3D medical image segmentation. The segmentation of medical images plays an important role in disease diagnosis, treatment planning and monitoring of disease progression. Differently from natural images, medical images can be multi-modal and usually have 3 or 4 dimensions that need to be simultaneously processed. Furthermore, they have a heterogeneous appearance caused by the distinct acquisition protocols, imagining equipment and large variability in terms of the location, size and shape of the anatomical structure of interest.

NAS approaches for medical image segmentation can be divided into two categories: algorithms that search for 2D architectures and produce the final segmentation by processing each slice independently [[16], [17], [18]] and algorithms that search for 3D segmentation architectures [[19], [20], [21], [22]]. These works focused on either optimizing the macro-structure of the network (best number of cells and their connection while keeping the configuration of the cell fixed) or the micro-structure of the network (best configuration for the cell while setting the number of cells and their connections arbitrarily). Ideally, both the micro- and macro-structures of the architecture should be optimized jointly to avoid having to resort to manual engineering to define the missing hyperparameter values. Very recently, Yu et al. [23] proposed a coarse-to-fine NAS that searches for the micro- and macro-structure of the network through a two-stage optimization problem. Since two optimization problems need to be solved in sequence to identify the micro- and macro-structure of the network, their approach is very time consuming. Furthermore, none of the aforementioned works have explicitly addressed the design of accurate as well as efficient architectures. Given the large volumetric data being processed, minimizing the complexity or size of a network is necessary to reduce its high computational cost and enable its application on real-world settings. Therefore, multiobjective NAS methods that construct accurate and efficient architectures while simultaneously optimizing the micro- and macro-structure of the architecture are needed for medical image applications.

In our previous work, we presented two multiobjective NAS methods for medical image segmentation. In [24], a multiobjective evolutionary based algorithm (MEA) is presented to adapt a semi-fixed 2D FCN architecture by maximizing the segmentation accuracy and reducing the network size. The segmentation is performed in a slice-wise manner and does not exploit the volumetric information available in the images. To incorporate volumetric information, an ensemble of 2D and 3D FCNs is introduced in [25,26] that is automatically designed using an improved multiobjective evolutionary based algorithm. However, their construction requires the search for both the 2D and 3D FCN architectures, which can be time-consuming for high-resolution images or large datasets. Moreover, neither of the proposed optimization approaches takes advantage of the information produced during the evolutionary search process. Considering that optimizing multiple objectives during the architectural search requires more iterations to approximate the Pareto Front, it is important to investigate algorithms that utilize all the information available to improve the convergence speed and reduce the search time.

In this paper, we propose an Efficient MultiObjective NAS framework for 3D medical image segmentation, called EMONAS-Net. EMONAS-Net has two key components: a novel search space that considers the configuration of the micro- and macro-structure of the segmentation architecture, and a Surrogate-assisted Multiobjective Evolutionary based Algorithm (SaMEA algorithm) that improves the efficiency of the architectural search. Regarding the search space, the micro-structure is represented by a directed acyclical graph (DAG) and the SaMEA algorithm determines the connections between the nodes and the most appropriate convolutional operation for each node. For the design of the macro-structure, the search space includes the hyperparameters that define the depth (number of decoder and encoder cells) and width (number of filters) of the architecture. The second component of our framework is the SaMEA algorithm, which searches for architectures that minimize the segmentation error and number of parameters of the network. We propose two strategies to improve the convergence and search time of the algorithm. First, information produced during the initial stages of the evolutionary search is used to increase the probability of selecting the best performing hyperparameter values and most promising subproblems. Secondly, a Random Forest model is applied as an inexpensive surrogate function to estimate a candidate's architecture performance and decrease the training time. The EMONAS-Net framework is tested on three medical image segmentations tasks from publically available datasets. These include the segmentation of the prostate in MR images from the MICCAI PROMISE12 challenge [27], the segmentation of the posterior and anterior parts of the hippocampus on MRI from the Medical Segmentation Decathlon challenge [28], and the segmentation of the left ventricle cavity, left ventricle myocardium and right ventricle cavity of cardiac MRI from the MICCAI ACDC challenge [29]. In all the datasets, the EMONAS-Net finds networks that perform better or similar to other NAS methods, while being significantly smaller and requiring considerably less computational time for the architecture search.

The contributions of this work are as follows:

• We propose a novel search space that simultaneously optimizes the micro- and macro-structure of the architecture providing three important advantages. First, the need for manual intervention is reduced as no pre-fixed templates for the cell or the depth/width of the macro architecture need to be defined beforehand. Second, the size of the architecture can be effectively minimized because all the hyperparameters that define the depth and width of the architecture in the micro and macro level are considered during the search. Finally, the total time to obtain the final architecture is reduced because no additional experiments or a second hyperparameter optimization problem is needed to determine the best configuration for the architecture.

• We propose a Surrogate-assisted Multiobjective Evolutionary based Algorithm (SaMEA) that takes advantage of the information generated during the initial generations of the evolutionary search to improve the exploration of the search space and convergence.

• We present an inexpensive Random Forest surrogate function that successfully ranks the performance of the candidate architectures during the evolutionary search and significantly reduces the time to automatically design deep neural networks. Furthermore, compared with other commonly used performance predictors such as neural networks, the proposed Random Forest does not need an extensive amount of training data and the number of hyperparameters to be set is very small and generalizable to various datasets.