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 authors propose a new consistency loss for improving model robustness to common image corruptions. They use a student-teacher training setup where only the student network uses batch normalization at training time. Improvements are shown on small scale corruption datasets (CIFAR-C), a single domain generalization dataset (VLCS), and RobustPointSet.