Augmented Normalization: Differentiating the Generalized Geometric Median

Tyler King, Ser-Nam Lim

Keywords: Batch Normalization, Distributionally Robust Optimization, Subpopulation Shift, Out-of-Distribution Robustness

TL;DR: We introduce Augmented Normalization (AugNorm), a plug-and-play normalization approach for deep neural networks that achieves comparable in-distribution test accuracy to BatchNorm while improving worst-case test accuracy under subpopulation shift.

Abstract: We introduce Augmented Normalization (AugNorm), a novel feature transformation method that normalizes around a generalized geometric median (GGM). Unlike traditional normalization techniques that require fixed statistics such as the mean or median, AugNorm formulates the GGM as the minimizer of a convex argmin function, enabling a smooth interpolation between generic test statistics such as the mean, median, and range. This yields a noise-robust test statistic that avoids optimization difficulties associated with median-based methods. On CIFAR-10, AugNorm matches BatchNorm in-distribution while outperforming median normalization. On CelebA, a subpopulation shift dataset, we show AugNorm strengthens out-of-distribution robustness by improving worst-case test accuracy. We extend our method by introducing a differentiable variant of AugNorm, where the test statistic becomes a trainable parameter. Our results indicate AugNorm is a simple and effective drop-in replacement for BatchNorm that complements existing robust training schemes for settings with distribution shift.

Submission Number: 74