Sandwich Batch Normalization

Xinyu Gong, Wuyang Chen, Tianlong Chen, Zhangyang Wang

Keywords: normalization

Abstract: We present Sandwich Batch Normalization ($\textbf{SaBN}$), a frustratingly easy improvement of Batch Normalization (BN) with only a few lines of code changes. SaBN is motivated by addressing the inherent $\textit{feature distribution heterogeneity}$ that one can be identified in many tasks, which can arise from model heterogeneity (dynamic architectures, model conditioning, etc.), or data heterogeneity (multiple input domains). A SaBN factorizes the BN affine layer into one shared $\textit{sandwich affine}$ layer, cascaded by several parallel $\textit{independent affine}$ layers. Its variants include further decomposing the normalization layer into multiple parallel ones, and extending similar ideas to instance normalization. We demonstrate the prevailing effectiveness of SaBN (as well as its variants) as a $\textbf{drop-in replacement in four tasks}$: neural architecture search (NAS), image generation, adversarial training, and style transfer. Leveraging SaBN immediately boosts two state-of-the-art weight-sharing NAS algorithms significantly on NAS-Bench-201; achieves better Inception Score and FID on CIFAR-10 and ImageNet conditional image generation with three state-of-the art GANs; substantially improves the robust and standard accuracy for adversarial defense; and produces superior arbitrary stylized results. We also provide visualizations and analysis to help understand why SaBN works. All our codes and pre-trained models will be released upon acceptance.

One-sentence Summary: We present Sandwich Batch Normalization, a plug-and-play module which is able to boost network performance on several tasks, including neural architecture search, conditional image generation, adversarial robustness and neural style transfer.

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics

Community Implementations: [ 2 code implementations](https://www.catalyzex.com/paper/arxiv:2102.11382/code)

Reviewed Version (pdf): https://openreview.net/references/pdf?id=SoENs8Q0oB