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

2026-02-25

AI Summary

This work proposes MNAS-Unet, a novel neural architecture search (NAS) framework for medical image segmentation that addresses the high computational cost and inefficiency of traditional NAS methods. By integrating Monte Carlo Tree Search (MCTS) into the search process for the first time in this domain, MNAS-Unet dynamically explores efficient network architectures while respecting resource constraints. The method introduces lightweight DownSC and UpSC units to enable rapid model adaptation. Experimental results on the PROMISE12, Ultrasound Nerve, and CHAOS datasets demonstrate that MNAS-Unet achieves superior segmentation accuracy with significantly improved efficiency: it reduces the search budget by 54%, uses only 0.6 million parameters, and incurs lower GPU memory consumption compared to existing approaches.

Technology Category

• Search and Optimization: Learning to Search
• Machine Learning: Deep Neural Architectures and Foundation Models
• Computer Vision: Segmentation

Application Category

• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

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.

Problem

• Medical Image Segmentation
• Neural Architecture Search
• Monte Carlo Tree Search
• Efficient Architecture Search
• Resource-Constrained Optimization

Innovation

• Monte Carlo Tree Search
• Neural Architecture Search
• Medical Image Segmentation
• Lightweight Model
• Efficient Architecture Search