Removing Batch Normalization Boosts Adversarial Training

TL;DR

Removing batch normalization layers in adversarial training, through the proposed NoFrost method, significantly improves clean and adversarial robustness of neural networks with minimal accuracy loss on clean samples.

Contribution

The paper introduces NoFrost, a normalizer-free adversarial training approach that eliminates batch normalization, enhancing robustness and clean accuracy compared to BN-based methods.

Findings

NoFrost achieves 74.06% clean accuracy on ImageNet with ResNet50.

NoFrost improves adversarial robustness to 23.56% against PGD attack.

NoFrost maintains high model smoothness and larger decision margins.

Abstract

Adversarial training (AT) defends deep neural networks against adversarial attacks. One challenge that limits its practical application is the performance degradation on clean samples. A major bottleneck identified by previous works is the widely used batch normalization (BN), which struggles to model the different statistics of clean and adversarial training samples in AT. Although the dominant approach is to extend BN to capture this mixture of distribution, we propose to completely eliminate this bottleneck by removing all BN layers in AT. Our normalizer-free robust training (NoFrost) method extends recent advances in normalizer-free networks to AT for its unexplored advantage on handling the mixture distribution challenge. We show that NoFrost achieves adversarial robustness with only a minor sacrifice on clean sample accuracy. On ImageNet with ResNet50, NoFrost achieves $74.06\%$ clean accuracy, which drops merely $2.00\%$ from standard training. In contrast, BN-based AT obtains $59.28\%$ clean accuracy, suffering a significant $16.78\%$ drop from standard training. In addition, NoFrost achieves a $23.56\%$ adversarial robustness against PGD attack, which improves the $13.57\%$ robustness in BN-based AT. We observe better model smoothness and larger decision margins from NoFrost, which make the models less sensitive to input perturbations and thus more robust. Moreover, when incorporating more data augmentations into NoFrost, it achieves comprehensive robustness against multiple distribution shifts. Code and pre-trained models are public at https://github.com/amazon-research/normalizer-free-robust-training.