Original: Distribution Mismatch Correction for Improved Robustness in Deep Neural Networks
Revised: Distribution Mismatch Correction for Improved Robustness in Deep Neural Networks

Original: Anonymous Author(s)

Original: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: Deep neural networks rely heavily on normalization methods to improve theirperformance and learning behavior.
Revised: Deep neural networks rely heavily on normalization methods to improve their performance and learning behavior.

Original: Although normalization methods spurred thedevelopment of increasingly deep and efﬁcient architectures, they also increasedthe vulnerability with respect to noise and input corruptions.
Revised: Although normalization methods spurred the development of increasingly deep and efﬁcient architectures, they also increased the vulnerability with respect to noise and input corruptions.

Original: In most applications,however, noise is ubiquitous and diverse; this can often lead to complete failureof machine learning systems as they fail to cope with mismatches between theinput distribution during training- and test-time.
Revised: In most applications, however, noise is ubiquitous and diverse; this can often lead to complete failure of machine learning systems as they fail to cope with mismatches between the input distribution during training- and test-time.

Original: The most common normalizationmethod, batch normalization, reduces the distribution shift during training but isagnostic to changes of the input distribution during test time.
Revised: The most common normalization method, batch normalization, reduces the distribution shift during training but is agnostic to changes of the input distribution during test time.

Original: Sample-based normalization methods can correct linear transformations of the activation distribution butcannot mitigate changes in the distribution shape; this makes the network vulnerable to distribution changes that cannot be reﬂected in the normalization parameters.
Revised: Sample-based normalization methods can correct linear transformations of the activation distribution but cannot mitigate changes in the distribution shape; this makes the network vulnerable to distribution changes that cannot be reﬂected in the normalization parameters.

Original: We propose an unsupervised non-parametric distribution correction method thatadapts the activation distribution of each layer.
Revised: We propose an unsupervised non-parametric distribution correction method that adapts the activation distribution of each layer.

Original: This reduces the mismatch betweenthe training and test-time distribution by minimizing the 1-D Wasserstein distance.
Revised: This reduces the mismatch between the training and test-time distribution by minimizing the 1-D Wasserstein distance.

Original: In our experiments, we empirically show that the proposed method effectively reduces the impact of intense image corruptions and thus improves the classiﬁcationperformance without the need for retraining or ﬁne-tuning the model.
Revised: In our experiments, we empirically show that the proposed method effectively reduces the impact of intense image corruptions and thus improves the classiﬁcation performance without the need for retraining or ﬁne-tuning the model.

Revised: An extended version of this paper can be found at https://arxiv.org/abs/2110.01955 .

Original:  Introduction
Revised:  Introduction

Original: Early on, Neural Networks (NNs) have proven to excel at interpolating between training data pointsbut to fail when extrapolating to regions not covered by the training data [2, 8].
Revised: Early on, Neural Networks (NNs) have proven to excel at interpolating between training data points but to fail when extrapolating to regions not covered by the training data [2, 8].

Original: The lack of sufﬁcientlylarge datasets, therefore, limited the application of NNs to tasks with well-known input distributions;this prevents any unpredictable behavior in the network that might stem from data with unknowndistribution.
Revised: The lack of sufﬁciently large datasets, therefore, limited the application of NNs to tasks with well-known input distributions; this prevents any unpredictable behavior in the network that might stem from data with unknown distribution.

Original: More precisely – given a model with parameters θ , output y , input x and a staticconditional probability p ( y | x , θ ) – the problem of evaluating samples x from a different distributionthan the training distribution is known as covariate shift [21, 24].
Revised: More precisely – given a model with parameters θ , output y , input x and a static conditional probability p ( y | x , θ ) – the problem of evaluating samples x from a different distribution than the training distribution is known as covariate shift [21, 24].

Original: In recent years, the rise of big dataand data augmentation techniques have alleviated the problem of distribution shifts via increasing thenumber of samples from the input space [13].
Revised: In recent years, the rise of big data and data augmentation techniques have alleviated the problem of distribution shifts via increasing the number of samples from the input space [13].

Original: The problem of covariate shift, however, remained– albeit in a different form: when training Deep Neural Networks (DNNs), each parameter updatecauses a distribution shift for the next mini-batch in the consecutive layers, resulting in convergenceproblems for DNNs.
Revised: The problem of covariate shift, however, remained – albeit in a different form: when training Deep Neural Networks (DNNs), each parameter update causes a distribution shift for the next mini-batch in the consecutive layers, resulting in convergence problems for DNNs.

Original: In particular with DNNs the problem is further exacerbated by a large numberof layers since the distribution shifts can occur internally (within the model) before every layer.
Revised: In particular with DNNs the problem is further exacerbated by a large number of layers since the distribution shifts can occur internally (within the model) before every layer.

Original: Thus,the main impact of the covariate shift moved from test-time to training-time.
Revised: Thus, the main impact of the covariate shift moved from test-time to training-time.

Original: The introduction of Batch normalization (BN) reduced such internal covariate shifts during trainingby matching the distribution of activations across batches and, in doing so, greatly improved theconvergence of deep Convolutional Neural Networks (CNNs) [10].
Revised: The introduction of Batch normalization (BN) reduced such internal covariate shifts during training by matching the distribution of activations across batches and, in doing so, greatly improved the convergence of deep Convolutional Neural Networks (CNNs) [10].

Original: BN, which is the most commonly used normalization method, estimates the normalization parameters, i.e. the mean and variance,based on the training set distributions.
Revised: BN, which is the most commonly used normalization method, estimates the normalization parameters, i.e. the mean and variance, based on the training set distributions.

Original: This makes BN inherently vulnerable to changes that aree.g. caused by image corruptions [3] and forces the models to extrapolate to regions not coveredduring training.
Revised: This makes BN inherently vulnerable to changes that are e.g. caused by image corruptions [3] and forces the models to extrapolate to regions not covered during training.

Original: To improve the corruption robustness of DNNs, data augmentation is used to improverobustness against speciﬁc types of corruption, but this concurrently reduces robustness against othertypes of corruptions [3, 7].
Revised: To improve the corruption robustness of DNNs, data augmentation is used to improve robustness against speciﬁc types of corruption, but this concurrently reduces robustness against other types of corruptions [3, 7].

Original: This highlights the importance of mitigating covariate shifts duringtest-time.
Revised: This highlights the importance of mitigating covariate shifts during test-time.

Original: Group Normalization (GN) and Filter Response Normalization (FRN) have been proposedto overcome the batch size dependence of BN (caused by insufﬁcient statistics for small batches).
Revised: Group Normalization (GN) and Filter Response Normalization (FRN) have been proposed to overcome the batch size dependence of BN (caused by insufﬁcient statistics for small batches).

Original: Both methods are more ﬂexible than BN as they compute the normalization parameter over individualsamples and are thus more robust against changes in the activation distributions.
Revised: Both methods are more ﬂexible than BN as they compute the normalization parameter over individual samples and are thus more robust against changes in the activation distributions.

Original: One importantlimitation of all the above normalization methods is that they can only correct for linear distributiontransformations (e.g. mean shift or variance scaling) but not for mismatches between the shape of thedistribution.
Revised: One important limitation of all the above normalization methods is that they can only correct for linear distribution transformations (e.g. mean shift or variance scaling) but not for mismatches between the shape of the distribution.

Original: We propose a non-parametric distribution correction method that utilizes the 1D-Wasserstein distanceand reduces distribution mismatches of arbitrary form during test-time.
Revised: We propose a non-parametric distribution correction method that utilizes the 1D-Wasserstein distance and reduces distribution mismatches of arbitrary form during test-time.

Original: This correction method canbe combined with other normalization methods and corrects for changes in the distribution shapein an unsupervised setting without the need for retraining or ﬁne-tuning of the models, in contrastto self-supervised methods which adapt at least some of the model parameters [14, 20, 25, 26].
Revised: This correction method can be combined with other normalization methods and corrects for changes in the distribution shape in an unsupervised setting without the need for retraining or ﬁne-tuning of the models, in contrast to self-supervised methods which adapt at least some of the model parameters [14, 20, 25, 26].

Original: Our proposed approach is an iterative procedure following an energy minimization scheme as usedin image denoising [5, 17, 19, 28].
Revised: Our proposed approach is an iterative procedure following an energy minimization scheme as used in image denoising [5, 17, 19, 28].

Original: It is agnostic to the speciﬁc type of noise and maps all noisyactivations to the target distribution of each layer.
Revised: It is agnostic to the speciﬁc type of noise and maps all noisy activations to the target distribution of each layer.

Original: We compare the target to the test-time distributionafter each activation layer and – if necessary – calculate corrections throughout the model.
Revised: We compare the target to the test-time distribution after each activation layer and – if necessary – calculate corrections throughout the model.

Original: Given thatthe target distribution is based on the training data, this effectively moves the test samples closer todata points seen during training and in doing so reduces the covariate shift.
Revised: Given that the target distribution is based on the training data, this effectively moves the test samples closer to data points seen during training and in doing so reduces the covariate shift.

Original: Consequently, the networkcan better process the features in subsequent layers.
Revised: Consequently, the network can better process the features in subsequent layers.

Original: A subsequent step minimizes the difference tothe original activation maps again and ensures that the proposed correction does not induce unwanteddistortions.
Revised: A subsequent step minimizes the difference to the original activation maps again and ensures that the proposed correction does not induce unwanted distortions.

Original: In our experiments, we empirically show that our proposed method improves robustness againsthigh-intensity noise of input corruptions.
Revised: In our experiments, we empirically show that our proposed method improves robustness against high-intensity noise of input corruptions.

Original: We evaluate and compare three normalization methods, i.e., BN, GN, and FRN on corrupted ImageNet (ILSVRC 2012).
Revised: We evaluate and compare three normalization methods, i.e., BN, GN, and FRN on corrupted ImageNet (ILSVRC 2012).

Original: Furthermore, we exemplary analyze theconvergence behavior of the proposed correction.
Revised: Furthermore, we exemplary analyze the convergence behavior of the proposed correction.

Original:  Related work
Revised:  Related work

Original: Many methods – such as weight regularization or dropout – aim to improve robustness in deeplearning models [4, 23].
Revised: Many methods – such as weight regularization or dropout – aim to improve robustness in deep learning models [4, 23].

Original: Alternatively, robustness can be improved by increasing the coverage of theinput space with data augmentation.
Revised: Alternatively, robustness can be improved by increasing the coverage of the input space with data augmentation.

Original: Therefore training samples are augmented by the applicationof afﬁne transformations or expected noise types and are then explicitly included in the trainingset [1, 6, 13, 22].
Revised: Therefore training samples are augmented by the application of afﬁne transformations or expected noise types and are then explicitly included in the training set [1, 6, 13, 22].

Original: But, as mentioned before, improving robustness to one corruption type via dataaugmentation can lead to a decrease of robustness against others [3, 7].
Revised: But, as mentioned before, improving robustness to one corruption type via data augmentation can lead to a decrease of robustness against others [3, 7].

Original: Other ways of improvingrobustness and prediction stability are representation learning techniques or capsule networks.
Revised: Other ways of improving robustness and prediction stability are representation learning techniques or capsule networks.

Original: Theseapproaches try to learn equivariant representations of features; i.e., conceptual representationsindependent of the position, orientation or context [11, 12, 18, 27].
Revised: These approaches try to learn equivariant representations of features; i.e., conceptual representations independent of the position, orientation or context [11, 12, 18, 27].

Original: Moreover, the choice of activationfunctions impacts the robustness of the network as well [15, 16].
Revised: Moreover, the choice of activation functions impacts the robustness of the network as well [15, 16].

Original: Recently, there has also beenincreasing interest in improving the robustness of normalization methods.
Revised: Recently, there has also been increasing interest in improving the robustness of normalization methods.

Original: The inﬂuence of inputcorruptions on networks using batch normalization has been investigated in [3]; this further ledto a domain adaption method for the normalization that improves the robustness of DNNs againstcorruptions.
Revised: The inﬂuence of input corruptions on networks using batch normalization has been investigated in [3]; this further led to a domain adaption method for the normalization that improves the robustness of DNNs against corruptions.

Original: A similar approach was taken by [20], which adapts the normalization parameters basedon test-time statistics.
Revised: A similar approach was taken by [20], which adapts the normalization parameters based on test-time statistics.

Original: Other approaches use self-supervised methods to retrain the model based ontest samples to reduce domain shifts [14, 25, 26].
Revised: Other approaches use self-supervised methods to retrain the model based on test samples to reduce domain shifts [14, 25, 26].

Original:  Distribution correction for Deep Neural Networks (DNNs)
Revised:  Distribution correction for Deep Neural Networks (DNNs)

Original: Machine learning and signal processing tasks frequently experience performance degradation causedby noise.
Revised: Machine learning and signal processing tasks frequently experience performance degradation caused by noise.

Original: As noise comes in miscellaneous forms, it is challenging to achieve general robustnessagainst arbitrary noise.
Revised: As noise comes in miscellaneous forms, it is challenging to achieve general robustness against arbitrary noise.

Original: Most existing normalization approaches use parametric approximations (e.g.
Revised: Most existing normalization approaches use parametric approximations (e.g.

Original: Gaussians), and can therefore only reﬂect changes in the distribution parameters (i.e., the mean and variance).
Revised: Gaussians), and can therefore only reﬂect changes in the distribution parameters (i.e., the mean and variance).

Original: Therefore, they often struggle in practice as they cannot correct for shape mismatches inthe activation distribution.
Revised: Therefore, they often struggle in practice as they cannot correct for shape mismatches in the activation distribution.

Original: In order to mitigate the distribution mismatches of the test-time activationsa , we must ﬁnd an effective way to suppress noisy activations while maintaining the classiﬁcationperformance in the subsequent layers.
Revised: In order to mitigate the distribution mismatches of the test-time activations a , we must ﬁnd an effective way to suppress noisy activations while maintaining the classiﬁcation performance in the subsequent layers.

Original: Therefore, we will formulate this problem as a probabilisticdenoising problem.
Revised: Therefore, we will formulate this problem as a probabilistic denoising problem.

Original: For our considerations, we recast the maximum a-posteriori problem into anequivalent energy minimization problem to simplify the optimization procedure.
Revised: For our considerations, we recast the maximum a-posteriori problem into an equivalent energy minimization problem to simplify the optimization procedure.

Original: The energy E (˜ a | a )is composed of two terms
Revised: The energy E (˜ a | a ) is composed of two terms

Original: R
Revised: R

Original: (˜ a ) (corresponding to the prior) and D ( a | ˜ a ) (corresponding to thelikelihood)
Revised: (˜ a ) (corresponding to the prior) and D ( a | ˜ a ) (corresponding to the likelihood)

Original: so that
Revised: so that

Original: [Equation]
Revised: [Equation]

Original: where ˜ a represents the corrected activations.
Revised: where ˜ a represents the corrected activations.

Original: This type of problem can be solved iteratively alternatingbetween prior and likelihood term.
Revised: This type of problem can be solved iteratively alternating between prior and likelihood term.

Original: Note that we choose independent values for the parameters ofthe prior-update, i.e., λ 1 , and of the likelihood-update, i.e., λ 2 , that implicitly determine the relativeimportance of the corresponding terms.
Revised: Note that we choose independent values for the parameters of the prior-update, i.e., λ 1 , and of the likelihood-update, i.e., λ 2 , that implicitly determine the relative importance of the corresponding terms.

Original: For this form, we must, at the one hand, specify a suitableprior term R (˜ a ) that reduces the covariate shift in each layer without restricting the network.
Revised: For this form, we must, at the one hand, specify a suitable prior term R (˜ a ) that reduces the covariate shift in each layer without restricting the network.

Original: The datalikelihood term D ( a | ˜ a ) , on the other hand, preserves the spatial correlations of the activation mapsand prevents the loss of valuable information.
Revised: The data likelihood term D ( a | ˜ a ) , on the other hand, preserves the spatial correlations of the activation maps and prevents the loss of valuable information.

Original: By minimizing both terms jointly, one can achieve anoptimal trade-off between minimizing the covariate shift and representing the available data.
Revised: By minimizing both terms jointly, one can achieve an optimal trade-off between minimizing the covariate shift and representing the available data.

Original: Furtherimplementation details are presented in the pseudocode in Appendix B.
Revised: Further implementation details are presented in the pseudocode in Appendix B.

Original: Typically, parametric distributions do not provide good representations of the activation distributionsin DNNs.
Revised: Typically, parametric distributions do not provide good representations of the activation distributions in DNNs.

Original: Consequently, any correction method that approximates the prior by a parametric targetdistribution q θ ( a ) distorts the shape of the true activations distribution p ( a ) .
Revised: Consequently, any correction method that approximates the prior by a parametric target distribution q θ ( a ) distorts the shape of the true activations distribution p ( a ) .

Original: Therefore, to reducedistortions, q ( a ) must be non-parametric as well.
Revised: Therefore, to reduce distortions, q ( a ) must be non-parametric as well.

Original: The Wasserstein distance proves to be particularlywell-suited for our correction method: not only does it allow us to effectively minimize the mismatchbetween the distributions during training- and test-time, but it also provides an elegant way ofrepresenting a non-parametric distribution in one dimension.
Revised: The Wasserstein distance proves to be particularly well-suited for our correction method: not only does it allow us to effectively minimize the mismatch between the distributions during training- and test-time, but it also provides an elegant way of representing a non-parametric distribution in one dimension.

Original: To ﬁnd a well-suited prior distribution, we must ﬁrst sort the N activations a for each sample m inthe training set in ascending order.
Revised: To ﬁnd a well-suited target distribution, we must ﬁrst sort the N activations a for each sample m in the training set in ascending order.

Original: We ﬂatten the channels and create a single distribution across theheight H , width W , and channel C dimension of the layer resulting in N = H × W × C activationvalues.
Revised: We ﬂatten the channels and create a single distribution across the height H , width W , and channel C dimension of the layer resulting in N = H × W × C activation values.

Original: Second, to minimize the non-parametric prior term, we need to calculate the Wassersteindistance between the activations distribution during test-time (i.e., p ( a ) ) and the target distributionq ( t ) obtained from the training set.
Revised: Second, to minimize the non-parametric prior term, we need to calculate the Wasserstein distance between the activations distribution during test-time (i.e., p ( a ) ) and the target distribution q ( t ) obtained from the training set.

Original: The one-dimensional Wasserstein distance between q ( t ) and p ( a ( m ) ) is given according to
Revised: The one-dimensional Wasserstein distance between q ( t ) and p ( a ( m ) ) is given according to

Original: [Equation]
Revised: [Equation]

Original: where t ( i ) are the sorted target values calculated using the Wasserstein barycenter of the training setand a (cid:48) ( m ) ( i ) are the sorted test-time activations.
Revised: where t ( i ) are the sorted target values calculated using the Wasserstein barycenter of the training set and a (cid:48) ( m ) ( i ) are the sorted test-time activations.

Original: Here we subtract the mean of the activations, as we donot care about the precise location of the distribution but primarily about its shape.
Revised: Here we subtract the mean of the activations, as we do not care about the precise location of the distribution but primarily about its shape.

Original: Let r = 1 ; then,the Wasserstein distance between p ( a ) and q ( t ) from (2) is minimized by updating the (unsorted)activation with index j according to ∆ j = t ( i ) − a ( i ) .
Revised: Let r = 1 ; then, the Wasserstein distance between p ( a ) and q ( t ) from (2) is minimized by updating the (unsorted) activation with index j according to ∆ j = t ( i ) − a ( i ) .

Original: We apply the correction after the ReLUactivation function; thus many activations are zero.
Revised: We apply the correction after the ReLU activation function; thus many activations are zero.

Original: Our updates must preserve this sparsity as theperformance will degrade otherwise.
Revised: Our updates must preserve this sparsity as the performance will degrade otherwise.

Original: Therefore, we explicitly enforce sparsity in the prior term R ,i.e., we prevent correcting activations with a = 0 by adding an inﬁnitely deep energy-well − δ ( a ) ,where δ ( · ) denotes the Dirac delta.
Revised: Therefore, we explicitly enforce sparsity in the prior term R , i.e., we prevent correcting activations with a = 0 by adding an inﬁnitely deep energy-well − δ ( a ) , where δ ( · ) denotes the Dirac delta.

Original: 1 Combining this sparsity term with the Wasserstein distanceﬁnally leads to the following prior term:
Revised: 1 Combining this sparsity term with the Wasserstein distance ﬁnally leads to the following prior term:

Original: [Equation]
Revised: [Equation]

Original: Minimizing only the prior term might have undesired side effects and destroy important structure inthe channels of the network, i.e the spatial correlation.
Revised: Minimizing only the prior term might have undesired side effects and destroy important structure in the channels of the network, i.e the spatial correlation.

Original: Therefore, the energy minimization needs toﬁnd a trade-off between matching the distributions and conserving the spatial correlation.
Revised: Therefore, the energy minimization needs to ﬁnd a trade-off between matching the distributions and conserving the spatial correlation.

Original: We achievethis by considering a likelihood term D ( a | ˜ a ) , modeled by
Revised: We achieve this by considering a likelihood term D ( a | ˜ a ) , modeled by

Original: [Equation]
Revised: [Equation]

Original: that conserves the structure in the data.
Revised: that conserves the structure in the data.

Original: This expression is also straighforward to minimize.
Revised: This expression is also straighforward to minimize.

Original:  Experiments
Revised:  Experiments

Original: In our experiments we ﬁrst analyze the convergence behavior of the correction algorithm on corrupted
Revised: In our experiments we ﬁrst analyze the convergence behavior of the correction algorithm on corrupted

Original: MNIST and present results for the corrupted classiﬁcation ILSVRC 2012 dataset (ImageNet-C) usingthree different normalization methods (BN, FRN, GN) for the models.
Revised: MNIST and present results for the corrupted classiﬁcation ILSVRC 2012 dataset (ImageNet-C) using three different normalization methods (BN, FRN, GN) for the models.

Original: Convergence behavior As our correction method is applied iteratively [Figure] we investigate its convergence behavior for the average classiﬁcation accuracy on the corrupted MNIST dataset.
Revised: Convergence behavior As our correction method is applied iteratively, we investigate its convergence behavior for the average classiﬁcation accuracy on the corrupted MNIST dataset.

Original: Here we varythe step size parameters λ 1 and λ 2 and show their inﬂuence.
Revised: Here we vary the step size parameters λ 1 and λ 2 and show their inﬂuence.

Original: In Figure 1 we see that BN, FRN
Revised: In Figure 1 we see that BN, FRN

Original: [Figure]
Revised: [Figure]

Original: and
Revised: and

Original: [Figure]
Revised: [Figure]

Original: without using the correction (at iteration 0).

Original: Also, BN is theonly model that can substantially improve its classiﬁcation performance by running the algorithm formore than 1 iteration.
Revised: Also, BN is the only model that can substantially improve its classiﬁcation performance by running the algorithm for more than 1 iteration.

Original: FRN shows consistent results after one iteration, whereas the performancesdiverge for more iterations.
Revised: FRN shows consistent results after one iteration, whereas the performances diverge for more iterations.

Original: GN also shows improvement within the second iteration for the parameterset λ 1 = 0 .
Revised: GN also shows improvement within the second iteration for the parameter set λ 1 = 0 .

Original: 25 , λ 2 = 0 .
Revised: 25 , λ 2 = 0 .

Original: 5 , but does not outperform the best model obtained with only one iteration.
Revised: 5 , but does not outperform the best model obtained with only one iteration.

Original: Corrupted ImageNet classiﬁcation The corrupted ILSVRC 2012 dataset contains 19 differentcorruption variants of the original dataset.
Revised: Corrupted ImageNet classiﬁcation The corrupted ILSVRC 2012 dataset contains 19 different corruption variants of the original dataset.

Original: It contains 50000 RGB images of size 224 × 224 percorruption that we used for our evaluations.
Revised: It contains 50000 RGB images of size 224 × 224 per corruption that we used for our evaluations.

Original: For the evaluation on the corrupted ImageNet (ILSVRC2012) dataset we again choose the same parameter set λ 1 = 0 .
Revised: For the evaluation on the corrupted ImageNet (ILSVRC 2012) dataset we again choose the same parameter set λ 1 = 0 .

Original: 5 and λ 2 = 0 .
Revised: 5 and λ 2 = 0 .

Original: 5 and N iter = 1 for allnormalization methods.
Revised: 5 and N iter = 1 for all normalization methods.

Original: We trained a single ResNet-50 model for 90 epochs using an SGD optimizerwith a base learning rate of 0.1 on the clean ImageNet dataset.
Revised: We trained a single ResNet-50 model for 90 epochs using an SGD optimizer with a base learning rate of 0.1 on the clean ImageNet dataset.

Original: The learning rate was decayed after 30,60, and 80 epochs, and a warm-up from 0.02 to 0.1 was used during the ﬁrst ﬁve epochs of training.
Revised: The learning rate was decayed after 30, 60, and 80 epochs, and a warm-up from 0.02 to 0.1 was used during the ﬁrst ﬁve epochs of training.

Original: [Table]
Revised: [Table]

Original: The results in Table 4 show that for the corrupted ImageNet dataset, all models regardless of theused norm achieve about 1% performance improvement by using the proposed correction method.
Revised: The results in Table 4 show that for the corrupted ImageNet dataset, all models regardless of the used norm achieve about 1% performance improvement by using the proposed correction method.

Original: Generally, we see that the more ﬂexible GN models are the most robust against image corruptionsoutperforming BN and FRN by > 3.0%.
Revised: Generally, we see that the more ﬂexible GN models are the most robust against image corruptions outperforming BN and FRN by > 3.0%.

Original: This supports our assumption that even for large amounts ofdata, covariate shift still inﬂuences performance for corrupted inputs.
Revised: This supports our assumption that even for large amounts of data, covariate shift still inﬂuences performance for corrupted inputs.

Original:  Conclusion and Outlook
Revised:  Conclusion and Outlook

Original: We proposed a non-parametric activation distribution correction method based on the Wassersteindistance.
Revised: We proposed a non-parametric activation distribution correction method based on the Wasserstein distance.

Original: It reduces the mismatch between test-time and training distributions of the activations within
Revised: It reduces the mismatch between test-time and training distributions of the activations within

Original: DNNs.
Revised: DNNs.

Original: The proposed method uses a maximum a-posteriori estimate, determined by minimizing theenergy with respect to a data likelihood term and a prior term based on the Wasserstein distance.
Revised: The proposed method uses a maximum a-posteriori estimate, determined by minimizing the energy with respect to a data likelihood term and a prior term based on the Wasserstein distance.

Original: Our proposed method works in an unsupervised setting and can be retroﬁtted into existing networkswithout retraining.
Revised: Our proposed method works in an unsupervised setting and can be retroﬁtted into existing networks without retraining.

Original: In our experiments, we showed that our correction algorithm can effectively reducethe mismatch between test-time and training distributions.
Revised: In our experiments, we showed that our correction algorithm can effectively reduce the mismatch between test-time and training distributions.

Original: This improves classiﬁcation performanceon corrupted input data.
Revised: This improves classiﬁcation performance on corrupted input data.

Original: For future applications, we want to extend our robustness evaluations andexplore the use of our algorithm for reducing the impact of parametric approximations.
Revised: For future applications, we want to extend our robustness evaluations and explore the use of our algorithm for reducing the impact of parametric approximations.

Original: A Analyzing the distributions
Revised: A Analyzing the distributions

Original: Since we assume a quasi-stationary distribution of the activations of the ﬂattened layer across thetraining set, we need to verify if this assumption holds for our ResNet models.
Revised: Since we assume a quasi-stationary distribution of the activations of the ﬂattened layer across the training set, we need to verify if this assumption holds for our ResNet models.

Original: This is necessary inorder to create a meaningful target distribution q ( t ) .
Revised: This is necessary in order to create a meaningful target distribution q ( t ) .

Original: Therefore, we analyze the variance σ 2 of thetarget values t over the training set:
Revised: Therefore, we analyze the variance σ 2 of the target values t over the training set:

Original: [Equation]
Revised: [Equation]

Original: In Figure 2, we see that although all normalization methods have a similar shape for the variance of their target value distribution, BN experiences a much higher variation at the tails of the distributions,having a value range 5 times higher than GN and FRN.
Revised: In Figure 2, we see that although all normalization methods have a similar shape for the variance of their target value distribution, BN experiences a much higher variation at the tails of the distributions, having a value range 5 times higher than GN and FRN.

Original: This observation is in line with the intuitionthat more ﬂexible normalization methods produce narrower activation distributions.
Revised: This observation is in line with the intuition that more ﬂexible normalization methods produce narrower activation distributions.

Original: Generally, we see that with the exception of the distribution-tails, the distribution of the targets ismore or less stationary.
Revised: Generally, we see that with the exception of the distribution-tails, the distribution of the targets is more or less stationary.

Original: This is an interesting observation, as one would expect that the activationdistributions for the entire layer have a higher dependency on the speciﬁc input.
Revised: This is an interesting observation, as one would expect that the activation distributions for the entire layer have a higher dependency on the speciﬁc input.

Original: Following thisobservation, we analyze this behavior by comparing the activation distributions for the most dissimilarexamples in the test set.
Revised: Following this observation, we analyze this behavior by comparing the activation distributions for the most dissimilar examples in the test set.

Original: Therefore we calculate the Wasserstein distances from a test example k to allother test examples m with respect to their channel distributions, utilizing the order statistics for eachchannel j separately and select the sample with the largest total value with index ˜ m ,˜
Revised: Therefore we calculate the Wasserstein distances from a test example k to all other test examples m with respect to their channel distributions, utilizing the order statistics for each channel j separately and select the sample with the largest total value with index ˜ m , [Equation] [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: Figure 3 shows that while the individual channel distributions are very distinct, the shape of thedistribution over all channels is more or less stable.
Revised: Figure 3 shows that while the individual channel distributions are very distinct, the shape of the distribution over all channels is more or less stable.

Original: This indicates that the inverse correlation ofchannels can conserve the overall shape of the distribution for each layer.
Revised: This indicates that the inverse correlation of channels can conserve the overall shape of the distribution for each layer.

Original: B Algorithm
Revised: B Algorithm

Original: The algorithm is described in the following pseudo-code.
Revised: The algorithm is described in the following pseudo-code.

Original: The algorithm is executed for every layerand returns the corrected activations ˜ a .
Revised: The algorithm is executed for every layer and returns the corrected activations ˜ a .

Original: The step-sizes λ 

Revised: The step-sizes λ

Original: and λ 2 determine the relative importance ofthe prior term versus the likelihood term.
Revised: and λ 2 determine the relative importance of the prior term versus the likelihood term.

Original: Algorithm
Revised: Algorithm

Original: Input: Unsorted activations a , sorted target values t , step-sizes λ 1 , λ 2 , and maximum number of iterations N iter
Revised: Input: Unsorted activations a , sorted target values t , step-sizes λ 1 , λ 2 , and maximum number of iterations N iter

Original: Output: Corrected activations ˜
Revised: Output: Corrected activations ˜

Original: a
Revised: a

Original: N ← len( a ) ˜ a ← a − 1 N (cid:80) i a i for n < N iter do [ a ( i ) ] , j ← sort( ˜ a ) for i <
Revised: N ← len( a ) ˜ a ← a − 1 N (cid:80) i a i for n < N iter do [ a ( i ) ] , j ← sort( ˜ a ) for i <

Original: N do j ← j i ˜ a j ← ˜ a j + λ 1 · ( t ( i ) − a ( i ) )
Revised: N do j ← j i ˜ a j ← ˜ a j + λ 1 · ( t ( i ) − a ( i ) )

Original: ˜ a j ← ˜ a j + λ 2 · ( a j − ˜ a j ) end for end for for j < N do if a j (cid:54) = 

Revised: ˜ a j ← ˜ a j + λ 2 · ( a j − ˜ a j ) end for end for for j < N do if a j (cid:54) =

Original: then ˜ a j ← ˜ a j else ˜ a j ← a j end if end for
Revised: then ˜ a j ← ˜ a j else ˜ a j ← a j end if end for

Original: C Supplemental material of experiments
Revised: C Supplemental material of experiments

Original: [Table]
Revised: [Table]