Improving Normalization with the James-Stein Estimator
Seyedalireza Khoshsirat Chandra Kambhamettu
Stein's paradox holds considerable sway in high-dimensional statistics,
highlighting that the sample mean, traditionally considered the de facto
estimator, might not be the most efficacious in higher dimensions. To address
this, the James-Stein estimator proposes an enhancement by steering the sample
means toward a more centralized mean vector. In this paper, first, we establish
that normalization layers in deep learning use inadmissible estimators for mean
and variance. Next, we introduce a novel method to employ the James-Stein
estimator to improve the estimation of mean and variance within normalization
layers. We evaluate our method on different computer vision tasks: image
classification, semantic segmentation, and 3D object classification. Through
these evaluations, it is evident that our improved normalization layers
consistently yield superior accuracy across all tasks without extra
computational burden. Moreover, recognizing that a plethora of shrinkage
estimators surpass the traditional estimator in performance, we study two other
prominent shrinkage estimators: Ridge and LASSO. Additionally, we provide
visual representations to intuitively demonstrate the impact of shrinkage on
the estimated layer statistics. Finally, we study the effect of regularization
and batch size on our modified batch normalization. The studies show that our
method is less sensitive to batch size and regularization, improving accuracy
under various setups.