Noise Injection: Improving Out-of-Distribution Generalization for Limited Size Datasets

TL;DR

The paper tackles poor out-of-distribution generalization in COVID-19 detection from chest X-rays due to shortcut learning on limited data. It investigates training-time noise injection across four noise types as a data augmentation strategy and analyzes how training-source diversity affects generalization to unseen sources. Across multiple metrics, the approach reduces the ID–OOD gap from about $0.10-0.20$ to $0.01-0.06$ on average over 10 seeds, demonstrating robustness improvements under data scarcity. The findings highlight the practical value of noise-based augmentation, while emphasizing that the choice of training data sources critically shapes generalization capabilities and should be carefully managed to learn generalizable biomarkers.

Abstract

Deep learned (DL) models for image recognition have been shown to fail to generalize to data from different devices, populations, etc. COVID-19 detection from Chest X-rays (CXRs), in particular, has been shown to fail to generalize to out-of-distribution (OOD) data from new clinical sources not covered in the training set. This occurs because models learn to exploit shortcuts - source-specific artifacts that do not translate to new distributions - rather than reasonable biomarkers to maximize performance on in-distribution (ID) data. Rendering the models more robust to distribution shifts, our study investigates the use of fundamental noise injection techniques (Gaussian, Speckle, Poisson, and Salt and Pepper) during training. Our empirical results demonstrate that this technique can significantly reduce the performance gap between ID and OOD evaluation from 0.10-0.20 to 0.01-0.06, based on results averaged over ten random seeds across key metrics such as AUC, F1, accuracy, recall and specificity. Our source code is publicly available at https://github.com/Duongmai127/Noisy-ood