Augmented Normalization: Differentiating the Generalized Geometric Median

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.