GradMask: Reduce Overfitting by Regularizing Saliency

Becks Simpson, Francis Dutil, Yoshua Bengio, Joseph Paul Cohen

Keywords: saliency, prediction attribution, regularization, medical imaging, deep learning

TL;DR: Regularizing saliency maps to prevent overfitting from incorrect feature attribution

Abstract: With too few samples or too many model parameters, overfitting can inhibit the ability to generalise predictions to new data. Within medical imaging, this can occur when features are incorrectly assigned importance such as distinct hospital specific artifacts, leading to poor performance on a new dataset from a different institution without those features, which is undesirable. Most regularization methods do not explicitly penalize the incorrect association of these features to the target class and hence fail to address this issue. We propose a regularization method, GradMask, which penalizes saliency maps inferred from the classifier gradients when they are not consistent with the lesion segmentation. This prevents non-tumor related features to contribute to the classification of unhealthy samples. We demonstrate that this method can improve test accuracy between 1-3% compared to the baseline without GradMask, showing that it has an impact on reducing overfitting.

Code Of Conduct: I have read and accept the code of conduct.