Original: C OUNTERBALANCING T EACHER : R EGULARIZING B ATCH N ORMALIZED M ODELS FOR R OBUSTNESS
Revised: C OUNTERBALANCING T EACHER : R EGULARIZING B ATCH N ORMALIZED M ODELS FOR R OBUSTNESS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Batch normalization (BN) is a ubiquitous technique for training deep neural networks that accelerates their convergence to reach higher accuracy.
Revised: Batch normalization (BN) is a ubiquitous technique for training deep neural networks that accelerates their convergence to reach higher accuracy.

Original: However, we demonstrate that BN comes with a fundamental drawback: it incentivizes the model to rely on frequent low-variance features that are highly speciﬁc to the training (in-domain) data, and thus fails to generalize to out-of-domain examples.
Revised: However, we demonstrate that BN comes with a fundamental drawback: it incentivizes the model to rely on frequent low-variance features that are highly speciﬁc to the training (in-domain) data, and thus fails to generalize to out-of-domain examples.

Original: In this work, we investigate this phenomenon by ﬁrst 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.
Revised: In this work, we investigate this phenomenon by ﬁrst 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.

Original: We then propose the Counterbalancing
Revised: We then propose the Counterbalancing

Original: Teacher (CT) method, which leverages a frozen copy of the same model without
Revised: Teacher (CT) method, which leverages a frozen copy of the same model without

Original: BN as a teacher to enforce the student network’s learning of robust representations by substantially adapting its weights through a consistency loss function.
Revised: BN as a teacher to enforce the student network’s learning of robust representations by substantially adapting its weights through a consistency loss function.

Original: This regularization signal helps CT perform well in unforeseen data shifts, even without information from the target domain as in prior works.
Revised: This regularization signal helps CT perform well in unforeseen data shifts, even without information from the target domain as in prior works.

Original: 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 efﬁcacy of
Revised: 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 efﬁcacy of

Original: CT by outperforming several methods on standard robustness benchmark datasets such as CIFAR-10-C, CIFAR-100-C, and VLCS.
Revised: CT by outperforming several methods on standard robustness benchmark datasets such as CIFAR-10-C, CIFAR-100-C, and VLCS.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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).
Revised: 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).

Original: It plays a critical role in stabilizing training dynamics for large models optimized with stochastic gradient descent, and has since spurred a ﬂurry of research in related modiﬁcations (Ba et al., 2016; Kingma and Ba, 2014) and its understanding (Gitman and Ginsburg, 2017; Santurkar et al., 2018;
Revised: It plays a critical role in stabilizing training dynamics for large models optimized with stochastic gradient descent, and has since spurred a ﬂurry of research in related modiﬁcations (Ba et al., 2016; Kingma and Ba, 2014) and its understanding (Gitman and Ginsburg, 2017; Santurkar et al., 2018;

Original: Luo et al., 2018; Kohler et al., 2018).
Revised: Luo et al., 2018; Kohler et al., 2018).

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: Speciﬁcally, we demonstrate that normalization incentivizes the model to exploit highly predictive, low-variance features (Geirhos et al., 2018; 2020), that lead to poor classiﬁcation accuracy when the test environment differs from that of training.
Revised: Speciﬁcally, we demonstrate that normalization incentivizes the model to exploit highly predictive, low-variance features (Geirhos et al., 2018; 2020), that lead to poor classiﬁcation accuracy when the test environment differs from that of training.

Original: Given the widespread use and beneﬁts of normalization, we desire a way to mitigate such drawbacks in models trained with BN.
Revised: Given the widespread use and beneﬁts of normalization, we desire a way to mitigate such drawbacks in models trained with BN.

Original: 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) signiﬁcantly impacts the estimated parameters.
Revised: 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) signiﬁcantly impacts the estimated parameters.

Original: Similar to recent work in the theory of deep learning (Khani and Liang, 2021; Raghunathan et al., 2020; Nakkiran, 2019; Hastie et al., 2019;
Revised: Similar to recent work in the theory of deep learning (Khani and Liang, 2021; Raghunathan et al., 2020; Nakkiran, 2019; Hastie et al., 2019;

Original: Liang et al., 2020), we study the min-norm solution in over-parametrized linear regression.
Revised: Liang et al., 2020), we study the min-norm solution in over-parametrized linear regression.

Original: Without normalization, the inductive bias selects a model that ﬁts training data and minimizes a ﬁxed 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.
Revised: Without normalization, the inductive bias selects a model that ﬁts training data and minimizes a ﬁxed 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.

Original: While such highly predictive features yield better performance in-domain where the features do not vary signiﬁcantly, they cause performance to plummet in OOD settings (e.g. data corruptions or missing features).
Revised: While such highly predictive features yield better performance in-domain where the features do not vary signiﬁcantly, they cause performance to plummet in OOD settings (e.g. data corruptions or missing features).

Original: This
Revised: This

Original: ℱ !
Revised: ℱ !

Original: 𝑥 𝑝 ℱ !
Revised: 𝑥 𝑝 ℱ !

Original: 𝒢 # ℒ !"
Revised: 𝒢 # ℒ !"

Original: (𝑟 𝒢 ! ,𝑟 ℱ " ) 𝑟 ℱ !
Revised: (𝑟 𝒢 ! ,𝑟 ℱ " ) 𝑟 ℱ !

Original: 𝑝 𝒢 " ↚ ↚ 𝑟
Revised: 𝑝 𝒢 " ↚ ↚ 𝑟

Original: 𝒢 " 𝑐 𝒢 " 𝑐 ℱ ! is in direct contrast to models trained without BN that assign equal weight to all input features, which help to reduce their overﬁtting on the training set.
Revised: 𝒢 " 𝑐 𝒢 " 𝑐 ℱ ! is in direct contrast to models trained without BN that assign equal weight to all input features, which help to reduce their overﬁtting on the training set.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: In particular, CT retains good performance in
Revised: In particular, CT retains good performance in

Original: OOD settings even without knowledge of statistics of the new domain.
Revised: OOD settings even without knowledge of statistics of the new domain.

Original: Our results mark a signiﬁcant 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 predeﬁned corruption functions (Hendrycks et al., 2019b).
Revised: Our results mark a signiﬁcant 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 predeﬁned corruption functions (Hendrycks et al., 2019b).

Original: As recent studies (Vasiljevic et al., 2016;
Revised: As recent studies (Vasiljevic et al., 2016;

Original: 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-speciﬁc properties, this motivates a shift towards developing models that are inherently robust, independent of data augmentation or input transformation.
Revised: 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-speciﬁc properties, this motivates a shift towards developing models that are inherently robust, independent of data augmentation or input transformation.

Original: 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).
Revised: 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).

Original: We further test CT on corrupted 3D point-cloud data (Taghanaki et al., 2020) and show it outperforms existing methods in terms of mean classiﬁcation accuracy over multiple test sets.
Revised: We further test CT on corrupted 3D point-cloud data (Taghanaki et al., 2020) and show it outperforms existing methods in terms of mean classiﬁcation accuracy over multiple test sets.

Original: To the best of our knowledge, this is the ﬁrst 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 classiﬁcation task.
Revised: To the best of our knowledge, this is the ﬁrst 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 classiﬁcation task.

Original: This is also the ﬁrst work to present a robust representation learning framework for common input distortions without additional data augmentation strategies or information derived from the target domain.
Revised: This is also the ﬁrst work to present a robust representation learning framework for common input distortions without additional data augmentation strategies or information derived from the target domain.

Original: In summary, our contribution is threefold:
Revised: In summary, our contribution is threefold:

Original: We provide theoretical justiﬁcations for why normalization encourages a model to exploit low-variance features, and empirically evaluate how this behavior can adversely affect downstream classiﬁcation accuracy.
Revised: We provide theoretical justiﬁcations for why normalization encourages a model to exploit low-variance features, and empirically evaluate how this behavior can adversely affect downstream classiﬁcation accuracy.

Original: We propose CT, a representation learning approach demonstrating
Revised: We propose CT, a representation learning approach demonstrating

Original: that regularizing representations of a batch-normalized network (i
Revised: that regularizing representations of a batch-normalized network (i

Original: e., student) using those from an unnormalized copy (i.e., teacher) can signiﬁcantly improve a model’s robustness.
Revised: e., student) using those from an unnormalized copy (i.e., teacher) can signiﬁcantly improve a model’s robustness.

Original: We experimentally verify the robustness of the representations learned by CT to input distortions and domain shift on a variety of tasks and models.
Revised: We experimentally verify the robustness of the representations learned by CT to input distortions and domain shift on a variety of tasks and models.

Original:  P ROBLEM STATEMENT AND ANALYSIS
Revised:  P ROBLEM STATEMENT AND ANALYSIS

Original: We assume a supervised classiﬁcation setting: given an input variable x ∈ X ⊆ R d , and a set of corresponding labels
Revised: We assume a supervised classiﬁcation setting: given an input variable x ∈ X ⊆ R d , and a set of corresponding labels

Original: y ∈ Y
Revised: y ∈ Y

Original: = { 1 , ..., k } , we aim to learn a classiﬁer f ζ : X −→ Y by minimizing the empirical risk:
Revised: = { 1 , ..., k } , we aim to learn a classiﬁer f ζ : X −→ Y by minimizing the empirical risk:

Original: [Equation]
Revised: [Equation]

Original: Here p d ( x, y ) is the underlying joint distribution where the dataset D = { ( x i , y i ) } ni =1 is sampled from.
Revised: Here p d ( x, y ) is the underlying joint distribution where the dataset D = { ( x i , y i ) } ni =1 is sampled from.

Original: In the following sections, we ﬁrst discuss the effect of (batch) normalization on the solutions found in underspeciﬁed regimes, then elaborate on our approach to optimize the aforementioned empirical risk.
Revised: In the following sections, we ﬁrst discuss the effect of (batch) normalization on the solutions found in underspeciﬁed regimes, then elaborate on our approach to optimize the aforementioned empirical risk.

Original: In this context, we refer to a problem as underspeciﬁed or overparameterized when degrees of freedom of a model is larger than the number of training samples.
Revised: In this context, we refer to a problem as underspeciﬁed or overparameterized when degrees of freedom of a model is larger than the number of training samples.

Original: 2. T HE EFFECT OF NORMALIZATION IN OVERPARAMETRIZED REGIMES
Revised: 2. T HE EFFECT OF NORMALIZATION IN OVERPARAMETRIZED REGIMES

Original: Modern deep learning frameworks usually incorporate many parameters (often larger than the number of training data points), which lead to underspeciﬁed regimes.
Revised: Modern deep learning frameworks usually incorporate many parameters (often larger than the number of training data points), which lead to underspeciﬁed regimes.

Original: 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).
Revised: 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).

Original: In the underspeciﬁed regime, the inductive bias of the estimation procedure, such as choosing parameters with the minimum norm, signiﬁcantly impacts the estimated parameters.
Revised: In the underspeciﬁed regime, the inductive bias of the estimation procedure, such as choosing parameters with the minimum norm, signiﬁcantly impacts the estimated parameters.

Original: In such regimes, we show that normalizing data incentivizes the model to rely on features with lower variance.
Revised: In such regimes, we show that normalizing data incentivizes the model to rely on features with lower variance.

Original: We analyze the effect of normalization on the min-norm solution in overparametrized noiseless linear regression.
Revised: We analyze the effect of normalization on the min-norm solution in overparametrized noiseless linear regression.

Original: 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).
Revised: 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).

Original: Let X ∈ R n × d denote training examples and Y ∈ R n denote their target.
Revised: Let X ∈ R n × d denote training examples and Y ∈ R n denote their target.

Original: Considering that we are in an over parametrized regime ( d > n ), there should be an equivalence class of solutions.
Revised: Considering that we are in an over parametrized regime ( d > n ), there should be an equivalence class of solutions.

Original: We assume that the inductive bias of the model is to choose the min-norm solution (the parameter with the minimum (cid:96) 2 norm).
Revised: We assume that the inductive bias of the model is to choose the min-norm solution (the parameter with the minimum (cid:96) 2 norm).

Original: This is in line with the recent speculation that the inductive bias in deep networks tends to ﬁnd a solution with minimum norm (Gunasekar et al., 2018).
Revised: This is in line with the recent speculation that the inductive bias in deep networks tends to ﬁnd a solution with minimum norm (Gunasekar et al., 2018).

Original: One can show that the convergence point of gradient descent run on the least-squares loss is the min-norm solution.
Revised: One can show that the convergence point of gradient descent run on the least-squares loss is the min-norm solution.

Original: Without normalization, the model chooses the min-norm solution which ﬁts the training data:
Revised: Without normalization, the model chooses the min-norm solution which ﬁts the training data:

Original: [Equation]
Revised: [Equation]

Original: Now we observe how normalization changes the estimated parameters.
Revised: Now we observe how normalization changes the estimated parameters.

Original: Let U be a diagonal matrix where U ii denotes the standard deviation of the i th feature.
Revised: Let U be a diagonal matrix where U ii denotes the standard deviation of the i th feature.

Original: By normalization, we transform X to
Revised: By normalization, we transform X to

Original: 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).
Revised: 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).

Original: In this case the model estimates ˆ β as follows:
Revised: In this case the model estimates ˆ β as follows:

Original: [Equation]
Revised: [Equation]

Original: and since we normalize data points at the test time as well, the estimated parameter used for prediction at the test time is ˆ θ =
Revised: and since we normalize data points at the test time as well, the estimated parameter used for prediction at the test time is ˆ θ =

Original: U − 1 ˆ β .
Revised: U − 1 ˆ β .

Original: Substituting θ instead of U − 1 β (thus Uθ = β ), we can write the equal formulation of 3 as:
Revised: Substituting θ instead of U − 1 β (thus Uθ = β ), we can write the equal formulation of 3 as:

Original: [Equation]
Revised: [Equation]

Original: For the same equivalence class of solutions a model with normalization (4) chooses different parameters in comparison to a model without normalization (2).
Revised: For the same equivalence class of solutions a model with normalization (4) chooses different parameters in comparison to a model without normalization (2).

Original: In particular, 2 chooses an interpolant with a minimum data independent norm.
Revised: In particular, 2 chooses an interpolant with a minimum data independent norm.

Original: 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.
Revised: 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.

Original: Note that projection of ˆ θ and ˆ ζ is the same in column space of training points.
Revised: Note that projection of ˆ θ and ˆ ζ is the same in column space of training points.

Original: Formally if Π = X (cid:62) ( XX (cid:62) ) − 1 X denote the column space of training points then Πˆ θ =
Revised: Formally if Π = X (cid:62) ( XX (cid:62) ) − 1 X denote the column space of training points then Πˆ θ =

Original: Πˆ ζ .
Revised: Πˆ ζ .

Original: However, their projections to the null space of training points ( I − Π ) are different.
Revised: However, their projections to the null space of training points ( I − Π ) are different.

Original: As a result as we have more data (smaller null space), ˆ θ and ˆ ζ become closer, and converge when n > d .
Revised: As a result as we have more data (smaller null space), ˆ θ and ˆ ζ become closer, and converge when n > d .

Original: Our analysis hold for classiﬁcation with max-margin, we only need to substitute Xθ
Revised: Our analysis hold for classiﬁcation with max-margin, we only need to substitute Xθ

Original: =
Revised: =

Original: Y by Y (cid:12) Xθ ≥ 1 (see
Revised: Y by Y (cid:12) Xθ ≥ 1 (see

Original: Appendix C.2 for details).
Revised: Appendix C.2 for details).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: How should we change the regularization such that it selects for a model that performs well both in- and out-of-domain?
Revised: How should we change the regularization such that it selects for a model that performs well both in- and out-of-domain?

Original: 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.
Revised: 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.

Original:  C OUNTERBALANCING T EACHER (CT) 3. T RAINING THE TEACHER ENCODER ( I . E . UNNORMALIZED ENCODER )
Revised:  C OUNTERBALANCING T EACHER (CT) 3. T RAINING THE TEACHER ENCODER ( I . E . UNNORMALIZED ENCODER )

Original: For an arbitrary neural network encoder F parametrized by ζ with one or more batch normalization layers, we ﬁrst create a clone of the network.
Revised: For an arbitrary neural network encoder F parametrized by ζ with one or more batch normalization layers, we ﬁrst create a clone of the network.

Original: We refer to the cloned network as G and its parameters as θ .
Revised: We refer to the cloned network as G and its parameters as θ .

Original: We then remove all the batch norm layers from the cloned network and train it using the classiﬁcation ( k class) objective of the task in hand L G cls .
Revised: We then remove all the batch norm layers from the cloned network and train it using the classiﬁcation ( k class) objective of the task in hand L G cls .

Original: In our experiments, we minimize the cross-entropy loss as a supervised classiﬁcation objective:
Revised: In our experiments, we minimize the cross-entropy loss as a supervised classiﬁcation objective:

Original: [Equation]
Revised: [Equation]

Original: henceforward we refer to G θ as the Counterbalancing Teacher or simply CT, and F ζ as the Student.
Revised: henceforward we refer to G θ as the Counterbalancing Teacher or simply CT, and F ζ as the Student.

Original: 3. R EGULARIZING FOR ROBUSTNESS USING THE COUNTERBALANCING TEACHER (CT) METHOD
Revised: 3. R EGULARIZING FOR ROBUSTNESS USING THE COUNTERBALANCING TEACHER (CT) METHOD

Original: We subsequently use CT with frozen parameters to regularize the training objective of the Student network.
Revised: We subsequently use CT with frozen parameters to regularize the training objective of the Student network.

Original: Inspired from the work of Huang and Belongie (2017), we propose the following objective to regularize the representations of F ζ (i.e. Student) using:
Revised: Inspired from the work of Huang and Belongie (2017), we propose the following objective to regularize the representations of F ζ (i.e. Student) using:

Original: [Equation]
Revised: [Equation]

Original: F ζ and G θ map the input x to latent representations r F θ ∈ R h and r G θ ∈ R h respectively.
Revised: F ζ and G θ map the input x to latent representations r F θ ∈ R h and r G θ ∈ R h respectively.

Original: h denotes the dimension of feature vectors.
Revised: h denotes the dimension of feature vectors.

Original: Finally, we add L ct to the cross entropy objective and minimize:
Revised: Finally, we add L ct to the cross entropy objective and minimize:

Original: [Equation]
Revised: [Equation]

Original: To calculate the ﬁrst and and second terms in Equation 6 (RHS), we l 2 normalize both r G θ and r F ζ .
Revised: To calculate the ﬁrst and and second terms in Equation 6 (RHS), we l 2 normalize both r G θ and r F ζ .

Original: Another variant of the model would be to train both G θ and F ζ simultaneously, however, as F ζ converges faster, it dominates G θ .
Revised: Another variant of the model would be to train both G θ and F ζ simultaneously, however, as F ζ converges faster, it dominates G θ .

Original: We found the current 2-steps setting working better i.e., training and freezing G θ in the ﬁrst step.
Revised: We found the current 2-steps setting working better i.e., training and freezing G θ in the ﬁrst step.

Original: 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. 
Revised: 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.

Revised: 7 controls the amount of regularization.

Original:  E XPERIMENTAL RESULTS
Revised:  E XPERIMENTAL RESULTS

Original: In this section, we are interested in empirically investigating the following questions:
Revised: In this section, we are interested in empirically investigating the following questions:

Original: Do batch and other normalization techniques lead to learning frequent low variance features?
Revised: Do batch and other normalization techniques lead to learning frequent low variance features?

Original: [Figure]
Revised: [Figure]

Original: 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
Revised: 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

Original: 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
Revised: 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

Original: 4. B ATCH NORMALIZATION IS A CAUSE OF FAILURE WHEN FREQUENT INPUT FEATURES ARE MISSING
Revised: 4. B ATCH NORMALIZATION IS A CAUSE OF FAILURE WHEN FREQUENT INPUT FEATURES ARE MISSING

Original: We start with a synthetic experiment to test the hypothesis that batch normalization leads to learning low variance dominant features.
Revised: We start with a synthetic experiment to test the hypothesis that batch normalization leads to learning low variance dominant features.

Original: We design a binary classiﬁcation 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 2b ﬁrst row) as “frequent” features.
Revised: We design a binary classiﬁcation 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 2b ﬁrst row) as “frequent” features.

Original: We examine this hypothesis both qualitatively (visualizing the important input features for the trained classiﬁers 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)
Revised: We examine this hypothesis both qualitatively (visualizing the important input features for the trained classiﬁers 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)

Original: None : contains no square.
Revised: None : contains no square.

Original: Ideally, a classiﬁer should not rely only on the dominant features to predict the class.
Revised: Ideally, a classiﬁer should not rely only on the dominant features to predict the class.

Original: However, as shown in the second row of Figure 2b, 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.
Revised: However, as shown in the second row of Figure 2b, 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.

Original: However, the same network without normalization takes into account other features as well.
Revised: However, the same network without normalization takes into account other features as well.

Original: We verify this by quantifying classiﬁcation accuracy on the different test sets.
Revised: We verify this by quantifying classiﬁcation accuracy on the different test sets.

Original: 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.
Revised: 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.

Original: 4. B ATCH NORMALIZATION IS A CAUSE OF FAILURE WHEN INPUT FEATURES ARE CORRUPTED
Revised: 4. B ATCH NORMALIZATION IS A CAUSE OF FAILURE WHEN INPUT FEATURES ARE CORRUPTED

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: (2019b).
Revised: (2019b).

Original: For simplicity and to directly study the effect of normalization under data shifts we do not use dropout and data augmentations in this experiment.
Revised: For simplicity and to directly study the effect of normalization under data shifts we do not use dropout and data augmentations in this experiment.

Original: We further replace batch normalization with other common normalization techniques such as layer normalization (LN)
Revised: We further replace batch normalization with other common normalization techniques such as layer normalization (LN)

Original: (Ba et al., 2016) and instance normalization (IN) (Ulyanov et al., 2016).
Revised: (Ba et al., 2016) and instance normalization (IN) (Ulyanov et al., 2016).

Original: As shown in Figure 3, the model with no normalization (NN) outperforms all the three types of normalization when evaluated on corrupted data.
Revised: As shown in Figure 3, the model with no normalization (NN) outperforms all the three types of normalization when evaluated on corrupted data.

Original: Whereas, on the clean data all three techniques do generally well.
Revised: Whereas, on the clean data all three techniques do generally well.

Original: With our CT model, corruption errors improves signiﬁcantly on WideResNet and ResNext, by ∼ 4% and ∼ 3% , respectively.
Revised: With our CT model, corruption errors improves signiﬁcantly on WideResNet and ResNext, by ∼ 4% and ∼ 3% , respectively.

Original: wBN
Revised: wBN

Original: NoBN
Revised: NoBN

Original: LN
Revised: LN

Original: IN
Revised: IN

Original: CT (ours)
Revised: CT (ours)

Original: C l a ss i f
Revised: C l a ss i f

Original: i c a t i o n e rr o r
Revised: i c a t i o n e rr o r

Original: Clean error
Revised: Clean error

Original: Curroption error (a) ResNext      wBN
Revised: Curroption error (a) ResNext      wBN

Original: NoBN
Revised: NoBN

Original: LN
Revised: LN

Original: IN
Revised: IN

Original: CT (ours)
Revised: CT (ours)

Original: C l a ss i f
Revised: C l a ss i f

Original: i c a t i o n e rr o r
Revised: i c a t i o n e rr o r

Original: Clean error
Revised: Clean error

Original: Curroption error (
Revised: Curroption error (

Original: [Figure]
Revised: [Figure]

Original: WideResNet
Revised: WideResNet

Original: 4. R EGULARIZING
Revised: 4. R EGULARIZING

Original: Here we compare our counterbalance teacher model with recent approaches on robustness, multidomain generalization, and corrupted point cloud classiﬁcation using different data sets and models.
Revised: Here we compare our counterbalance teacher model with recent approaches on robustness, multidomain generalization, and corrupted point cloud classiﬁcation using different data sets and models.

Original: Robustness to common data corruptions.
Revised: Robustness to common data corruptions.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: However, we show that these approaches work only if BN statistics are adapted to a particular corruption.
Revised: However, we show that these approaches work only if BN statistics are adapted to a particular corruption.

Original: This is unrealistic, especially at inference time, as we often do not have prior knowledge of corruption types that may or may not occur.
Revised: This is unrealistic, especially at inference time, as we often do not have prior knowledge of corruption types that may or may not occur.

Original: 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.
Revised: 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.

Original: CT vs. covariate shift adaptation methods for robustness.
Revised: CT vs. covariate shift adaptation methods for robustness.

Original: We consider three scenarios for adaptation of batch norm statistics; Let C
Revised: We consider three scenarios for adaptation of batch norm statistics; Let C

Original: = { c i | i
Revised: = { c i | i

Original: ∈
Revised: ∈

Original: 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):
Revised: 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):

Original: (I) Adapt-one-test-one : BN statistics are adapted to a batch of samples from each corruption type
Revised: (I) Adapt-one-test-one : BN statistics are adapted to a batch of samples from each corruption type

Original: c i at a time, and classiﬁcation error is computed on the same c i .
Revised: c i at a time, and classiﬁcation error is computed on the same c i .

Original: This is repeated for each corruption type and classiﬁcation error is averaged over each type (Benz et al., 2021; Schneider et al., 2020).
Revised: This is repeated for each corruption type and classiﬁcation error is averaged over each type (Benz et al., 2021; Schneider et al., 2020).

Original: However, it is not trivial to differentiate samples based on corruption type and adapt each separately at test time.
Revised: However, it is not trivial to differentiate samples based on corruption type and adapt each separately at test time.

Original: (II) Adapt-one-test-all : BN statistics are adapted to a batch of samples from a randomly selected c i , and the classiﬁcation error is averaged over C .
Revised: (II) Adapt-one-test-all : BN statistics are adapted to a batch of samples from a randomly selected c i , and the classiﬁcation error is averaged over C .

Original: This experiment reveals the caveats of adaptation on a single corruption, where model fails to generalize to other types of corruptions.
Revised: This experiment reveals the caveats of adaptation on a single corruption, where model fails to generalize to other types of corruptions.

Original: (III) Adapt-all-test-all : BN statistics are adapted to a batch of samples from a random combination of samples from C .
Revised: (III) Adapt-all-test-all : BN statistics are adapted to a batch of samples from a random combination of samples from C .

Original: The goal is to aid adaptation of statistics by looking at samples from “all” corruption types.
Revised: The goal is to aid adaptation of statistics by looking at samples from “all” corruption types.

Original: The classiﬁcation error is then averaged over C .
Revised: The classiﬁcation error is then averaged over C .

Original: This scenario is more realistic than (I) and (II), but still requires a batch of samples from the target domain.
Revised: This scenario is more realistic than (I) and (II), but still requires a batch of samples from the target domain.

Original: In Figure 4, we shed light on a more realistic scenario of batch norm statistics adaptation.
Revised: In Figure 4, we shed light on a more realistic scenario of batch norm statistics adaptation.

Original: We assume access to some extra samples, but not their corruption types (i.e., adapt-all-test-all ).
Revised: We assume access to some extra samples, but not their corruption types (i.e., adapt-all-test-all ).

Original: Surprisingly, we regardless of batch size, such methods perform signiﬁcantly worse on both clean and corruption errors compared to our CT method on CIFAR-10-C with WideResNet model.
Revised: Surprisingly, we regardless of batch size, such methods perform signiﬁcantly worse on both clean and corruption errors compared to our CT method on CIFAR-10-C with WideResNet model.

Original: Further analysis on how adaptation-based methods cause generalization failures are included in Table 6.
Revised: Further analysis on how adaptation-based methods cause generalization failures are included in Table 6.

Original: [Figure]
Revised: [Figure]

Original: SimCLR
Revised: SimCLR

Original: BYOL
Revised: BYOL

Original: Barlow Twins
Revised: Barlow Twins

Original: CT (ours)
Revised: CT (ours)

Original: CIFAR
Revised: CIFAR

Original: [Table]
Revised: [Table]

Original: AllConvNet
Revised: AllConvNet

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table]
Revised: [Table]

Original: WideResNet
Revised: WideResNet

Original: [Table]
Revised: [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: WideResNet
Revised: WideResNet

Original: [Table]
Revised: [Table]

Original: Comparing the robustness of CT to self-supervised methods on common input corruptions.
Revised: Comparing the robustness of CT to self-supervised methods on common input corruptions.

Original: In this experiment, we evaluate how recent self-supervised and contrastive learning methods perform on unseen input corruptions:
Revised: In this experiment, we evaluate how recent self-supervised and contrastive learning methods perform on unseen input corruptions:

Original: SimCLR (Chen et al., 2020), BYOL (Grill et al., 2020), and Barlow
Revised: SimCLR (Chen et al., 2020), BYOL (Grill et al., 2020), and Barlow

Original: Twins (Zbontar et al., 2021).
Revised: Twins (Zbontar et al., 2021).

Original: As reported in Table 1, our method outperforms all self-supervised methods (with ﬁne-tuned encoders) based on the mean corruption error for both CIFAR-10-C and CIFAR-100-C by a large margin.
Revised: As reported in Table 1, our method outperforms all self-supervised methods (with ﬁne-tuned encoders) based on the mean corruption error for both CIFAR-10-C and CIFAR-100-C by a large margin.

Original: For all the methods in Table 1 we use the same individual augmentations (same operations as in AutoAugment) as in other experiments.
Revised: For all the methods in Table 1 we use the same individual augmentations (same operations as in AutoAugment) as in other experiments.

Original: See appendix B. for linear evaluation results.
Revised: See appendix B. for linear evaluation results.

Original: 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 classiﬁcation head on top of the pre-trained encodings.
Revised: 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 classiﬁcation head on top of the pre-trained encodings.

Original: To make it a fair comparison, our modiﬁcation 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 classiﬁcation task.
Revised: To make it a fair comparison, our modiﬁcation 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 classiﬁcation task.

Original: In this way, both the encoder and classiﬁer beneﬁt from labeled supervision.
Revised: In this way, both the encoder and classiﬁer beneﬁt from labeled supervision.

Original: Comparing CT to data augmentation-based methods.
Revised: Comparing CT to data augmentation-based methods.

Original: Hendrycks et al.
Revised: Hendrycks et al.

Original: (2019b) has studied the effectiveness of different data augmentation techniques on making a model robust to common data corruptions.
Revised: (2019b) has studied the effectiveness of different data augmentation techniques on making a model robust to common data corruptions.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: [Table]
Revised: [Table]

Original: CT’s performance in domain generalization.
Revised: CT’s performance in domain generalization.

Original: In this experiment, we evaluate CT on domain generalization using the VLCS benchmark (Torralba and Efros, 2011).
Revised: In this experiment, we evaluate CT on domain generalization using the VLCS benchmark (Torralba and Efros, 2011).

Original: VLCS consists of images from ﬁve object categories shared by the PASCAL VOC 2007, LabelMe, Caltech, and Sun datasets, which are considered to be four separate domains.
Revised: VLCS consists of images from ﬁve object categories shared by the PASCAL VOC 2007, LabelMe, Caltech, and Sun datasets, which are considered to be four separate domains.

Original: 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.
Revised: 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.

Original: We use ResNet-18 as the backbone to make a fair comparison with the state-of-the-art.
Revised: We use ResNet-18 as the backbone to make a fair comparison with the state-of-the-art.

Original: As summarized in Table 3, CT outperforms the state-of-the-art on 3 out of 4 domains and by 1.83% on average.
Revised: As summarized in Table 3, CT outperforms the state-of-the-art on 3 out of 4 domains and by 1.83% on average.

Original: [Table]
Revised: [Table]

Original: CT’s classiﬁcation performance on corrupted 3D point cloud data.
Revised: CT’s classiﬁcation performance on corrupted 3D point cloud data.

Original: In this experiment, we use the RobustPointSet dataset Taghanaki et al.
Revised: In this experiment, we use the RobustPointSet dataset Taghanaki et al.

Original: (2020) which is created for analysis of point classiﬁers in terms of robustness to 3D corruptions.
Revised: (2020) which is created for analysis of point classiﬁers in terms of robustness to 3D corruptions.

Original: We follow the same training-domain validation setting as in Taghanaki et al.
Revised: We follow the same training-domain validation setting as in Taghanaki et al.

Original: 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.
Revised: 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.

Original: We then test the models on the six distorted unseen test sets.
Revised: We then test the models on the six distorted unseen test sets.

Original: This experiment shows the vulnerability of the models trained on original data to unseen input transformations.
Revised: This experiment shows the vulnerability of the models trained on original data to unseen input transformations.

Original: As shown in Table 4, PointNet-CT signiﬁcantly improves classifcation accuracy on most of the shifted test sets, and 1.87% on average compared to original PointNet.
Revised: As shown in Table 4, PointNet-CT signiﬁcantly improves classifcation accuracy on most of the shifted test sets, and 1.87% on average compared to original PointNet.

Original: [Table]
Revised: [Table]

Original:  R ELATED WORK
Revised:  R ELATED WORK

Original: Reducing classiﬁcation error on common input corruptions.
Revised: Reducing classiﬁcation error on common input corruptions.

Original: Hendrycks et al.
Revised: Hendrycks et al.

Original: (2019b) proposed a data augmentation technique which mixes multiple augmented samples using random weights and showed improvements on robust accuracy.
Revised: (2019b) proposed a data augmentation technique which mixes multiple augmented samples using random weights and showed improvements on robust accuracy.

Original: Rusak et al.
Revised: Rusak et al.

Original: (2020) proposed to mix Guassian noise with adversarial samples for robustness to common distortions.
Revised: (2020) proposed to mix Guassian noise with adversarial samples for robustness to common distortions.

Original: However, as in the real world, the space of distortions and their mixture is not ﬁnite, it is not trivial to devise a set of augmentation types that will make a model robust to all distortions.
Revised: However, as in the real world, the space of distortions and their mixture is not ﬁnite, it is not trivial to devise a set of augmentation types that will make a model robust to all distortions.

Original: Moreover, making a model robust to certain corruptions via augmentation does not generalize to others.
Revised: Moreover, making a model robust to certain corruptions via augmentation does not generalize to others.

Original: Dodge and Karam (2017) leverage an ensemble approach for robustness against distortions, but they also assume corruptions are known beforehand.
Revised: Dodge and Karam (2017) leverage an ensemble approach for robustness against distortions, but they also assume corruptions are known beforehand.

Original: Batch normalization’s vulnerability.
Revised: Batch normalization’s vulnerability.

Original: Several works have studied the effect of batch normalization in the context of robustness to adversarial perturbations.
Revised: Several works have studied the effect of batch normalization in the context of robustness to adversarial perturbations.

Original: However, as the most relevant works to ours, Schneider et al.
Revised: However, as the most relevant works to ours, Schneider et al.

Original: and Benz et al.
Revised: and Benz et al.

Original: (2021) suggested improving robustness to common data corruptions by updating the batch normalization parameters using statistics calculated form extra test distribution samples.
Revised: (2021) suggested improving robustness to common data corruptions by updating the batch normalization parameters using statistics calculated form extra test distribution samples.

Original: 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 speciﬁc distortion and tested on the same distortion type.
Revised: 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 speciﬁc distortion and tested on the same distortion type.

Original: Teacher-student models in the context of robustness.
Revised: Teacher-student models in the context of robustness.

Original: 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.
Revised: 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.

Original: In contrast, our CT method is augmentation-independent.
Revised: In contrast, our CT method is augmentation-independent.

Original: Adversarially trained models might not necessarily be robust to common data corruptions, as we showed in Table 2, our CT method performs signiﬁcantly 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).
Revised: Adversarially trained models might not necessarily be robust to common data corruptions, as we showed in Table 2, our CT method performs signiﬁcantly 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).

Original: 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.
Revised: 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.

Original: Similarities to self-supervised methods.
Revised: Similarities to self-supervised methods.

Original: Our CT method is similar to self-supervised models (SSL) (Chen et al., 2020;
Revised: Our CT method is similar to self-supervised models (SSL) (Chen et al., 2020;

Original: Grill et al., 2020;
Revised: Grill et al., 2020;

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Hendrycks et al.
Revised: Hendrycks et al.

Original: (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.
Revised: (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.

Original: Note that our method does not employ any self-supervised loss function.
Revised: Note that our method does not employ any self-supervised loss function.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this work, we investigated the effect of BN on model robustness.
Revised: In this work, we investigated the effect of BN on model robustness.

Original: We provided a theoretical justiﬁcation 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.
Revised: We provided a theoretical justiﬁcation 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.

Original: 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.
Revised: 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.

Original: Empirically, we demonstrated that our method’s learned representations are robust to common corruptions on a suite of domain generalization tasks.
Revised: Empirically, we demonstrated that our method’s learned representations are robust to common corruptions on a suite of domain generalization tasks.

Original: A PPENDIX A A DDITIONAL EXPERIMENTAL DETAILS
Revised: A PPENDIX A A DDITIONAL EXPERIMENTAL DETAILS

Original: We implement our CT variants using Tensorﬂow Abadi et al.
Revised: We implement our CT variants using Tensorﬂow Abadi et al.

Original: A. MNIST MISSING FEATURES EXPERIMENT
Revised: A. MNIST MISSING FEATURES EXPERIMENT

Original: We trained the MNIST models (architecture below) with SGD optimizer (with Tensorﬂow’s default parameters) and batch size of 128 for 100 epochs.
Revised: We trained the MNIST models (architecture below) with SGD optimizer (with Tensorﬂow’s default parameters) and batch size of 128 for 100 epochs.

Original: For the NoBN model, we simply removed the BN layers and trained all models without any data augmentation.
Revised: For the NoBN model, we simply removed the BN layers and trained all models without any data augmentation.

Original: [Table]
Revised: [Table]

Original: A.
Revised: A.

Original: [Table]
Revised: [Table]

Original: C OMPARING CT TO OTHER METHODS ON COMMON INPUT CORRUPTION
Revised: C OMPARING CT TO OTHER METHODS ON COMMON INPUT CORRUPTION

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: A. A PPLYING CT TO POINT CLOUDS
Revised: A. A PPLYING CT TO POINT CLOUDS

Original: 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.
Revised: 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.

Original: A. A PPLYING CT TO DOMAIN GENERALIZATION
Revised: A. A PPLYING CT TO DOMAIN GENERALIZATION

Original: We trained CT with SGD optimizer with learning rate of 0.001 and momentum of 0.9.
Revised: We trained CT with SGD optimizer with learning rate of 0.001 and momentum of 0.9.

Original: We set batch size to 32 and image size to 224x224, and train for 500 epochs.
Revised: We set batch size to 32 and image size to 224x224, and train for 500 epochs.

Original: A. S ELF - SUPERVISED BASELINES
Revised: A. S ELF - SUPERVISED BASELINES

Original: We train the self-supervised methods with batch size 256 for 1000 epochs while maintaining the original hyper-parameters.
Revised: We train the self-supervised methods with batch size 256 for 1000 epochs while maintaining the original hyper-parameters.

Original: We ﬁne-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 l 2 coefﬁcient set to 0.0005 to avoid overﬁtting.
Revised: We ﬁne-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 l 2 coefﬁcient set to 0.0005 to avoid overﬁtting.

Original: B A DDITIONAL EXPERIMENTAL RESULTS B. B ATCH - NORM ADAPTATION ADDITIONAL RESULTS
Revised: B A DDITIONAL EXPERIMENTAL RESULTS B. B ATCH - NORM ADAPTATION ADDITIONAL RESULTS

Original: Batch norm adaptation-based results for robustness to common input corruptions on CIFAR-10-C.
Revised: Batch norm adaptation-based results for robustness to common input corruptions on CIFAR-10-C.

Original: [Table]
Revised: [Table]

Original: L INEAR EVALUATION OF SELF - SUPERVISED MODELS
Revised: L INEAR EVALUATION OF SELF - SUPERVISED MODELS

Original: Self-supervised linear evaluation results where encoders are frozen and only classiﬁcation layer is trained.
Revised: Self-supervised linear evaluation results where encoders are frozen and only classiﬁcation layer is trained.

Original: [Table]
Revised: [Table]

Original: P REDICTIVE PERFORMANCE OF MNIST DATASET VARIABLES .
Revised: P REDICTIVE PERFORMANCE OF MNIST DATASET VARIABLES .

Original: The following analysis is performed to determine whether the dark and light pixels in the MNIST experiment are discriminative.
Revised: The following analysis is performed to determine whether the dark and light pixels in the MNIST experiment are discriminative.

Original: We plot histograms of all pixel values for each class separately.
Revised: We plot histograms of all pixel values for each class separately.

Original: The histograms look very similar, as shown in Figure 5, indicating that the dark or light pixels are not predictive alone.
Revised: The histograms look very similar, as shown in Figure 5, indicating that the dark or light pixels are not predictive alone.

Original: Therefore we add red or blue squares/variables which are both low variance and predictive.
Revised: Therefore we add red or blue squares/variables which are both low variance and predictive.

Original: C A DDITIONAL ANALYSIS ON THE EFFECT OF NORMALIZATION IN OVERPARAMETERIZED REGIME C. S UBTRACTING MEAN DOES NOT CHANGE THE ESTIMATED PARAMETERS
Revised: C A DDITIONAL ANALYSIS ON THE EFFECT OF NORMALIZATION IN OVERPARAMETERIZED REGIME C. S UBTRACTING MEAN DOES NOT CHANGE THE ESTIMATED PARAMETERS

Original: In Section 2.1, we show that dividing each feature by its variance incentivize the model to rely more on features with low variance.
Revised: In Section 2.1, we show that dividing each feature by its variance incentivize the model to rely more on features with low variance.

Original: Here we show that subtracting mean from each feature and target does not change the estimated parameter in overparametrized regime.
Revised: Here we show that subtracting mean from each feature and target does not change the estimated parameter in overparametrized regime.

Original: Recall that X ∈ R n × d denote the training examples and Y ∈ R n denote their target.
Revised: Recall that X ∈ R n × d denote the training examples and Y ∈ R n denote their target.

Original: Let µ X ∈ R
Revised: Let µ X ∈ R

Original: d be the mean of features in training data, and µ Y ∈ R denote the mean of the their targets.
Revised: d be the mean of features in training data, and µ Y ∈ R denote the mean of the their targets.

Original: 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:
Revised: 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:

Original: [Equation]
Revised: [Equation]

Original: C. T HE EFFECT OF NORMALIZATION ON MAX - MARGIN CLASSIFIERS
Revised: C. T HE EFFECT OF NORMALIZATION ON MAX - MARGIN CLASSIFIERS

Original: In Section 2.1, we analyze the effect of normalization in linear regression.
Revised: In Section 2.1, we analyze the effect of normalization in linear regression.

Original: Here we show that the analysis hold for the max-margin classiﬁers as well.
Revised: Here we show that the analysis hold for the max-margin classiﬁers as well.

Original: Soudry et al., Soudry et al.
Revised: Soudry et al., Soudry et al.

Original: (2018), show that without any explicit regularization, gradient descent on logistic loss converges to the L maximum margin separator for all linearly separable datasets.
Revised: (2018), show that without any explicit regularization, gradient descent on logistic loss converges to the L maximum margin separator for all linearly separable datasets.

Original: 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).
Revised: 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).

Original: Without normalization we have:
Revised: Without normalization we have:

Original: [Equation]
Revised: [Equation]

Original: Recall that U is a diagonal matrix where U ii is the standard deviation of the i th feature.
Revised: Recall that U is a diagonal matrix where U ii is the standard deviation of the i th feature.

Original: After normalization we have:
Revised: After normalization we have:

Original: [Equation]
Revised: [Equation]

Original: At the test time we predict sign ( XU − 1 ˆ β ) , substituting U − 1 β with θ we have:
Revised: At the test time we predict sign ( XU − 1 ˆ β ) , substituting U − 1 β with θ we have:

Original: [Equation]
Revised: [Equation]

Original: Similar to linear regression scenario normalizing data change the inductive bias of estimator to choose a model that rely more on low variance features.
Revised: Similar to linear regression scenario normalizing data change the inductive bias of estimator to choose a model that rely more on low variance features.

Original: 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.
Revised: 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.