Counterbalancing Teacher: Regularizing Batch Normalized Models for Robustness

Batch normalization ( BN ) is a ubiquitous technique for training deep neural networks that accelerates their convergence to reach higher accuracy . However , we demonstrate that BN comes with a fundamental drawback : it incentivizes the model to rely on frequent low-variance features that are highly specific to the training ( in-domain ) data , and thus fails to generalize to out-of-domain examples . In this work , we investigate this phenomenon by first showing that removing BN layers across a wide range of architectures leads to lower out-of-domain and corruption errors at the cost of higher in-domain error . We then propose the Counterbalancing Teacher ( CT ) method , which leverages a frozen copy of the same model without BN as a teacher to enforce the student network ’ s learning of robust representations by substantially adapting its weights through a consistency loss function . This regularization signal helps CT perform well in unforeseen data shifts , even without information from the target domain as in prior works . We theoretically show in an overparameterized linear regression setting why normalization leads a model ’ s reliance on such in-domain features , and empirically demonstrate the efficacy of CT by outperforming several methods on standard robustness benchmark datasets such as CIFAR-10-C , CIFAR-100-C , and VLCS . 1 INTRODUCTION . Batch normalization ( BN ) , a neural network layer that normalizes input features by aggregating batch statistics during training , is a key component for accelerating convergence in the modern deep learning toolbox ( Ioffe and Szegedy , 2015 ; Santurkar et al. , 2018 ; Bjorck et al. , 2018 ) . It plays a critical role in stabilizing training dynamics for large models optimized with stochastic gradient descent , and has since spurred a flurry of research in related modifications ( Ba et al. , 2016 ; Kingma and Ba , 2014 ) and its understanding ( Gitman and Ginsburg , 2017 ; Santurkar et al. , 2018 ; Luo et al. , 2018 ; Kohler et al. , 2018 ) . Despite its advantages , BN has recently been shown to be a source of vulnerability to adversarial perturbations ( Galloway et al. , 2019 ; Benz et al. , 2021 ) . In our work , we take this observation one step further and demonstrate that BN also compromises a model ’ s out-of-domain ( OOD ) generalization capabilities . Specifically , we demonstrate that normalization incentivizes the model to exploit highly predictive , low-variance features ( Geirhos et al. , 2018 ; 2020 ) , that lead to poor classification accuracy when the test environment differs from that of training . Given the widespread use and benefits of normalization , we desire a way to mitigate such drawbacks in models trained with BN . To better understand this phenomenon , we investigate the effect of normalization in over-parametrized regimes , where there exist multiple solutions and inductive bias ( e.g. , minimizing the norm of the weights ) significantly impacts the estimated parameters . Similar to recent work in the theory of deep learning ( Khani and Liang , 2021 ; Raghunathan et al. , 2020 ; Nakkiran , 2019 ; Hastie et al. , 2019 ; Liang et al. , 2020 ) , we study the min-norm solution in over-parametrized linear regression . Without normalization , the inductive bias selects a model that fits training data and minimizes a fixed norm independent of data ; with normalization , the same inductive bias selects a model that minimizes a data-dependent norm , leading the model to rely more on low-variance features . While such highly predictive features yield better performance in-domain where the features do not vary significantly , they cause performance to plummet in OOD settings ( e.g . data corruptions or missing features ) . This is in direct contrast to models trained without BN that assign equal weight to all input features , which help to reduce their overfitting on the training set . Drawing inspiration from our observation and the knowledge distillation literature ( Hinton et al. , 2015 ; Romero et al. , 2014 ) , we propose a simple teacher-student model to combine the best of both worlds : we leverage features derived both from a network without BN ( teacher ) and its clone with BN ( student ) to learn representations that achieve high standard and robust accuracies . We incorporate a regularization term in the loss function which encourages the features learned from the student encoder to have similar statistics and structure to those learned from the teacher ; we name this model the Counterbalancing Teacher ( CT ) and show that it helps in achieving both higher robust and clean accuracy compared to a ( batch ) normalized model . In particular , CT retains good performance in OOD settings even without knowledge of statistics of the new domain . Our results mark a significant improvement over prior works , which have tackled similar problems by either : ( a ) modifying the statistics of a trained model ( Schneider et al. , 2020 ; Benz et al. , 2021 ) using privileged information from the target domain ; or ( b ) augmenting the training data using a set of predefined corruption functions ( Hendrycks et al. , 2019b ) . As recent studies ( Vasiljevic et al. , 2016 ; Geirhos et al. , 2018 ; Taghanaki et al. , 2020 ) show that such approaches often fail to generalize due to the tendency of neural networks to memorize data-specific properties , this motivates a shift towards developing models that are inherently robust , independent of data augmentation or input transformation . Empirically , we demonstrate that CT outperforms most existing data augmentation-based techniques and covariate shift adaption-based methods ( which require information from the test set ) on mean corruption error on CIFAR10-C and CIFAR100-C ( Hendrycks and Dietterich , 2019 ) , and achieves state-of-the-art performance in domain generalization on the VLCS dataset ( Torralba and Efros , 2011 ) . We further test CT on corrupted 3D point-cloud data ( Taghanaki et al. , 2020 ) and show it outperforms existing methods in terms of mean classification accuracy over multiple test sets . To the best of our knowledge , this is the first work to explore both theoretically and empirically why BN leads to a model ’ s over-reliance on frequent , low-variance features , which can adversely affect its performance on the downstream classification task . This is also the first work to present a robust representation learning framework for common input distortions without additional data augmentation strategies or information derived from the target domain . In summary , our contribution is threefold : 1 . We provide theoretical justifications for why normalization encourages a model to exploit low-variance features , and empirically evaluate how this behavior can adversely affect downstream classification accuracy . 2 . We propose CT , a representation learning approach demonstrating that regularizing representations of a batch-normalized network ( i.e. , student ) using those from an unnormalized copy ( i.e. , teacher ) can significantly improve a model ’ s robustness . 3 . We experimentally verify the robustness of the representations learned by CT to input distortions and domain shift on a variety of tasks and models . 2 PROBLEM STATEMENT AND ANALYSIS . We assume a supervised classification setting : given an input variable x ∈ X ⊆ Rd , and a set of corresponding labels y ∈ Y = { 1 , ... , k } , we aim to learn a classifier fζ : X −→ Y by minimizing the empirical risk : ζ = argmin ζ Ex , y∼pd ( x , y ) [ ` ( x , y ; ζ ) ] ≈ argmin ζ n∑ i=1 ` ( fζ ( xi ) , yi ) ( 1 ) Here pd ( x , y ) is the underlying joint distribution where the dataset D = { ( xi , yi ) } ni=1 is sampled from . In the following sections , we first discuss the effect of ( batch ) normalization on the solutions found in underspecified regimes , then elaborate on our approach to optimize the aforementioned empirical risk . In this context , we refer to a problem as underspecified or overparameterized when degrees of freedom of a model is larger than the number of training samples . 2.1 THE EFFECT OF NORMALIZATION IN OVERPARAMETRIZED REGIMES . Modern deep learning frameworks usually incorporate many parameters ( often larger than the number of training data points ) , which lead to underspecified regimes . In other words , many distinct solutions solve the problem equally i.e. , have the same training or even held-out loss ( D ’ Amour et al. , 2020 ) . In the underspecified regime , the inductive bias of the estimation procedure , such as choosing parameters with the minimum norm , significantly impacts the estimated parameters . In such regimes , we show that normalizing data incentivizes the model to rely on features with lower variance . We analyze the effect of normalization on the min-norm solution in overparametrized noiseless linear regression . This setup has been studied in many recent works for understanding some phenomena in deep networks ( Khani and Liang , 2021 ; Raghunathan et al. , 2020 ; Nakkiran , 2019 ) . Let X ∈ Rn×d denote training examples and Y ∈ Rn denote their target . Considering that we are in an over parametrized regime ( d > n ) , there should be an equivalence class of solutions . We assume that the inductive bias of the model is to choose the min-norm solution ( the parameter with the minimum ` 2 norm ) . This is in line with the recent speculation that the inductive bias in deep networks tends to find a solution with minimum norm ( Gunasekar et al. , 2018 ) . One can show that the convergence point of gradient descent run on the least-squares loss is the min-norm solution . Without normalization , the model chooses the min-norm solution which fits the training data : ζ̂ = arg min ζ ‖ζ‖22 s.t . Xζ = Y . ( 2 ) Now we observe how normalization changes the estimated parameters . Let U be a diagonal matrix where Uii denotes the standard deviation of the ith feature . By normalization , we transform X to XU−1 ( for simplicity , we assume the mean of each feature is 0 , we can show that transforming points do not change the estimated parameter , see Appendix C.1 for details ) . In this case the model estimates β̂ as follows : β̂ = arg min β ‖β‖22 s.t . XU−1β = Y , ( 3 ) and since we normalize data points at the test time as well , the estimated parameter used for prediction at the test time is θ̂ = U−1β̂ . Substituting θ instead of U−1β ( thus Uθ = β ) , we can write the equal formulation of 3 as : θ̂ = arg min θ ‖Uθ‖22 s.t . Xθ = Y . ( 4 ) For the same equivalence class of solutions a model with normalization ( 4 ) chooses different parameters in comparison to a model without normalization ( 2 ) . In particular , 2 chooses an interpolant with a minimum data independent norm . On the other hand , 4 , chooses an interpolant with a minimum data-dependent norm , which incentives the model to assign higher weights to low variance features . Note that projection of θ̂ and ζ̂ is the same in column space of training points . Formally if Π = X > ( XX > ) −1X denote the column space of training points then Πθ̂ = Πζ̂ . However , their projections to the null space of training points ( I − Π ) are different . As a result as we have more data ( smaller null space ) , θ̂ and ζ̂ become closer , and converge when n > d. Our analysis hold for classification with max-margin , we only need to substitute Xθ = Y by Y Xθ ≥ 1 ( see Appendix C.2 for details ) . We conjecture that minimizing the data-dependent norm in each layer leads to reliance on low variance ( frequent ) features , which can result in a better in-domain generalization as these feature do not exhibit high variations . Nonetheless , in a new domain where some ( or all ) of the training-domain features are missing or altered ( e.g. , when there is some data-agnostic corruption such as Gaussian noise ) , 2 performs better as its inductive bias is data independent . How should we change the regularization such that it selects for a model that performs well both in- and out-of-domain ? Inspired by this analysis , we introduce a simple , yet powerful two-step approach that combines the normalized and unnormalized copies of the same network for robust representation learning .

Summary: The paper aims to address the problem of out-of-domain generalization with batch normalization layers. It first identifies the reliance on low-variance features in batch normalization, and proposes a counterbalancing teacher approach to distill from a BN-model to a non-BN model. Experiments verify the hypothesis of the paper and show that the proposed approach outperforms the various baselines on robustness to data corruption and domain generalization. The proposed approach is also evaluated on a 3D point cloud dataset.