Optimizing Neural Network Architecture for Medical Image Segmentation Using Monte Carlo Tree Search

TL;DR

A novel medical image segmentation framework, MNAS-Unet, which combines Monte Carlo Tree Search (MCTS) and Neural Architecture Search (NAS), which suggests that MNAS-Unet can improve search efficiency while maintaining competitive segmentation accuracy under practical resource constraints.

Abstract

This paper proposes a novel medical image segmentation framework, MNAS-Unet, which combines Monte Carlo Tree Search (MCTS) and Neural Architecture Search (NAS). MNAS-Unet dynamically explores promising network architectures through MCTS, significantly enhancing the efficiency and accuracy of architecture search. It also optimizes the DownSC and UpSC unit structures, enabling fast and precise model adjustments. Experimental results demonstrate that MNAS-Unet outperforms NAS-Unet and other state-of-the-art models in segmentation accuracy on several medical image datasets, including PROMISE12, Ultrasound Nerve, and CHAOS. Furthermore, compared with NAS-Unet, MNAS-Unet reduces the architecture search budget by 54% (early stopping at 139 epochs versus 300 epochs under the same search setting), while achieving a lightweight model with only 0.6M parameters and lower GPU memory consumption, which further improves its practical applicability. These results suggest that MNAS-Unet can improve search efficiency while maintaining competitive segmentation accuracy under practical resource constraints.

Figures (8)

• Figure 1: Neural architecture search.
• Figure 2: Monte Carlo Tree Search process (adapted from chaslot2008parallel).
• Figure 3: (a) Architecture of U-Net. (b) The U-shaped structure of the MNAS-Unet is depicted with rectangles, each representing a cell architecture subject to optimization. The specific operations of DownSC and UpSC are integrated into MNAS-Unet, corresponding to downward and upward movements, respectively. A blue arrow indicates the progression of the feature map (input image). Additionally, the light purple arrow indicates a transformation operation within UpSC that is also subject to automatic search. Adapted from weng2019unet.
• Figure 4: A depiction of the cellular architecture is provided. A red arrow indicates a reduction process, typically represented by max pooling. A green arrow marks a standard operation, which might be an identity or a convolution that maintains the dimensions of the feature map. Moreover, a black arrow highlights the concatenation operation. Adapted from weng2019unet.
• Figure 5: Overview of MNAS-Unet search procedures.