Improving medical image segmentation in pre-trained U-Nets using Shapley-guided pruning of adaptive skip-connection modules.
Authors
Affiliations (2)
Affiliations (2)
- Institute of Computer Science and Mathematics, Maria Curie Sklodowska University, Akademicka 9 Street, 20-033, Lublin, Poland.
- Institute of Computer Science and Mathematics, Maria Curie Sklodowska University, Akademicka 9 Street, 20-033, Lublin, Poland. [email protected].
Abstract
Neural architecture search (NAS) can improve medical image segmentation, but its practical use is limited by computational cost and instability of the discovered architectures, particularly when applied to pre-trained models and limited data. We present Shapley-guided pruning as a practical validation-guided extension of retrofit NAS for pre-trained U-Nets. Rather than defining a new NAS family, the method keeps the IAC search space and PC-DARTS-style supernet optimization, while adding iterative pruning driven by Shapley value estimates computed on held-out validation data. By progressively removing low-impact components while preserving learned architecture parameters, the approach reduces search space complexity and improves the reliability of the final discrete architecture. We evaluate the method on four public benchmarks (ACDC, BraTS, KiTS, and AMOS) in a controlled 2D slice-based, [Formula: see text], single-GPU setting. Within this protocol, the proposed approach improves or matches strong baselines in most comparisons, accelerates search by up to four times, and yields more stable operation choices across runs. The findings support Shapley-guided pruning as a practical retrofit-NAS mechanism under resource constraints, without implying direct clinical competitiveness with high-resolution or fully 3D segmentation pipelines.