AC-UNet: Adaptive Connection UNet for White Matter Tract Segmentation Through Neural Architecture Search

Yanming Zhu, Ari Tchetchenian, Xuefei Yin, Alan Wee-Chung Liew, Yang Song, Erik Meijering

Abstract: White matter tract segmentation in diffusion magnetic resonance images is crucial for brain health analysis. A prevailing method for this task is deep learning using U-shaped networks. Several variants of UNet have been proposed to improve its skip connections, resulting in segmentation improvements. In this paper, we propose a novel neural architecture search-based solution for skip connection optimization, named AC-UNet (Adaptive Connection UNet). It can automatically identify the optimal skip connections of a U-shaped network to construct the best architecture for the task. Moreover, we propose to search non-repeatable operations for each layer, which further extends the exploration of feature aggregation for better segmentation. Experimental results on the largest public tractograms dataset demonstrate the superiority of AC-UNet over mainstream UNet-based architectures.