Sandwich Batch Normalization: A Drop-In Replacement for Feature Distribution Heterogeneity

TL;DR

Sandwich Batch Normalization (SaBN) is a simple yet effective modification of BN that addresses feature distribution heterogeneity, improving performance across tasks like image generation, NAS, adversarial training, and style transfer.

Contribution

SaBN introduces a novel affine layer structure that promotes balanced gradients and can be easily integrated into existing models, enhancing their performance.

Findings

Improves Inception Score and FID on CIFAR-10 and ImageNet GANs

Boosts NAS-Bench-201 performance significantly

Enhances adversarial robustness and style transfer quality

Abstract

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.