Towards Generalization Of Medical Imaging Ai Models: Sharpness-Aware Minimizers And Beyond

Learning representations to perform a task well on an unseen test dataset is the central problem of generalization. Numerous techniques to improve generalization ability of deep neural networks are continuously being developed. Each of these techniques work on different ingredients of a neural network training. In this paper, we investigate the interplay of three different disciplines: (i) data augmentation, concerning training data, (ii) transfer learning, concerning initialization of weights and (iii) sharpness aware minimization (SAM), concerning the optimization algorithm used for learning weights. We consider two representative problems in medical imaging: (a) a difficult task of cardiac view classification on ultrasound images and (b) COVID-19 detection from chest X-ray images, and perform extensive experimental analysis to methodically expose the improvement gained in generalization due to each of the aforementioned techniques individually, and also when used in combination. Our empirical studies suggest that sharpness aware minimization improves generalization by 5-10%, over and above the gain obtained by other methods. For the cardiac view classification problem, SAM generalizes well even when there is a distribution shift in test data. Further, we scrutinize the models under various settings to understand the geometry of loss landscape.