Learning Faster without Deeper Networks: A*-Inspired Batch Selection for Efficient CNN Training
Common practice when training Convolutional Neural Networks (CNNs) is to utilize randomly shuffled mini-batches. During this stage of the training process, two limitations are faced. Firstly the slowing down of convergence, and secondly the limited learning signal contribution due to large quantities of samples that are considered easy while training CNNs.In our work, we aim to address these inefficiencies by proposing A*-Inspired Batch Selection (A*-BS). This is a lightweight and model-agnostic training strategy that formulates mini-batch scheduling as a heuristic search problem. In this type of batch selection process each batch is treated as a node in a search space and ranked using an A*-like score by combining a loss-based difficulty measure with a reuse penalty. The two main benefits of this method are: informative gradient updates encouragement and batch diversity selection throughout the training process. The proposed method does not modify network architectures or optimization algorithms and can be seamlessly integrated into existing training pipelines. We evaluate the approach on the twelve 2D classification tasks of the MedMNIST-v2 benchmark, using a deliberately simple architecture of approximately 2.25 × 105 parameters and comparing against the ResNet-18 and ResNet-50 baselines reported by the benchmark.On half of these tasks, the lightweight model combined with A*-BS reaches higher accuracy and AUC than both ResNet baselines, with relative gains of up to 15%. Moreover, an ablation under identical CNN architecture and hyper-parameters shows that A*-BS outperforms random batch shuffling on all twelve MedMNIST tasks. Wall-clock measurements further show that the lightweight CNN with A*-BS trains substantially faster than ResNet-18 and ResNet-50 on identical hardware.Indicative in these results is that intelligent batch ordering can partially compensate for reduced architectural complexity, thus offering a computationally efficient alternative to deeper models. Furthermore, the reliability of the proposed strategy is reinforced by its strong performance with a simple CNN, even when compared to deeper and more sophisticated architectures.