Efficient Search of Implantable Adaptive Cells for Medical Image Segmentation

TL;DR

This paper introduces IAC-LTH, a faster NAS method for implantable adaptive cells in U-Net, which maintains segmentation quality while significantly reducing search time across multiple medical imaging benchmarks.

Contribution

It proposes a stability-based pruning criterion that allows early stopping of NAS, making adaptive skip-module design more practical for medical image segmentation.

Findings

IAC-LTH reduces NAS cost by up to 16x while maintaining or improving segmentation performance.

Early-stabilizing operations can predict final architecture quality, enabling efficient search.

The method is effective across various datasets, architectures, and training settings.

Abstract

Purpose: Adaptive skip modules can improve medical image segmentation, but searching for them is computationally costly. Implantable Adaptive Cells (IACs) are compact NAS modules inserted into U-Net skip connections, reducing the search space compared with full-network NAS. However, the original IAC framework still requires a 200-epoch differentiable search for each backbone and dataset. Methods: We analyzed the temporal behavior of operations and edges within IAC cells during differentiable search on public medical image segmentation benchmarks. We found that operations selected in the final discrete cell typically emerge among the strongest candidates early in training, and their architecture parameters stabilize well before the final epoch. Based on this, we propose a Jensen--Shannon-divergence-based stability criterion that tracks per-edge operation-importance distributions and progressively prunes low-importance operations during search. The accelerated framework is called IAC-LTH. Results: Across four public benchmarks (ACDC, BraTS, KiTS, AMOS), several 2-D U-Net backbones, and a 2-D nnU-Net pipeline, IAC-LTH discovers IAC cells whose patient-level segmentation performance matches and sometimes slightly exceeds that of cells found by the original full-length search, while reducing wall-clock NAS cost by 3.7x to 16x across datasets and backbones. These results are consistent across architectures, benchmarks, and both non-augmented and augmented training settings, while preserving the gains of IAC-equipped U-Nets over strong attention-based and dense-skip baselines. Conclusion: Competitive IAC architectures can be identified from early-stabilizing operations without running the full search, making adaptive skip-module design more practical for medical image segmentation under realistic computational constraints.