A Lightweight Neural Architecture Search Model for Medical Image Classification
Lunchen Xie, Eugenio Lomurno, Matteo Gambella, Danilo Ardagna, Manuel Roveri, Matteo Matteucci, Qingjiang Shi
TL;DR
This paper tackles the high computational cost of neural architecture search for medical image classification. It introduces ZO-DARTS+, a differentiable NAS method that uses sparsemax and a sparsity-promoting annealing schedule, along with a zeroth-order gradient estimator to avoid expensive Hessian calculations. Across five MedMNIST datasets, ZO-DARTS+ achieves accuracy comparable to state-of-the-art NAS methods while reducing search time by up to three times and often completing search by around the 40th epoch. The study demonstrates that sparse, interpretable operation selections can tailor architectures to dataset characteristics, suggesting practical impact for rapid, domain-specific model design in medical imaging.
Abstract
Accurate classification of medical images is essential for modern diagnostics. Deep learning advancements led clinicians to increasingly use sophisticated models to make faster and more accurate decisions, sometimes replacing human judgment. However, model development is costly and repetitive. Neural Architecture Search (NAS) provides solutions by automating the design of deep learning architectures. This paper presents ZO-DARTS+, a differentiable NAS algorithm that improves search efficiency through a novel method of generating sparse probabilities by bi-level optimization. Experiments on five public medical datasets show that ZO-DARTS+ matches the accuracy of state-of-the-art solutions while reducing search times by up to three times.