Admissibility of Stein Shrinkage for Batch Normalization in the Presence of Adversarial Attacks
This work addresses robustness in deep learning for image processing tasks, but it is incremental as it applies an existing statistical method to a known bottleneck in adversarial settings.
The paper tackles the vulnerability of batch normalization to adversarial attacks by proving that Stein shrinkage estimators for mean and variance dominate sample estimators under sub-Gaussian attack models, leading to state-of-the-art performance in image classification and segmentation tasks with and without attacks.
Batch normalization (BN) is a ubiquitous operation in deep neural networks used primarily to achieve stability and regularization during network training. BN involves feature map centering and scaling using sample means and variances, respectively. Since these statistics are being estimated across the feature maps within a batch, this problem is ideally suited for the application of Stein's shrinkage estimation, which leads to a better, in the mean-squared-error sense, estimate of the mean and variance of the batch. In this paper, we prove that the Stein shrinkage estimator for the mean and variance dominates over the sample mean and variance estimators in the presence of adversarial attacks when modeling these attacks using sub-Gaussian distributions. This facilitates and justifies the application of Stein shrinkage to estimate the mean and variance parameters in BN and use it in image classification (segmentation) tasks with and without adversarial attacks. We present SOTA performance results using this Stein corrected batch norm in a standard ResNet architecture applied to the task of image classification using CIFAR-10 data, 3D CNN on PPMI (neuroimaging) data and image segmentation using HRNet on Cityscape data with and without adversarial attacks.