🤖 AI Summary
This work addresses the challenge of efficiently searching for U-Net architectures that balance performance and compactness under resource constraints. The authors propose a training-free method for automatic U-Net lightweighting, which leverages sensitivity analysis during initialization to directly identify optimal ultra-lightweight architectures without requiring labeled data. For the first time, this approach precisely locates the stability boundary for channel-width compression without any training, substantially reducing model selection overhead. Its core components include a Jacobian-based sensitivity metric, total variation analysis, and unlabeled-image-driven initialization evaluation, all integrated into the nnU-Net framework. Evaluated across six medical imaging datasets, the resulting XTinyU-Net achieves accuracy comparable to nnU-Net with 400–1600× fewer parameters and outperforms existing lightweight models with 5–72× fewer parameters.
📝 Abstract
While U-Net architectures remain the gold standard for medical image segmentation, their deployment in resource-constrained environments demands aggressive model compression. However, finding an optimally efficient configuration is computationally prohibitive, typically requiring exhaustive train-and-evaluate cycles to find the smallest model that maintains peak performance. In this paper, we introduce a training-free selection framework to automatically identify ultralightweight, dataset-specific U-Net configurations directly at initialization. We observe that systematically scaling down U-Net channel width induces a sharp transition from a stable performance plateau to representational capacity collapse. To pinpoint this boundary without training, we propose a Jacobian-based sensitivity metric that scores discrete, width-capped U-Net variants using a small set of unlabeled images. By analyzing the total variation of this sensitivity curve, we isolate the smallest stable configuration, which we denote as XTinyU-Net. Evaluated across six diverse medical datasets within the nnU-Net framework, XTinyU-Net achieves segmentation accuracy comparable to the heavy nnU-Net baseline with 400x-1600x fewer parameters, and outperforms contemporary lightweight architectures while utilizing 5x-72x fewer parameters. Code is publicly accessible on https://github.com/alvinkimbowa/nntinyunet.git.