Auto-nnU-Net: Towards Automated Medical Image Segmentation

2025-05-22

AI Summary

Existing AutoML approaches for medical image segmentation (MIS) lack sufficient automation under diverse task requirements and resource-constrained settings, particularly due to limited joint optimization of hyperparameters, neural architecture, and hierarchical design. Method: We propose the first fully automated nnU-Net variant enabling simultaneous hyperparameter optimization (HPO), differentiable neural architecture search (NAS), and hierarchical NAS (HNAS). We introduce PriorBand—a resource-aware regularization strategy—that dynamically balances segmentation accuracy and training cost via Bayesian optimization, differentiable search, and a novel hierarchical search space. Results: Evaluated on 10 datasets from the Medical Segmentation Decathlon, our method significantly improves Dice scores on six datasets (average +2.1%), matches performance on four, and maintains controllable training overhead—demonstrating practical deployability. This work establishes the first end-to-end, full-stack automated modeling framework for MIS, eliminating manual tuning and fixed-architecture limitations.

Technology Category

• Machine Learning: Auto ML and Hyperparameter Tuning
• Computer Vision: Segmentation
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Medical Image Segmentation (MIS) includes diverse tasks, from bone to organ segmentation, each with its own challenges in finding the best segmentation model. The state-of-the-art AutoML-related MIS-framework nnU-Net automates many aspects of model configuration but remains constrained by fixed hyperparameters and heuristic design choices. As a full-AutoML framework for MIS, we propose Auto-nnU-Net, a novel nnU-Net variant enabling hyperparameter optimization (HPO), neural architecture search (NAS), and hierarchical NAS (HNAS). Additionally, we propose Regularized PriorBand to balance model accuracy with the computational resources required for training, addressing the resource constraints often faced in real-world medical settings that limit the feasibility of extensive training procedures. We evaluate our approach across diverse MIS datasets from the well-established Medical Segmentation Decathlon, analyzing the impact of AutoML techniques on segmentation performance, computational efficiency, and model design choices. The results demonstrate that our AutoML approach substantially improves the segmentation performance of nnU-Net on 6 out of 10 datasets and is on par on the other datasets while maintaining practical resource requirements. Our code is available at https://github.com/LUH-AI/AutonnUNet.

Problem

• Automating hyperparameter and architecture search for medical image segmentation
• Balancing model accuracy with computational resource constraints
• Improving segmentation performance across diverse medical datasets

Innovation

• Enables hyperparameter optimization and neural architecture search
• Introduces Regularized PriorBand for resource-efficient training
• Improves segmentation performance while maintaining practical resources