Year: 2021

Abstract:
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.

text: 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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 PROBLEM STATEMENT AND ANALYSIS
text: 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.

heading: 2.1 THE EFFECT OF NORMALIZATION IN OVERPARAMETRIZED REGIMES
text: 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.

heading: 3 COUNTERBALANCING TEACHER (CT)

heading: 3.1 TRAINING THE TEACHER ENCODER (I.E. UNNORMALIZED ENCODER)
text: For an arbitrary neural network encoder F parametrized by ζ with one or more batch normalization layers, we first create a clone of the network. We refer to the cloned network as G and its parameters as θ. We then remove all the batch norm layers from the cloned network and train it using the classification (k class) objective of the task in hand LGcls. In our experiments, we minimize the cross-entropy loss as a supervised classification objective:
LGcls(pGθ , y) = − k∑ i yi log(piGθ ), (5)
henceforward we refer to Gθ as the Counterbalancing Teacher or simply CT, and Fζ as the Student.

heading: 3.2 REGULARIZING FOR ROBUSTNESS USING THE COUNTERBALANCING TEACHER (CT) METHOD
text: We subsequently use CT with frozen parameters to regularize the training objective of the Student network. Inspired from the work of Huang and Belongie (2017), we propose the following objective to regularize the representations of Fζ (i.e. Student) using:
Lct ( rGθ , rFζ ) = 1
h h∑ i=1 ( riGθ − r i Fζ )2 (6)
Fζ and Gθ map the input x to latent representations rFθ ∈ Rh and rGθ ∈ Rh respectively. h denotes the dimension of feature vectors. Finally, we add Lct to the cross entropy objective and minimize:
L(rGθ , rFζ , pFζ , y) = LFcls(pFζ , y) + λLct(rGθ , rFζ ) (7)
To calculate the first and and second terms in Equation 6 (RHS), we l2 normalize both rGθ and rFζ . Another variant of the model would be to train both Gθ and Fζ simultaneously, however, as Fζ converges faster, it dominates Gθ. We found the current 2-steps setting working better i.e., training and freezing Gθ in the first step. Similar to knowledge distillation, a limitation of our CT method can be training an extra encoder as a regularizer, however, note that the encoder is removed during inference, therefore the processing time at inference will be similar to that of the original model. The λ in Eq. 7 controls the amount of regularization.

heading: 4 EXPERIMENTAL RESULTS
text: In this section, we are interested in empirically investigating the following questions:
1. Do batch and other normalization techniques lead to learning frequent low variance features?
2. Would a trained model with normalization fail if the dominant features are missing or corrupted, and how would an unnormalized network behave in the same scenario?
3. Does our CT method lead to a reasonable in- and out-of-domain accuracy as well as low corruption error across different architectures and data modalities?

heading: 4.1 BATCH NORMALIZATION IS A CAUSE OF FAILURE WHEN FREQUENT INPUT FEATURES ARE MISSING
text: We start with a synthetic experiment to test the hypothesis that batch normalization leads to learning low variance dominant features. We design a binary classification problem by selecting two digits of “2” and “3” from the MNIST dataset and add a red and a blue square to all samples of class “3” (Figure 2c first row) as “frequent” features.
We examine this hypothesis both qualitatively (visualizing the important input features for the trained classifiers using GradCAM (Selvaraju et al., 2017)) and quantitatively by evaluating on multiple test sets: 1) RB: similar to training samples, digit “3” includes both the red and blue squares, 2) R: digit “3” includes only red square, 3) B: digit “3” includes only blue square, and 4) None: contains no square.
Ideally, a classifier should not rely only on the dominant features to predict the class. However, as shown in the second row of Figure 2c, the network trained with batch normalization quickly picks up those low variance clues, which in this case is either the red or the blue square. However, the same network without normalization takes into account other features as well. We verify this by quantifying classification accuracy on the different test sets. As demonstrated in Figure 2a, since the model relies mostly on the colored squares, it drastically fails when those features are missing at test time (the “None” test set) which is not the case for the network without normalization.

heading: 4.2 BATCH NORMALIZATION IS A CAUSE OF FAILURE WHEN INPUT FEATURES ARE CORRUPTED
text: In this experiment, we focus on a scenario where models are trained with clean data while the test data is shifted by common data corruptions. Here, we use the WideResNet 40-2 and AllConvNet architectures and CIFAR-10-C dataset to be consistent with the work of Hendrycks et al. (2019b). For simplicity and to directly study the effect of normalization under data shifts we do not use dropout and data augmentations in this experiment. We further replace batch normalization with other common normalization techniques such as layer normalization (LN) (Ba et al., 2016) and instance normalization (IN) (Ulyanov et al., 2016). As shown in Figure 3, the model with no normalization (NN) outperforms all the three types of normalization when evaluated on corrupted data. Whereas, on the clean data all three techniques do generally well. With our CT model, corruption errors improves significantly on WideResNet and ResNext, by ∼ 4% and ∼ 3%, respectively.

heading: 4.3 REGULARIZING A BATCH NORMALIZED MODEL USING CT LEADS TO A HIGHER ROBUSTNESS
text: Here we compare our counterbalance teacher model with recent approaches on robustness, multidomain generalization, and corrupted point cloud classification using different data sets and models.
Robustness to common data corruptions. We compare the robustness of our CT model to two main groups of the recent approaches on robustness: 1) covariate shift adaptation and 2) data augmentation-based methods on CIFAR-10-C and CIFAR-100-C (Hendrycks and Dietterich, 2019). Covariate shift adaptation approaches (Benz et al., 2021; Schneider et al., 2020) are designed with the assumption that there exist some samples from (or close to test distribution) which can be used to calibrate the parameters of the batch normalization layers. However, we show that these approaches work only if BN statistics are adapted to a particular corruption. This is unrealistic, especially at inference time, as we often do not have prior knowledge of corruption types that may or may not occur. Additionally, we show adapting batch norm statistics to a single or a few distortion types can often fail to generalize to unseen distortions (Geirhos et al., 2018; Vasiljevic et al., 2016) which is also the case for data augmentation-based methods.
CT vs. covariate shift adaptation methods for robustness. We consider three scenarios for adaptation of batch norm statistics; Let C = {ci | i ∈ Z} be the set of corruptions that can be applied to the input of the network (we assume that the corruption severity level (Hendrycks and Dietterich, 2019) is 3 in all scenarios):
(I) Adapt-one-test-one: BN statistics are adapted to a batch of samples from each corruption type ci at a time, and classification error is computed on the same ci. This is repeated for each corruption type and classification error is averaged over each type (Benz et al., 2021; Schneider et al., 2020). However, it is not trivial to differentiate samples based on corruption type and adapt each separately at test time.
(II) Adapt-one-test-all: BN statistics are adapted to a batch of samples from a randomly selected ci, and the classification error is averaged over C. This experiment reveals the caveats of adaptation on a single corruption, where model fails to generalize to other types of corruptions.
(III) Adapt-all-test-all: BN statistics are adapted to a batch of samples from a random combination of samples from C. The goal is to aid adaptation of statistics by looking at samples from “all” corruption types. The classification error is then averaged over C. This scenario is more realistic than (I) and (II), but still requires a batch of samples from the target domain.
In Figure 4, we shed light on a more realistic scenario of batch norm statistics adaptation. We assume access to some extra samples, but not their corruption types (i.e., adapt-all-test-all). Surprisingly, we regardless of batch size, such methods perform significantly worse on both clean and corruption errors compared to our CT method on CIFAR-10-C with WideResNet model. Further analysis on how adaptation-based methods cause generalization failures are included in Table 6.
Comparing the robustness of CT to self-supervised methods on common input corruptions. In this experiment, we evaluate how recent self-supervised and contrastive learning methods perform on unseen input corruptions: SimCLR (Chen et al., 2020), BYOL (Grill et al., 2020), and Barlow Twins (Zbontar et al., 2021). As reported in Table 1, our method outperforms all self-supervised methods (with fine-tuned encoders) based on the mean corruption error for both CIFAR-10-C and CIFAR-100-C by a large margin. For all the methods in Table 1 we use the same individual augmentations (same operations as in AutoAugment) as in other experiments. See appendix B.2 for linear evaluation results. Concretely, we note that these self-supervised representation learning methods proceed in two stages: pre-training an encoder via self-supervision; and learning a linear classification head on top of the pre-trained encodings. To make it a fair comparison, our modification was to convert this two-stage process into a single joint-training procedure, by using an objective that combines these two terms: (1) a contrastive loss term that uses label information to select positive and negative examples for training the encoder (similar to Supervised Contrastive Learning by (Khosla et al., 2020)); and (2) a cross-entropy loss term for the downstream classification task. In this way, both the encoder and classifier benefit from labeled supervision.
Comparing CT to data augmentation-based methods. Hendrycks et al. (2019b) has studied the effectiveness of different data augmentation techniques on making a model robust to common data corruptions. They proposed to create a weighted mixture of input samples using multiple different augmentation types (AugMix) and showed this method can outperform other augmentationbased methods on robustness to common corruptions. Since our CT model is inherently data augmentation-independent, without any augmentation, it outperforms (by achieving 25.7% mCE while the original mCE is 39.6%) several recent and widely used augmentation-based approaches such as Standard (Hendrycks et al., 2019b), Cutout (DeVries and Taylor, 2017), CutMix (Yun et al., 2019), and Adversarial Training (Madry et al., 2017) on CIFAR-10-C with WideResNet. When CT is used with a few simple data augmentations (AutoAugment (Cubuk et al., 2018) individual operations, but randomly), as shown in Table 2, it outperforms most of the data augmentation-based methods including the recent and advanced ones such as Mixup (Zhang et al., 2017) and AutoAugment (Cubuk et al., 2018), while it achieves comparable results to AugMix without the Jensen-Shannon divergence (JSD) loss on CIFAR-10-C and CIFAR-100-C datasets with different networks.
CT’s performance in domain generalization. In this experiment, we evaluate CT on domain generalization using the VLCS benchmark (Torralba and Efros, 2011). VLCS consists of images from five object categories shared by the PASCAL VOC 2007, LabelMe, Caltech, and Sun datasets, which are considered to be four separate domains. We follow the standard evaluation strategy used in (Carlucci et al., 2019), where we partition each domain into a train (70%) and test set (30%) by random selection from the overall dataset. We use ResNet-18 as the backbone to make a fair comparison with the state-of-the-art. As summarized in Table 3, CT outperforms the state-of-the-art on 3 out of 4 domains and by 1.83% on average.
CT’s classification performance on corrupted 3D point cloud data. In this experiment, we use the RobustPointSet dataset Taghanaki et al. (2020) which is created for analysis of point classifiers in terms of robustness to 3D corruptions. We follow the same training-domain validation setting as in Taghanaki et al. (2020). We train each model without data augmentation on the clean training set and select the best performing checkpoint on the clean validation set for each method. We then test the models on the six distorted unseen test sets. This experiment shows the vulnerability of the models trained on original data to unseen input transformations. As shown in Table 4, PointNet-CT significantly improves classifcation accuracy on most of the shifted test sets, and 1.87% on average compared to original PointNet.

heading: 5 RELATED WORK
text: Reducing classification error on common input corruptions. Hendrycks et al. (2019b) proposed a data augmentation technique which mixes multiple augmented samples using random weights and showed improvements on robust accuracy. Rusak et al. (2020) proposed to mix Guassian noise with adversarial samples for robustness to common distortions. However, as in the real world, the space of distortions and their mixture is not finite, it is not trivial to devise a set of augmentation types that will make a model robust to all distortions. Moreover, making a model robust to certain corruptions via augmentation does not generalize to others. Dodge and Karam (2017) leverage an ensemble approach for robustness against distortions, but they also assume corruptions are known beforehand.
Batch normalization’s vulnerability. Several works have studied the effect of batch normalization in the context of robustness to adversarial perturbations. However, as the most relevant works to ours, Schneider et al. (2020) and Benz et al. (2021) suggested improving robustness to common data corruptions by updating the batch normalization parameters using statistics calculated form extra test distribution samples. However, extra test samples are not always available, and similar to data augmentation-based methods, these approaches work only when a model is adapted to a specific distortion and tested on the same distortion type.
Teacher-student models in the context of robustness. The focus of the few existing methods on leveraging teacher-student models for robustness has been on adversarial perturbations for developing adversarially robust models either by training teacher encoders with adversarial (Goldblum et al., 2020; Papernot et al., 2016) or augmented (Arani et al., 2021) examples. In contrast, our CT method is augmentation-independent. Adversarially trained models might not necessarily be robust to common data corruptions, as we showed in Table 2, our CT method performs significantly better (at least by 12%) on both CIFAR-10-C and CIFAR-100-C datasets across different models compared to the adversarially trained model (AdvT). Even if the augmented samples are crafted by common data augmentations rather than adding adversarial noise, the robustness obtained by the augmentations does not generalize to unseen corruptions, even when they are from the same family of the augmentations a model has seen (Vasiljevic et al., 2016) during training.
Similarities to self-supervised methods. Our CT method is similar to self-supervised models (SSL) (Chen et al., 2020; Grill et al., 2020; He et al., 2020) in the sense that SSL also leverages a copy of the same model but for encoding a different view of the input. In SSL, the copy of the model is the same as the main model, while we discard batch normalization layers in our CT model from one of the copies. Hendrycks et al. (2019a) showed that SSL can improve error on distorted data, however their method—and in general existing self-supervised methods—are highly dependent on data-augmentation (inductive bias) for creating different views while our approach does not require extra views of the data. Note that our method does not employ any self-supervised loss function.

heading: 6 CONCLUSION
text: In this work, we investigated the effect of BN on model robustness. We provided a theoretical justification in the overparameterized linear regime for how BN compromises OOD generalization performance by encouraging the model to rely on highly predictive, low-variance input features. Then, we proposed a robust representation learning framework – Counterbalancing Teacher (CT) – which leverages a frozen copy of the same model without BN as a teacher to regularize the features learned by a student network to improve generalization. Empirically, we demonstrated that our method’s learned representations are robust to common corruptions on a suite of domain generalization tasks.

heading: APPENDIX

heading: A ADDITIONAL EXPERIMENTAL DETAILS
text: We implement our CT variants using Tensorflow Abadi et al. (2016).

heading: A.1 MNIST MISSING FEATURES EXPERIMENT
text: We trained the MNIST models (architecture below) with SGD optimizer (with Tensorflow’s default parameters) and batch size of 128 for 100 epochs. For the NoBN model, we simply removed the BN layers and trained all models without any data augmentation.

heading: A.2 COMPARING CT TO OTHER METHODS ON COMMON INPUT CORRUPTION
text: For both CIFAR-10-C and CIFAR-100-C we train Gθ using Adam optimizer with learning rate of 0.0001 and batch size of 64. In the second step, we train Fζ using SGD with Nesterov, initial learning rate of 0.1 and decaying to 0.00001 using cosine scheduler for 300 epochs.

heading: A.3 APPLYING CT TO POINT CLOUDS
text: We trained both Gθ and Fζ with Adam optimizer and learning rate of 0.001 (divided by 10 at epochs 50 and 75) and batch size of 32 for 500 epochs.

heading: A.4 APPLYING CT TO DOMAIN GENERALIZATION
text: We trained CT with SGD optimizer with learning rate of 0.001 and momentum of 0.9. We set batch size to 32 and image size to 224x224, and train for 500 epochs.

heading: A.5 SELF-SUPERVISED BASELINES
text: We train the self-supervised methods with batch size 256 for 1000 epochs while maintaining the original hyper-parameters. We fine-tune all self-supervised methods for 350 epochs with Adam Kingma and Ba (2014) and with starting learning rate of 0.001 with a cosine decay, and l2 coefficient set to 0.0005 to avoid overfitting.

heading: B ADDITIONAL EXPERIMENTAL RESULTS

heading: B.1 BATCH-NORM ADAPTATION ADDITIONAL RESULTS
text: Batch norm adaptation-based results for robustness to common input corruptions on CIFAR-10-C.

heading: B.2 LINEAR EVALUATION OF SELF-SUPERVISED MODELS
text: Self-supervised linear evaluation results where encoders are frozen and only classification layer is trained.

heading: B.3 PREDICTIVE PERFORMANCE OF MNIST DATASET VARIABLES.
text: The following analysis is performed to determine whether the dark and light pixels in the MNIST experiment are discriminative. We plot histograms of all pixel values for each class separately. The histograms look very similar, as shown in Figure 5, indicating that the dark or light pixels are not predictive alone. Therefore we add red or blue squares/variables which are both low variance and predictive.

heading: C ADDITIONAL ANALYSIS ON THE EFFECT OF NORMALIZATION IN OVERPARAMETERIZED REGIME

heading: C.1 SUBTRACTING MEAN DOES NOT CHANGE THE ESTIMATED PARAMETERS
text: In Section 2.1, we show that dividing each feature by its variance incentivize the model to rely more on features with low variance. Here we show that subtracting mean from each feature and target does not change the estimated parameter in overparametrized regime. Recall that X ∈ Rn×d denote the training examples and Y ∈ Rn denote their target. Let µX ∈ Rd be the mean of features in training data, and µY ∈ R denote the mean of the their targets. Since µX is a linear combination of the rows of X and row operation does not change the projection matrix; therefore, the following two linear programs have the same solution:
ζ̂ = arg min ‖ζ‖ s.t. Xζ = Y
θ̂ = arg min ‖θ‖ s.t. (X − µX)θ = Y −~1µY

heading: C.2 THE EFFECT OF NORMALIZATION ON MAX-MARGIN CLASSIFIERS
text: In Section 2.1, we analyze the effect of normalization in linear regression. Here we show that the analysis hold for the max-margin classifiers as well. Soudry et al., Soudry et al. (2018), show that without any explicit regularization, gradient descent on logistic loss converges to the L2 maximum margin separator for all linearly separable datasets. In the overparametrized regime data are completely separable, there are also some investigation that as d increases all the data points serve a support vectors Narang (2020).
Without normalization we have:
ζ̂ = arg min ‖ζ‖22 s.t. Y Xζ ≥ 1 (8)
Recall that U is a diagonal matrix where Uii is the standard deviation of the ith feature. After normalization we have:
β̂ = arg min ‖β‖22 s.t. Y XU−1β ≥ 1 (9)
At the test time we predict sign(XU−1β̂), substituting U−1β with θ we have:
θ̂ = arg min ‖Uθ‖22 s.t. Y Xθ ≥ 1 (10)
Similar to linear regression scenario normalizing data change the inductive bias of estimator to choose a model that rely more on low variance features. This data-dependent norm lead to good performance in-domain where such features exhibit low variation but will perform poorly out of domain data are corrupted independently or are sampled from a different distribution with different statistic than the training distribution.

Review summaries:
The authors study the effects of BatchNorm layers on model robustness.  They demonstrate that models trained with BatchNorm are likely to rely on low-variance features which are performant in-domain, but are not helpful for out of domain data.  Leveraging this observation, they propose a student-teacher method (Counterbalancing Teacher) that creates a copy of the same model with all BatchNorm layers removed.  The demonstrate that this leads to very competitive performance with other techniques for reducing the performance hit of out domain evaluation: self-supervised methods and data augmentation methods. 

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. 

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.