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

arXiv: https://arxiv.org/abs/2102.11382 · HTML: https://mlsanity.com/paper/2102.11382/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 61 of its 73 references that are in the map.
- First posted on arXiv 22 Feb 2021.

## Builds on (most cited first, 30 of 61)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Towards Deep Learning Models Resistant to Adversarial Attacks](https://mlsanity.com/paper/1706.06083/) (2017)
- [A Style-Based Generator Architecture for Generative Adversarial Networks](https://mlsanity.com/paper/1812.04948/) (2018)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Progressive Growing of GANs for Improved Quality, Stability, and Variation](https://mlsanity.com/paper/1710.10196/) (2017)
- [Improved Techniques for Training GANs](https://mlsanity.com/paper/1606.03498/) (2016)
- [Large Scale GAN Training for High Fidelity Natural Image Synthesis](https://mlsanity.com/paper/1809.11096/) (2018)
- [Improved Training of Wasserstein GANs](https://mlsanity.com/paper/1704.00028/) (2017)
- [Arbitrary Style Transfer in Real-time with Adaptive Instance Normalization](https://mlsanity.com/paper/1703.06868/) (2017)
- [DARTS: Differentiable Architecture Search](https://mlsanity.com/paper/1806.09055/) (2018)
- [Spectral Normalization for Generative Adversarial Networks](https://mlsanity.com/paper/1802.05957/) (2018)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)
- [Neural Architecture Search with Reinforcement Learning](https://mlsanity.com/paper/1611.01578/) (2016)
- [Semantic Image Synthesis with Spatially-Adaptive Normalization](https://mlsanity.com/paper/1903.07291/) (2019)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Instance Normalization: The Missing Ingredient for Fast Stylization](https://mlsanity.com/paper/1607.08022/) (2016)
- [Efficient Neural Architecture Search via Parameter Sharing](https://mlsanity.com/paper/1802.03268/) (2018)
- [Self-Attention Generative Adversarial Networks](https://mlsanity.com/paper/1805.08318/) (2018)
- [Wasserstein GAN](https://mlsanity.com/paper/1701.07875/) (2017)
- [Feature Squeezing: Detecting Adversarial Examples in Deep Neural Networks](https://mlsanity.com/paper/1704.01155/) (2017)
- [How Does Batch Normalization Help Optimization?](https://mlsanity.com/paper/1805.11604/) (2018)
- [NAS-Bench-101: Towards Reproducible Neural Architecture Search](https://mlsanity.com/paper/1902.09635/) (2019)
- [A Learned Representation For Artistic Style](https://mlsanity.com/paper/1610.07629/) (2016)
- [Defense against Adversarial Attacks Using High-Level Representation Guided Denoiser](https://mlsanity.com/paper/1712.02976/) (2017)
- [Slimmable Neural Networks](https://mlsanity.com/paper/1812.08928/) (2018)
- [Revisiting Batch Normalization For Practical Domain Adaptation](https://mlsanity.com/paper/1603.04779/) (2016)
- [Adversarial Examples Improve Image Recognition](https://mlsanity.com/paper/1911.09665/) (2019)

## Built on by (most cited first, 7 of 7)

- [Dataset Distillation via the Wasserstein Metric](https://mlsanity.com/paper/2311.18531/) (2023)
- [STIV: Scalable Text and Image Conditioned Video Generation](https://mlsanity.com/paper/2412.07730/) (2024)
- [U-Sleep's resilience to AASM guidelines](https://mlsanity.com/paper/2209.11173/) (2022)
- [DiT-Air: Revisiting the Efficiency of Diffusion Model Architecture Design in Text to Image Generation](https://mlsanity.com/paper/2503.10618/) (2025)
- [Semantic Routing: Exploring Multi-Layer LLM Feature Weighting for Diffusion Transformers](https://mlsanity.com/paper/2602.03510/) (2026)
- [GPAS: Accelerating Convergence of LLM Pretraining via Gradient-Preserving Activation Scaling](https://mlsanity.com/paper/2506.22049/) (2025)
- [Towards Understanding Dual BN In Hybrid Adversarial Training](https://mlsanity.com/paper/2403.19150/) (2024)

## Abstract (arXiv)

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.