Neural architecture search for 3D biomedical image classification

Abstract

3D medical image classification is crucial for improving diagnostic accuracy and treatment planning, but it encounters challenges due to the complexity and variability of volumetric data. While 3D Convolutional Neural Networks offer potential solutions, designing effective architectures is complex and resource-intensive. Neural Architecture Search automates this process, optimizing network designs for specific tasks, thereby improving model performance. This study introduces a novel extension of the PBC-NAS method for 3D medical image classification, aiming to balance prediction accuracy and model complexity. We focus on optimizing neural network architectures using Neural Architecture Search for six different 3D datasets from MedMNIST3D, including OrganMNIST3D, NoduleMNIST3D, FractureMNIST3D, AdrenalMNIST3D, VesselMNIST3D, and SynapseMNIST3D, which are derived from real-world clinical imaging datasets. We have compared our method with state-of-the-art handcrafted networks, AutoML frameworks and recent NAS studies in terms of prediction performance and model complexity. The proposed NAS methods demonstrate superior performance compared to state-of-the-art handcrafted networks and AutoML frameworks. Our proposed model (Ours #3 $$^\dagger$$ ) achieves the highest average Area Under the Curve (AUC) of 0.915 and accuracy (ACC) of 0.847 (best result across three independent runs), outperforming all handcrafted networks and AutoML frameworks. Compared to other NAS-based methods, all proposed models achieve higher average AUC scores, and it is important to note that they do not rely on data augmentation, pre-processing, or feature selection, unlike the competing NAS methods which do use data augmentation during training. The study also highlights significant reductions in computational complexity, with FLOPs reduced by up to 45.51 times and parameters by up to 211 times compared to ResNet models. An ablation study reveals that while fine-tuning a model optimized for one dataset can achieve competitive results on other datasets, dataset-specific NAS is crucial for optimal performance. Despite this, the ablation results still outperform ResNets and AutoML frameworks in terms of average AUC and ACC. The study concludes that the proposed NAS approach effectively optimizes neural network architectures for complex 3D medical image classification tasks, achieving state-of-the-art performance without data augmentation.