Active learning boosts 3D medical image segmentation on single GPUs

Active-DiNTS: Active Differentiable Network Topology Search

Computer Vision and Pattern RecognitionArtificial IntelligenceMachine LearningNeural and Evolutionary Computing

Summary

Finding good computer designs for segmenting 3D medical images usually needs lots of labeled data and powerful computers. The authors introduce a method called Active-DiNTS that actively picks the most uncertain images to label, while simultaneously searching for the best network design. This process works efficiently on just one GPU and uses fewer labeled images. It improves accuracy notably on brain tumor MRI segmentation compared to earlier methods, and speeds up the search for network designs by over twenty times. This makes designing and training 3D image segmentation models more practical under real-world limits.

What this means in practice

  • For medical image engineers: Improve 3D segmentation models by jointly selecting data to label and architecture design under limited hardware and annotation resources.
  • For machine learning platform developers: Integrate active data selection with neural architecture search to speed up model development for complex 3D imaging tasks on single GPUs.

Authors

Gean Trindade Pereira, Thierry Urruty, Muriel Visani, André C. P. L. F. de Carvalho

Abstract

Neural Architecture Search (NAS) has proved to be a strong alternative to manual network design, but applying it to 3D medical image segmentation is limited by two well-known costs, large annotation budgets and multi-GPU clusters. Thus, this paper introduces Active-DiNTS (Active Differentiable Network Topology Search), an approach that embeds pool-based Active Learning (AL) into a bi-level differentiable topology search to perform architecture discovery and label curation jointly. At each query round, unlabeled MRI volumes are ranked by one of three uncertainty signals (Entropy, Variance, or Standard Deviation), and only the top-ranked volumes are sent to an oracle for annotation. The new labels feed two interlocked stages. Network weights are updated in an outer loop, while the macro/micro topology of a U-Net-style backbone is refined in an inner loop. Three AL regimes (weights-only, topology-only, joint) expose the speed-accuracy trade-off. Evaluations on the Medical Segmentation Decathlon (MSD) Task01 BrainTumour benchmark showed that Active-DiNTS surpasses DiNTS, C2FNAS, and nnU-Net in Dice, with gains of about 10 percentage points on Edema and 5 points on Non-Enhancing core, on a single GPU and using a fraction of the labeled volumes. The discovered architectures are denser and more FLOP-heavy than prior baselines, but remain competitive in trainable parameters and peak memory; the fastest search variant finishes in under 0.25 GPU-days, over 27x faster than the eight-GPU DiNTS search. Together, these results indicate that pairing differentiable NAS with active data acquisition is a practical recipe for accurate 3D segmentation under realistic constraints.