  We present Sandwich Batch Normalization (SaBN), a frustratingly easy
improvement of Batch Normalization (BN) with only a few lines of code changes.
SaBN is motivated by addressing the inherent feature distribution heterogeneity
that one can be identified in many tasks, which can arise from data
heterogeneity (multiple input domains) or model heterogeneity (dynamic
architectures, model conditioning, etc.). Our SaBN factorizes the BN affine
layer into one shared sandwich affine layer, cascaded by several parallel
independent affine layers. Concrete analysis reveals that, during optimization,
SaBN promotes balanced gradient norms while still preserving diverse gradient
directions -- a property that many application tasks seem to favor. We
demonstrate the prevailing effectiveness of SaBN as a drop-in replacement in
four tasks: conditional image generation, neural architecture search (NAS),
adversarial training, and arbitrary style transfer. Leveraging SaBN immediately
achieves better Inception Score and FID on CIFAR-10 and ImageNet conditional
image generation with three state-of-the-art GANs; boosts the performance of a
state-of-the-art weight-sharing NAS algorithm significantly on NAS-Bench-201;
substantially improves the robust and standard accuracies for adversarial
defense; and produces superior arbitrary stylized results. We also provide
visualizations and analysis to help understand why SaBN works. Codes are
available at: https://github.com/VITA-Group/Sandwich-Batch-Normalization.
