XTinyU-Net: Training-Free U-Net Scaling via Initialization-Time Sensitivity

TL;DR

This paper introduces XTinyU-Net, a training-free method to identify ultra-lightweight, dataset-specific U-Net configurations for medical image segmentation by analyzing initialization-time sensitivity, achieving high accuracy with minimal parameters.

Contribution

A novel Jacobian-based sensitivity metric enables automatic selection of the smallest stable U-Net configuration without training, reducing computational costs significantly.

Findings

XTinyU-Net matches nnU-Net accuracy with 400-1600x fewer parameters.

Outperforms other lightweight architectures with 5-72x fewer parameters.

Effective across six diverse medical datasets.

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.