MC Layer Normalization for calibrated uncertainty in Deep Learning
Authors: Thomas Frick
Diego Antognini
Ioana Giurgiu
Benjamin F Grewe
Cristiano Malossi
Rong J.B. Zhu
Mattia Rigotti
Venue: Accepted by TMLR ()

Abstract: Efficiently estimating the uncertainty of neural network predictions has become an increasingly important challenge as machine learning models are adopted for high-stakes industrial applications where shifts in data distribution may occur. Thus, calibrated prediction uncertainty is crucial to determine when to trust a model's output and when to discard them as implausible. We propose a novel deep learning module - MC Layer Normalization - that acts as a drop-in replacement for Layer Normalization blocks and endows a neural network with uncertainty estimation capabilities. Our method is motivated from an approximate Bayesian perspective, but it is simple to deploy with no significant computational overhead thanks to an efficient one-shot approximation of Monte Carlo integration at prediction time. To evaluate the effectiveness of our module, we conduct experiments in two distinct settings. First, we investigate its potential to replace existing methods such as MC-Dropout and Prediction-Time Batch Normalization. Second, we explore its suitability for use cases where such conventional modules are either unsuitable or sub-optimal for certain tasks (as is the case with modules based on Batch Normalization, which is incompatible for instance with transformers). We empirically demonstrate the competitiveness of our module in terms of prediction accuracy and uncertainty calibration on established out-of-distribution image classification benchmarks, as well as its flexibility by applying it on tasks and architectures where previous methods are unsuitable.

Published in Transactions on Machine Learning Research (02/2024)
MC Layer Normalization for calibrated uncertainty in Deep Learning
Thomas Frick IBM Research, ETH Zurich
fri@zuirch.ibm
Diego Antognini IBM Research
diego. btagg..nnttmmiiii@immm.comm
Ioana Giurgiu
igi@zarichhiüm. www.trrric.ibb.c..
IBM Research
Benjamin Grewe ETH Zurich
bgrewe@ethz.ck
Cristiano Malossi IBM Research
acm@zurich. ibm...immooooo. com
Rong J.B. Zhu
rongzhu@fudan.ccu edw.cc
Pudan University
Mattia Rigotti IBM Research
mrg@zurich.ihm www.rricch......mmmmmmmmmm com
Reviewed OpenReview: https: https:///peerreviau net/ forum? id= 6G3IC13E0C
openreview.
Abstract
Efficiently estimating the uncertainty of neural network predictions has become increas- ingly important challenge machine learning models adopted for high-stakes industrial applications where shifts in data distribution may occur. Thus, calibrated prediction uncer- tainty is crucial to determine when to trust model's output and when to discard them as implausible. We propose a novel deep lcarning module MC Layer Normalization that acts as a drop-in replacement for Layer Normalization blocks and endows a neural network with uncertainty estimation capabilities. Our method motivated from an approximate Bayesian perspective, but it simple to deploy with no significant computational over- head thanks efficient onc-shot approximation of Monte Carlo integration at prediction lime. To evaluate the elfectiveness of our module, conduct experiments in two distinct settings. First, we investigate its potential to replace existing methods such as MC-Dropout and Prediction-Time Batch Normalization. Second, we explore its suitability for use cases where such conventional modules are either unsuitable or sub-optimal for certain tasks (as the case with modules based on Batch Normalization, which incompatible for instance with transformers). We empirically demonstrate the competitiveness of our module in terms of prediction accuracy and uncertainty calibration on established out-of-distriibtiion image classification benchmarks, well as its flexibility by applying it tasks and architectures where previous methods are unsuitable. Code implementing our MC-LayerNorm module can be found here https://githun wwwp:ittuu...oo..eee httpuau.coo////mmmaaaeeeaaee

---

Published in Transactions on Machine Learning Research (02/2024)
1 Introduction
Endowing neural networks with a mechanism for efficient estimation of prediction uncertainly is a challenging problem that is acquiring increasing attention as these models are deployed in critical decision making scttings. In high-stakes real-world applications such autonomous driving (Bojarski et al. 2016), robotics (Sinderhauf et al., 2018), medical diagnosis (Djuric et al. 2017 Esteva et al. 2017] models are often operating in an out-of-distribution situation compared to the training data. In such scenarios, calibrated prediction uncertainty is crucial to meaningfully compare competing predictions and decide when to trust them or when to reject them as implausible.
Bayesian methods such as Bayesian Deep Neural Networks (DNNs) offer a principled formalism to compute prediction uncertainty (MacKay 1992). Their disadvantage is that obtaining uncertainty measures over complex models such as large neural networks quickly becomes intractable because of the computational challenge of estimating and updating posteriors over their parameters. To overcome these issues, researchers have recently proposod approximate Bayesian methods that make use of variational approximations by conveniently leveraging sampling mechanisms that are inherently present in modern DNNs such as Dropout and Batch Normalization. MC-Dropout (Gal & Ghahramani] 2016), for instance cleverly exploits the fact that Dropout noise can be interpreted as sampling mechanism over a variational distribution approximating the posterior of the DNNs parameters given the training data. Following this work, the stochastic nature of the mini-batch sampling process exposed by Batch Normalization (BatchNorm) layers has also been used to perform approximate Bayesian inference for uncertainty estimation (Teye et al. and domain-adaptive prediction-time calibration of uncertainty (Nado et al. 2021).
2018 Mukhoti et al. 2020).
In this paper, we propose a new deep learning module that fits neatly into the line of research on sampling- based deep learning layers for estimating prediction uncertainty in neural networks. In particular, our module can be used as a drop-in replacement for Layer Normalization to scamlessly add uncertainty calibration capabilities to neural network. Our module, named MC Layer Normalization (MC-LayerNorm) consists of stochastic variant of Layer Normalization (LayerNorm) that subsamples features when computing the normalization statistics used to normalize the input features. While it offers complementary functionality to related modules such as MC-Dropout and Prediction-Tinee Batch Norm, it can also be used in training settings and with architectures where the latter two modules are not technically applicable or result in sub- optimal performance. In addition, MC-LayerNorm inherits the advantages of LayerNorm over BatchNorm, including the fact that its training behavior does not depend on the mini-batch size and that it invariant to single input data re-scaling. This last property is especially appealing as it allows MC-LayerNorm to perform zero-shot domain adaptation at prediction time.
We ralustrate how MC-LayerNorm can be theoretically motivated from an approximate Bayesian perspective while being easy to deploy in practice without imposing significant computational overhead. We then demon- strate the effectiveness of our module through empirical evaluation in two distinct settings: First, we show its competitive performance in terms of prediction accuracy and uncertainty calibration compared to state-of- the-art alternatives. To that end, we conduct experiments on established benchmarks for out-of-listributtio (OOD) image classification with convolutional networks. Secondly, we highlight the applicability of MC- LayerNorm in scenarios where competing methods cannot be used due to incompatibility or suboptimality with the architectures at hand. To this end, we present experiments using the Vision Transformers the same OOD image classification benchmarks as in the first setting. As the Vision Transformer architecture prohibits the use of Batch Norm, we show the superiority of our method compared lo existing uncertainty cal- ibration methods that can be applied. In addition, we investigate the use case of click- thickhhhrough rate prediction on the Critco datasct, where the current state-of-tho-art model does not include Dropout or BarchNorm lay- ers but does include LayerNorm blocks. As a result, the use of MC-Dropout or Prediction Time Norch is impossible, while MC-LayerNorm is natural drop-in replacement.
2

---

Published in Transactions on Machine Learning Research (02/2024)
2
Related Work
Our proposed normalization layer is very much inspired by Monte Carlo Batch Normalization and the related work in approximate Bayesian inference in deep learning such as MC-Dropout, as well as the prediction-ting adaptive normalization method Prediction-Tiie
Batch Normalization.
Momte Carlo Dropout (MC-Dropout, Gal & Ghahramani [016];; Gal et al. 2017 approximates Bayesian
inference in neural networks for uncertainty estimates using dropout. Usually, Dropout sets a portion of the input features to zero at training time while not dropping any input features at test time. MC-Dropout introduces stochasticity at inference time by using the training behavior at test time. Similarly, Monte Carlo Batch Normalization (MCBN, Teye et al. (2018), Mukhoti et al. (2020)) introduces an alternative based on the properties of the batch normalization statistics leveraging training-time batch statistics test time. Instead of utilizing the running batch statistics at inference time, they propose to approximate Bayesian inference by rumning multiple forward passes with distinet sets of training-tinn time normalization statistics. Prediction-Time Batch Normalization is proposed by Nado et al. (2021) as countermeasure to the covariate shift that comes with out-of-distribution data samples. The method works by discarding the running batch statistics of batch normalization layers and instead uses the batch statistics of each separate batch at test time. Consequently, this effectively counteracts the covariate shift and significantly improves model calibration.
Masksembles (Durasov et al. 2021 relies on fixed number of binary masks instead of randomly dropping parts of the network as in MC-Dropout. The masks are parameterized in a way that allows to change corre- lations between individual sub-models by controlling the overlap between the masks. This leads to a simple and easy to implement method with performance on par with Ensembles at a fraction of the cost. Tempera- ture Sealing (Guo et al. 2017: Platt et al.) 1999) is a post-hoc method which improves calibration after the initial training of the model by Luning sofumax temperature parameter on the validation/ validation/calibratioo ccalibration set.
2023 proposed a different approach to incorporate Bayesian viewpoint into neural network training by deriving training objective which includes a regularization term that corresponds to performing function-space Empirical Bayes estimation. Contrary lo the previous methods this approach requires to change the training objective, and therefore departs from our goal of designing lightweight method that can be implemented as a simple drop-in replacement that otherwise leaves the training and prediction-time procedure unchanged.
Recently, Rudner et al.
A different line of related works, which also aim at efficiently estimating uncertainty of neural network pre- dictions, but completely depart from a Bayesian viewpoint, are frequentist methods which include conformal prediction methods (Vovk et al. 2005), which have recently been successfully applied to deep learning set- tings (Bates et al. 2021 Zhu & Rigotti, 2021 Angelopoulos et al. 2022). They have also been developed in the direction of accommodating covariate shift, making them interesting for OOD prediction Tibshirani
et al. 2019
3 MC Layer Normalization
The main contribution of this paper is to propose a new normalization module that, similar to Monte Carlo Batch Normalization (MCBN, Teye et al. (2018), Mukhoti et al. (2020)), also allows for uncertainty estimation, in particular through Monte Carlo sampling over a source of randomness. Crucially, instead of sampling randomness originating from the stochasticity of mini-batcha such as in MCBN, we propose to use a variant of Layer Normalization where we inject stochasticity by subsampling features when computing the normalization statistics used to normalize the feature vector. Here we detail this procedure by first summarizing the original Layer Normalization module using a similar notation as the original paper [Ba ltt al. 2016).
We consider the 1-th hidden layer in feed-forward neural network, and let a' be the vector representation of the summed inputs to the neurons (preactivations) in that layer. These preactivations are computed through
3

---

Published in Transactions on Machine Learning Research (02/2024)
matrix-vector multiplication of the weight matrix W (wly) and the inputs to the layer ht' as:
dt = Σ wiggt
1 s(a: + b/),
(1)
where f()) is element-wise activation function such as ReLU and if is a bias parameter.
Layer Normalization (Ba et al. 2016) was propose as method to mitigate the covariate shift of correlated inputs as an alternative to Batch Normalization (Toffe & Szegedy 2015), and consists in normalizing the hidden units a given layer as a = (a: - n²) /o' using the mean j and the variance of of their preactivations computed as follows:
Σ 0
(od)² = Ni Σ a -d)"
(2)
N
where N₁ is the number of units in layer I. One of the practical advantages of Layer Normalization over Batch Normalization is that it does not rely on any assumption about the size of the mini-batch and can therefore be used with batch size of 1.
Our MC Layer Normalization (MC-LayerNorm) layer consists in modifying eq. (2 by running the averages
only over a random subset of units. In particular, we define a set S C [N.] (I....,NI) obtained by sampling a fixed fraction f of preactivations (or, equivalently, by dropping them with drop rate 1 - f), and compute the normalization statistics over the sampled units as follows:
Σ d VES:
(g²)²
1
Σ iES!
(at TT)²
(3)
S
S
where [Si] denotes the size of the set St.
As result, the normalized preactivations computed by MC-LayerNorm and outputs of layer are:
f(a) क = /o! f(a b!).
(4)
Notice that jr' and of are random variables whose realization is determined by the subset of randomly sampled indices St, and are at and the outputs hi+1. This observation motivates the following probabilistic view of architectures endowed with Nayerryrroo
Probabilistic view of MC-LayerNoom.
Assuming without loss of generality supervised learning setting on training dataset D ((x,, 1:) we formulate the goal of supervised learning training parameters 0 (which include the weights W and biases if of each layer 1) to maximizing the likelihood of the dataset D under the predictive probability
Pe (v(p)r) softmax(fo(z)).
(5)
where fo(-) is parametrized as neural network (see e.g. Gal & Ghahramani 2015
A network (z) endowed with MC-LayerNorm layers, can itself be modeled as a distribution of networks fo[a] (7t,, {+),)"))) a)),, where (7',5"), the set of realizations of the random statistics in eq. (3) in all layers I. As we show below, an equivalent way of seeing this to think sampled network fo(a ( J() for a given set (7',5'), as a corresponding network fa(x), where now is sampled from a distribution ô go (0) defined appropriately. In the next section show that the distribution go (θ) converges in specific sense towards Gaussian distribution around the parameters 0.
Approximate normality of MC-LayerNorm networks.
Our main theoretical result follows directly from:
4

---

Published in Transactions on Machine Learning Research (02/2024)
Theorem 3.1. Given network fo(x) with MC-LayerNorm layers, denote by fo(z) corresponding network with all MC-LayerNorm (eq. 3 replaced by regular LayerNorm (eq. (2) Replacing LayerNorm
in fo(x) with MC-LayerNorm induces distribution of models fa(a)) with ~ qo In addition, go(0) is asymptotically normal oround the parameters of fo(x).
Proof. The proof of theorem 3.1 is sketched in appendix A.3
Next, we show how to use this result to approximate Bayesian inference with deep neural networks.
Approximate Bayesian inference with MC-LayerNorm.
The theorem 3.1 allows us to leverage the methods developed by Gal & Ghahramani 2016), who use stochas- tic regularization techniques to perform practical approximate inference in the space of neural networks fa (-) by virtue of the fact that they introducc a random variable that allows for sampling over networks. Specifi- cally, the idea starts from the goal of computing the predictive distribution plylæ, D) over outputs y given a new input æ and the training dataset D using Bayesian Model Averaging (Hoeting et al. 1999 Wilson & Iz- mailov 2020), but by replacing the intractable posterior p(0[D) over parameters with tractable variational approximation (0), then marginalizing over A using Monte Carlo integration: Pa. (y/z) with 0n ~ q (0).
plylor, D)-
= po(y)=)p(01D)d0 po(ulz)) (0)
N
In practice, in mini-batch Stochastic Grading Descent (SGD) training setting, for each training point and each MC-LayerNorm layer sample subset S to while computing the SGD step. This will implement the process of optimizing (0) towards (#).
At prediction time, what eq- (6) says is to implement Monte Carlo integration by running the network forward pass times given an imput data x, each time with different realizations of the subsets St The obtained outputs are then averaged obtain the final prediction for input
There is however additional approximation of the inference process for a network with MC-LayerNorm that allows to compute predictions even more efficiently by exploiting the second part of theorem 3.1 Because, according to the theorem the distribution 0 ~ (0) is approximately normal around the parameters of a corresponding network with Layer Normalization, we can approximate the Monte Carlo integration in eq. 66 with:
N Pa. (y|x) with On ~ -0(0) sofumax(Ja, (x)) ze softmax(FF(z)).
N
N Σ
(7)
n=1
where we used oq. (5) and Laplace's approximation basod on the asymptotic normality of (0) around 0 for the model [o( with regular LayerNorm instead of MC-LayerNorm. What this suggests that prediction time cheap further onc-shot approximation of the MC integration in eq. (6). We will refer to this modality of use of MC-LayerNorm at prediction time as One-shot approzimation.
simply run the model by replacing MC-LayerNorm with regular LayerNorm
We summarize the previous results in pseudocode snippets detailing the use of MC-LayerNonn in practice for training neural networks with SGD with backpropagation (algorithm D and prodiction time (algorithm 2 Notice that as for regular LayerNorm, in the case of convolutional inputs MC-LayerNorn normalizes its inputs both channel and spatial dimensions.
5

---

Published in Transactions on Machine Learning Research (02/2024)
Algorithm MC-LayerNorm module (training mode)
Algorithm 2 MC-LayerNorn module (eval mode)
Input: input vector a = (ai,.... as...., .a)
Input: input vector
Input: (optional) number of MC samples Ne
(a1,...., Oir...,... aN)
if mc integration then
Input: (parameter) fraction of sampled units f
// MC integration:
for n = 1 to N. do a MC_ KC_LayerrVrm(()
Sample: Select S C N] with
\// Run Algorithm 1}
JJ indices at random
end for Output: ä= 1 N.
Compute:
Ex a"
11 == TST Eies a;
else
Compute:
// One-shot approximation: Compute: μ N Σ di
8² 18 Eies (a: (1)²
Output: a = (a),.... ds....,aN ,NN
Compute: 0² Di= (a, (a. u)) 1)²
with a. "
Output: (01..... asy..... .....).,)))).) with ā. 2:22
end if
4 Results
We empirically evaluate MC-LayerNorm on a suite of established classification benchmarks. In addition to assessing the accuracy of the predictions, we verify the improved uncertainty calibration characteristics of MC-LayerNorm. We focus in particular on the distributions differ due to covariate shift. (Shimodaira) 2000), i.e. when the marginal distributions of features are different, Ptrain
(OOD) setting where training and test
(z), but conditional label distributions are preserved Ptrain (t/z) Ptest (z) The expectation is then that properly calibrated uncertainly will be reflected in predictions that tend to be less confident on the OOD imputs that are not well represented in the training distribution.
We conduct experiments for the two distinct application settings mentioned above: Firstly, we investigate MC-LayerNorm's potential as a viable alternative to existing methods and ils application to real-world scenarios with the sparse occurrence of distifiiist samples. Secondly, we explore MC-LayerNorm an option for accurate uncertainty estimates in applications where competing methods cannot be used they are incompatible with the architectures (e.g., if no BatchNorm layers are present in the model, can't apply MC- NoC-Ba.chNorm).
Following previous work (Ovadia et al. 2019 Nado et al. 2021), we quantify the calibration of the uncertainly of predictions using two measures: expected calibration error (ECE) and Brier score.
ECE (lower is better) quantifics the difference between the confidence of a model and its accuracy computed on bins of samples sorted by confidence (Naeini et al., 2015). Coucretely, we group our predictions into bins B, according to confidence, compute the average prediction acc(B,) within bins, then compute ECE- Σ; N |ace(B;) - (..). As an artifact due confidence binning, ECE can over-emphasize the tails of the probabilities (Quinonero-Co Canirrr---Caaaaaa.aa al. 2006), which is why the literature typically also monitors another calibration metric like the Brier score.
Brier score (lower is better) is the squared distance between the vector of predicted probabilities and the onc-hot encoded true labels (Brier] 1950). It is guaranteed to be a proper scoring rule, i.e. it decreases monotonically to zero as the predicted probabilities approach the true targets distribution (Gneiting & Raftery 2007). Brier score plots will be presented in the Appendix.
Image Classification
We follow the benchmarking methodology in Nado al. (2021)) model calibration metrics (accuracy, ECE, and Brier score) are evaluated on CIFAR-10-C, TinylmageNet-C. C, and ImageNet-C introduced in Hendrycks & Dietterich 2019) (CCA4.0 license). Both datasets corrupted versions of their original test sets (CIFAR-

---

Published in Transactions on Machine Learning Research (02/2024)
10 Krizhevsky 2009 TinylmageNet (Le & Yang 2017); and ImageNet Dengg et al. 2009 applying one of 15 corruptions (e-g., frost, motion blur, rain).
generated by
We assess model calibration for three different scenarios:
1. In-distribution (Test): The original test set of the corresponding dataset is used as baseline for all metrics.
2. Out-of-distribution (OOD):
The metrics are evaluated on the corrupted C-variant test sets (severity 5). This favors Prodiction-Tiie Time BatchNorm as the batches exclusively consist of out-of-distributie samples. We use batch size of 128 (following the insight from Nado el. al. (2021) that shows that Prediction- Time BatchNorm performance degrades only for batch sizes below 128).
3. Zero-shot prediction-time domain adaptation (Mix): This scenario simulates the situation where a deployed model suddenly encounters OOD samples and still hasn't gathered enough observations to re-adapt (as Prediction- Time BatchNorm needs to do on multiple OOD samples). We call this "Mix" emphasize that the models are trained on in-distribution data but lested on OOD samples. Concretely, we create batches of size N consisting of N in-distribution samples and a single distiiiiin sample (we the same batch size in the OOD setting N 128). The calibration metrics then measured only the out-of-distribution samples.
All experiments consist of fine-tuning pretrained model with the respective uncertainty calibration method applied. We train 3 models for each normalization method we select the configuration with the best average accuracy on the validation set. No data augmentations are applicd during training other than a horizontal flip of the images.
configuration (see Appendix A.4). Then for each
0.8
0.6
accuracy
accuracy
0.6
BerchMorm
layerform
0.4
LayerNorm
Masioemale
Masksembe
0.4
Predictionnl BN
Prediction BN
MC-NN
MC-3N
Mc-D-opputt
MC-Dropout
MC-NN MC Approx Joun)
MC-LN MC Apoms (Ouns)
MC LN One Shot (Gurg)
0.0
LN Onc Shot (Durs)
Test
00D
Mix
Test
OOD
Mix
0.8
0.6
ece
0.4
0.2
0.2
0.0
Test
0OD
Mix
Test
0OD
Mix
(a) CIFAR-10-C
(b) TinyImageNet-C
Figure CIFAR-10-C and TinylmageNet-CC with ConvNext: Calibration for in-distribution (Test), out-of-distribbutio (OOD), and zero-shot prediction-time domain adaptation (Mix) sctting: We compare LayerNorm, BatchNorm, Prodiction-Time BatchNorm, MC-Dropout (f=0.9, 10 MC iterations), MC-BatchNorm (10 MC iterations), Masksembles, well as our novel method, MC-LayerNorm (MC and One-Shot approximation, f=0.7) on the expected calibration error (ECE) and accuracy. While Prediction-Tine BatchNorm is superior in the full out-of-distribution (OOD) setting, MC-LayerNorn out- performs all other methods in the zero-shot prediction-time domain adaptation (Mix) use case. Error bars are rucste from evaluation runs for the 15 corruptions applied the original test. set and over three training runs.

---

Published in Transactions on Machine Learning Research (02/2024)
0.9
accuracy
0.8
accuracy
0.6
0.7
0.4
0.6
0.5
MC4N Ore-shot (Dura)
layefform
0.2
LayarNorm
WC-X One-Sner (Duns)
15ce mg
13ca
MCNN 7:5ca ing [Duns)
MC-U + T-Scaiing (Oun))
0.0
Test
0OD
Mix
Test
OOD
Mix
0.4
0.6
ece
ece
0.4
0.2
0.2
0.0
0.0
Test
0OD
Mix
Test
OOD
Mix
(a) CIFAR-10-C
(b) IinyImageNet-C
Figure 2: Temperature Scaling ConvNext for CIFAR-10-C and InngmeageNet-C: Calibration for in-distribution (Test), out-of-distribution (OOD), and zero-shot prediction-time domain adaptation (Mix) setting: We compare LayerNorm and our novel method, MC-LayerNorm (MC and Onc-Shot approximation, f-0.7) with combination of cach respective method and Temperature Scaling on the expected calibration error (ECE) and accuracy. The Temperature scaled versions outperform their respective non-post-hoc calibrated methods. Error bars are constructed from evaluation runs for the 15 corruptions applied to the original test set and over three training runs.
Out-Of-Distribution Image Classification ConvNext
We evaluate the first sctting by empirically comparing MC Layer Normalization (MC integration and One-shot approximation mode) with Prediction-Time Batch Normalization introduced in Nado al. (2021)
Additionally, we report results for MC-Dropout (Gal & Ghahramani 2016 Gal et al. 2017 NorBatchNorm
(Teye et al. 2018 Mukhoti et al.) 2020), and Masksembles (Durasov et al., 2021). For both MC-Dropout
and Masksembles we use the optimal parameter configuration mentioned Durasov el. al. (202T). Finally, we also investigate the combination MC-TayerNorm with post-hoc Temperature Scaling (Guo et al. 2017 Platt al. 1999).
We focus evaluation the ConvNext architecture Liu et al. 2022 architecture for which both BatchNorm and LayerNorm normalization has been shown to reach similar image classification performance. This equality of classification accuracy leads to fair comparison of uncertainty calibration methods independent of the used normalization layer. For CIFAR-10-C, we limit ourselves to the pre-trained architecture size "pico" from the excellent timm library (Wightman et al. 2023 (Apache-2.0 licensc), while for IingInageNet-C, we make use of the "tiny" variant.
it is the only neural network
Figure I tables and [ shows results for CIFAR-10-C and TinyImageNet-C. While MC-LayerNorm per- forms slightly worse (CIFAR-10) or similarly (TinylmageNee) in terms of calibration (ECE) compared to Prediction-Tine BatchNorm in the OOD sctting, it outperforms all other methods in the zero-shot prediction- time domain adaptation (Mix) setting. Prediction-Time BatchNorm loses much of its real-world (Mix) per- formance because relies heavily on batches consisting entirely of LayerNorm methods do not roly on batch statistics as they calculate normalization statistics across channels and spatial dimensions. As result, our method's calibration abilities are independent of batch size and work for small but also for mixed in- and out-of-distribution batches.
samples. In contrast,
8

---

Published in Transactions on Machine Learning Research (02/2024)
1.0
0.8
0.8
Masksemble
0.6
Masksemble
accuracy
LayerNorm
accuracy
LayerNorm
0.6
MC- Dropout
MC-Dropout
0.4
MC-LN MC Approx. (Ours)
MC-LN MC Approx. (Ours)
0.4
MC-LN One-Shot (Ours)
MC-LN One-Shot (Ours)
0.2
T-Scaling
T-Scaling
MC-LN T-Scaling (Ours)
T-Scaling (Ours)
0.0
Test
OOD
Test
COD
0.6
0.5
ece
0.4
0.4
0.3
0.2
0.2
0.1
0.0
Test
OOD
Test
OOD
(a) CIFAR-10-C
(b) Tiny TiyInngeNNet-C
Figure 3: CIFAR-10-C and TinyImageNee--C with Vision Transformer: Calibration for in- distribution (Test) and out-of-distributiin (OOD): We compare LayerNorm, MC-Dropout (f=0.9, MC iterations), MC-BatchNorm (10 MC iterations). Masksembles, Temperature Scaling as well as our novel method MC-LayerNorm (MC and One-Shot approximation, f=0.6) on the expected calibration error (ECE) and accuracy. MC-LayerNorm in combination with Temperature scaling outperforms all other methods. Error bars constructed from evaluation runs for the 15 corruptions applied to the original test set and over three training runs.
Figure 5 and 6 illustrate calibration performance comparisons for varying fractions of sampled features of the MC-LayerNorm. The results show that MC-LayerNorm is stable with respect to its only hyper- parameter, the fraction of subsampled features. The model performance and calibration remain stable even for coufigurations where the normalization is calculated from only 60% of the input features.
Next, we evaluate the necessity of enhancing approximate Bayesian methods with post-hoc calibration, Figure shows the calibration performance of LayerNorm and MC-LayerNorm well the combination of these methods with post-hoc Temperature Scaling. We observe that the Temperature scaled versions surpass their respective post-hoc calibrated counterparts. Temperature scaling alone shows strong effect it equals (CIFAR-10) or outperforms (Tinyllmageeet) MC-LayerNorm for both the OOD and Mix setting, However, the combined effect of Temperature Scaling and MC-LayerNorm is noticeably more robust and results a better overall calibration, indicating practically interesting synergistic effects between the two methods. This pattern is also evident in experiments with Vision Transformers, as shown in Figure
We approximate Bayesian estimation methods (e.g., MC-LayerNorm) and post-hoc Temperature Scaling complementary and synergistic. For one, post-hoc calibration is not always possible, it requires held-out in-distribution calibration datasct, which might not always be available. In such a case, one must necessarily rely on approaches like approximate Bayesian estimation methods. Second, post-hoe calibration only works the base model is already at least partially calibrated to the extent that the miscalibration pattern consistent across the range of logits. The better the basc model calibrated, the more effective post-hoc calibration methods will be. Here, believe that approximate Bayesian models can certainly be crucial to enable further calibration through post-hoc methods.

---

Published in Transactions on Machine Learning Research (02/2024)
0.20
0.15
AUC
0.75
0.10
0.70
0.05
LayerNorm
MC-LayerNom (Ours)
0.00
Test
Test
2
shift
shift
Figure 4: Critco dataset experiment. Area under the Curve (AUC) and expected calibration error (ECE) for four shift intensities: vanilla LayerNorm is compared to MC-LayerNorm (f-0.7) while applying increasing amounts noise to the test set. MC-LayerNorm shows a clear advantage, retaining more of the original model performance as the data shifts further out-of-distribution (plots show average over 10 runs; error bars indicate standard errors around the mean).
Rut-Of-Distin Of-Ditiiiiittionn Image Classification Vision Transformer
We evaluate setting 2 by rumning experiments using the Vision Transformer (Dosovitskiy el al. 2021 archi- tecture (ViT). As discussed in more detail in Shen et al. (2020), batch norm is not used for ViIs for variety of historical and performance reasons. Therefore, endowing them with MC-BatchNorm or Prediction Time BatchNorm would require a relatively radical architectural alteration which will typically enforce change in the training procedure (specifically, to take into account the effect of mini-batch size on the training with BatchNorm), and compromise on the performance as state-of-the-ars is empirically achieved without a BatchNorm layer. Thus MC-LayerNorm is the only norm layer based uncertainty calibration method that is applicable to this architecture. We use the pre-trained architecture implementations from timm (Wightman et al.] 2023) (Apache-2.0 license) and use the size "vit tiny" and "vit small" for CIFAR10 and TinyIm- ageNet/ ImageNet respectively. We compare our method's performance (in MC integration and One-shot approximation mode) versus MC-Dropout (Gal & Ghahramani 2016 Gal et al. 2017), Temperature Scaling (Guo al.] 2017 Platt et al.] 1999) and Masksembles (Durasov et al. 2021). For both MC-Dropout and Masksembles we use the optimal parameter configuration mentioned Durasov et al. (2021). Additionally, we also report results for combination of Temperature Scaling and MC-LayerNNorn.
Figure 3 tables 3 and table 4 shows the superiority of our method when evaluated in terms of uncertainty calibration and accuracy for the Test and the OOD setting. As non of the evaluated methods are batch de- pendent, unlike the BatchNorm methods from the previous experiment, we forgo evaluation on the zero-shot prediction-time domain adaptation use case in this section. MC-LayerNorm outperforms classical Layer- Norm, MC-Dropout and Maskscmbles in terms of ECE and accuracy on both sctting. While Temperature Scaling outperforms MC-LayerNorm, we can show that combination of MC-LayyeNorm and post-hoc Tem- perature scaling reaches maximum calibration performance. These results consistent with the additional experiments in Appendix A.2.1
Training VITs with the Masksemble approach was unfortunately difficult as non of the models we trained converged to a reasonable classification performance. We hypothesize that ViT token width is too narrow for this method to converge as it splits the token into subsets with S overlap.
7b shows calibration performance comparisons for varying fractions of sampled features of the MC-LayerNorm.
Figure
Criteo Display Advertising Challenge
Here, we showcase the value our method in another situation where the competing methods cannot be used as they are incompatible with the architectures. For the experiments, we train models on the Criteo Display Advertising Challenge (CritcoLabs) 2014). We leverage the current state-of-the-art model, MaskNet (Wang al.) 2021). relying on the implementation from Zhu (002)), Apache-2.0 license). This
10

---

Published in Transactions on Machine Learning Research (02/2024)
architecture does not include any BatchNorm layers, meaning that it cannot be naturally extended with PPreddiionnnTime BatchNorm without modifying the training procedure (specifically, to take into account the effect of mini-batch size on the training with BatchNorm), and without incurring a loss in performance (as SOTA is empirically achieved without a BatchNorm layer). MaskNet however includes LayerNorm modules, meaning that we can straightforwardly replace them by dropping in MC-LayerNNorm without changing anything else in the architecture or training. The Criteo challenge is therefore also meant as use case to emphasize the advantages of our approach also beyond vision tasks, in particular its flexibility drop-in replacement layer for LayerNorm (a normalization layer which is almost ubiquitous in modern deep learning) and its ease of use since it does not require any other modifications to the models architecture. We train two types of models: the original MaskNet model containing vanilla LayerNorm and a patched model for which we replace all LayerNorm layers with MC-LayerNorm. At prediction time, we run the MC- LayerNorm models using the One-shot approzimation, making them computationally as cheap as the regular MaskNet models to evaluate. The trained models are evaluated on the original test set and a corrupted version of the test set. We simulate out-of-distribution data by applying corruption in the form of Gaussian noise on the numerical features. Shift intensities [6.25% 12.5%, 25%, 50%] are based on units of the standard deviation of the original feature values. Finally, measure the difference in model performance between the original models with LayerNorm and the adjusted models with MC-LayerNorm.
Figure shows results on the Criteo dataset demonstrating a clear advantage of MC-LayerNorn when it come the considered metrics. Here we displays the performance of MC-LayerNorm for f=0.7 (the curves are robust within the interval f=0.6-0.8, with a gradual progressive degradation as is varied above or below this range). The Figure shows that as we shift. the tests set further out-of-distribution, our method retains more of the original calibration and model performance compared to the baseline of the traditional LayerNorm.
5 Conclusion and Discussion
We presented novel normalization layer, MC-LayerNorm, which improves model calibration for in- and out-of-distribution (OOD) data by injecting stochasticity through subsampling features when computing the normalization statisties. Our method can be easily deployed without significant computational overhead by replacing the LayerNorm blocks in an already trained model and then fine-tuning it. We empirically demonstrated the application of MC-LayerNorm for two distinct scttings: First, we have shown its potential as an allernative to existing methods for OOD classification and its superiority compared to Batch Norm based alternatives when it comes to mixed in- and -iitr--iiittiittionn batches. Secondly, we presented experiments for use cases where application of conventional methods are either unsuitable or technically infeasible. In this setting, MC-LayerNorn was shown to outperform alternatives on OOD classification using Vision Transformers and on the CTR Criteo dataset using MaskNet.
In other words, MC-LayerNorm is able to perform zero-shot prediction-time domain adaptation, which gives it an advantage in real-world OOD use cases where a deployed model is being exposed for the first time to new OOD samples and has not had time to properly adapt to the covariate shift.
Limitations and Future Work. In addition to the C-variants of the considered datasets we focused on in this study, future work could validate our method on their respective A-variants, which consist of real-world adversarial examples. An additional direction for future work is MC-LayerNorm's potential impact on other model performance metrics, such as the stability of training convergence. For example, during experimen- tation, we noticed improved training convergence toward the same performance metrics (reduced variability across runs with different hyper-parameter settings like different learning rate values) when employing MC- LayerNorm.
In terms of broader impact, as alrcady discussed, being able to provide deep neural networks with calibrated prediction uncertainly would make their meployment potentially safer and trusted, particularly in high-stakes applications where it is crucial to be able to reliably take consequential decisions based on them. Therefore, we argue that the line of rescarch to which our work contributes is an important step towards trying to systenatically mitigate potential negative socictal effects due to the involuntary misuse of deep learning technologies.
11

---

Published in Transactions on Machine Learning Research (02/2024)
References
Anastasios Angelopoulos, Stephen Bates, Jitendra Malik, and Michael I. Jordan. Uncertainty Sets for Image Classifiers using Conformal Prediction, September 2022.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer Normalization, July 2016.
Stephen Bates, Anastasios Angelopoulos, Lihua Lei, Jitendra Malik, and Michael Jordan. Distribution-free, Risk-controlling Prediction Sets. Journal of the ACM, 68(6):43:1 43:34, September 2021. ISSN 0004-5411. doi: 10.1145/3477 3478535.
Mariusz Bojarski, Davide Del Testa, Daniel Dworakowski, Bernhard Firner, Beat Flepp, Prasoon Goyal, Lawrence D. Jackel, Mathew Monfort, Urs Muller, and Jiakai Zhang. End to end learning for self-driving cars. arxiv preprint or 1xii:1604.007366 2016.
Glenn W. Brier. Verification of forecasts expressed in terms probability. Monthly weather review, 78(1):
1-3, 1950.
CriteoLabs. Display Advertising C http:////./../..///////////////r 2014.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Ti Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp. 248 255, 2009. doi: D.1109/CVVR.2209522666
Ugljesa Djurie, Gelareh Zadeh, Kenneth Aldape, and Phedias Diamandis. Precision histology: How deep learning is poised to revitalize histomorphology for personalized cancer care. NPJ precision oncology, 1 (1):1-5, 2017.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby, An image is worth 16x16 words: Transformers for image recognition at scale, 2021.
Nikita Durasov, Timur Bagautdinov, Pierre Baque, and Pascal Fua. Masksembles for Uncertainty Estima- tion, June 2021.
Andre Esteva, Brett Kuprel, Roberto A. Novoa, Justin Ko, Susan M. Swetter, Helen M. Blau, and Sebastian Thrun. Dermatologist-level classification of skin cancer with deep neural networks. nature, 542(7639): 115-118, 2017.
Yarin Gal and Zoubin Ghahramani. Bayesian convolutional neural networks with Bernoulli approximate variational inference. arXiu preprint arXiv:1506.021 2015.
Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainly deep learning. In International Conference on Machine Learning, pp. 1050-1059. PMLR, 2016.
Yarin Gal, Riashat, Islam, and Zoubin Ghahramani. Deep bayesian active learning with image data. Tn International Conference on Machine Learning, pp. 1183 1192. PMLR, 2017.
Tilmann Gneiting and Adrian F.. Raftery. Strictly proper scoring rules, prediction, and estimation. Journal of the American statistioal Association, 102(477):359 378, 2007.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In International Conference on Machine Learning, pp. 1321-1330. PMLR, 2017.
Jaroslav Hájek. Limiting distributions in simple random sampling from finite population. Publications of the Mathematical Institute of the Hungarian Academy of Sciences, 5:361-374, 1960.
Dan Hendrycks and Thomas G. Dietterich. Benchmarking Neural Network Robustness to Common Cor- ruptions and Surface Variations, April 2019. Comment: Supersoded by Benchmarking Neural Network Robustness Common Corruptions and Perturbations. arXiv:1903.12261
12

---

Published in Transactions on Machine Learning Research (02/2024)
Jennifer A. Hoeting, David Madigan, Adrian E. Raftery, and Chris T. Volinsky. Bayesian model averaging: A tutorial (with comments by M. Clyde, David Draper and EI George, and rejoinder by the authors. Statistical science, 14(4):382 417, 1999.
Sergey Ioffe and Christian Szegody. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning, pp. 448-456. PMLR, 2015.
Alex Krizhevsky. Learning Multiple Layers of Features from Tiny Images. 2009.
Ya Le and Xuan Yang. Tiny ImageNet Visual Recognition Challenge. 2017.
Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Sain-
ing Xie. A ConvNet for the 2020s, March 2022. Comment: CVPR 2022; Code:
http://lhu...oommmeeookreeeeee NeX..
David JC MacKay. A practical Bayesian framework for backpropagation networks. Newral computation, 4 (3):448-47) 472, 1992.
Jishnu Mukhoti, Puncet K. Dokania, Philip HS Torr, and Yarin Gal. On Batch Normalisation for Approxi- mate Bayesian Inference. arxio preprint urXiv:2012. 13220, 2020.
Zachary Nado, Shreyas Padhy, D. Sculley, Alexander D' D'Amour, Balaji Lakshminarayanan, and Jasper Snoek. Evaluating Prediction- Time Batch Normalization for Robustness under Covariate Shift, Jamuary 2021.
Mahdi Pakdaman Naeini, Gregory F. Cooper, and Milos Hauskrecht. Obtaining Well Calibrated Probabil- ities Using Bayesian Binning. Proceedings of the ... AAAI Conference Artificial Intelligence. AAAI Conference on Artificial Intelligence, 2015:2901-2900 January 2015. ISSN 2159-5399.
Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, David Sculley, Sebastian Nowozin, Joshua Dillon, Balaji Takshminarayanan, and Jasper Snock. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. Advances in neural information processing systems, 32, 2019.
John Platt et al. Probabilistic outputs for support vector machines and comparisons to regularized likelihood methods. Advances in large margin classifiers, 10(3):61-74 1999.
Geoff Pleiss. Cpleiss/temperatuure tetpreeepprrrtturressaannnnn January 2024.
Joaquin Quinonero-Candee Carl Edward Rasmussen, Fabian Sinz, Olivier Bousquet, and Bernhard Scholkopf. Evaluating predictive uncertainty challenge. Lecture Notes in Computer Science, 3944:1-27, 2006.
Tim GJ Rudner, Sanyam Kapoor, Shikai Qiu, and Andrew Gordon Wilson Function-space regularization neural networks: probabilistic perspective. In International Conference on Machine Learning, pp. PMLR, 2023.
Sheng Shen, Zhewei Yao, Amir Gholami, Michael W. Mahoney, and Kurt Keutzer. Powernorm: Rethinking batch normalization in transformers, 2020.
Hidetoshi Shimodaira. Improving predictive inference under covariate shift by weighting the log- likeliiold function. Journal of Statistical Planning and Inference, 90(2)2227-244, October 2000. ISSN 0378-3758.
doi:
Niko Sünderhauf, Oliver Brock, Walter Scheirer, Raia Hadsell, Dieter Fox, Jürgen Leitner, Ben Uperoft, Pieter Abbeel, Wolfram Burgard, and Michael Milford. The limits and potentials of deep learning for robotics. The International journal of robotics research, 37(4-5):40--220, 2018.
Mattias Teye, Hossein Azizpour, and Kevin Smith. Bayesian uncertainty estimation for batch normalized deep networks. In International Conference Machine Learning, pp. 4907-4916. PMLR, 2018.
13

---

Published in Transactions on Machine Learning Research (02/2024)
Ryan J Tibshirani, Rina Foygel Barber, Emmanuel Candes, and Aaditya Ramdas. Conformal Prediction Un- der Covariate Shift. In Advances in Newral Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
A. Vaart. Asymptotic Stotistics. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 1998.
Vladimir Vovk Alexander Gammerman, and Glenn Shafer. Algorithmic Learning in a Random World. Beenner-Veerlg, New York, 2005. ISBN 978-0-387-00152-4. doi: 10.1007/b106715.
Zhiqiang Wang, Qingyun She, and Junlin Zhang. MaskNet: Introducing Feature-Win Multiplication to
CTR Ranking Models by Instance-Guided Mask, July 2021. Comment: In Proceedings of DLP-KDD 2021. Singapggeo.r.. arXiv admin note: text overlap with arXiv:2006.12753
Ross Wightman, Nathan Raw, Alexander Soare, Aman Arora, Chris Ha, Christoph Reich, Jakub Kacz- marzyk, mrT23, Mike, SeeFin, contrastive, Mohammed Rizin, Hyeongehan Kim, Csaba Kertész, Dushyant Mehta, Guillem Cucurull, Kushajveer Singh, han, Yuki Tatsunami, Andrew Lavin, Juntang Zhuang, Matthijs Hollemans, Mohamed Rashad, Sepehr Sameni, Vyacheslav Shults, Lucain, Xiao Wang, Yonghye Kwon, Yusuke Uchida, and Zhun Zhong. Rwightman frighhtman//yttorr-imng-nooooodo V0.8.6dev0 Release. Zenodo, January 2023.
Andrew G. Wilson and Pavel Izmailov. Bayesian deep learning and probabilistic perspective of general- ization. Advances in neural information processing systems, 33:4697-4708, 2020.
Jamie Zhu. FuxiCTR: A configurable, tunable, and reproducible library for CTR prediction. hwapn/uuepaaaapeee pai/FuxiCTR, January 2023.
R. Zhu and M. Rigotti. Deep Bandits Show-Off: Simple and Efficient Exploration with Deep Networks. Advances in Newral Information Processing Systems (NeurIPS), 35, 2021.
14

---

Published in Transactions on Machine Learning Research (02/2024)
A
Appendix
A.1
Additional results
A.1.1
Comparison of MC-LayerNorm drop rates for ConvNext
accuracy
MC-LayerNorm f=0.9
0.6
MC-LayerNorm f=0.8
MC-LayerNorm f=0.7
0.4
MC-LayerNorm f=0.6
Test
OOD
Mix
1.0
brier
0.5
Test
0OD
Mix
0.4
0.2
0.0
Test
OOD
Figure 5: CIFAR-10-C with ConvNext:
Comparison of calibration between different MC-LayerNorm (One-Shot approximation) subsampling fractions for in-distribution (Test), out-of-distribation (OOD) and the zero-shot prediction-ting sitietion-time domain adaptation (Mix) setting.
15

---

Published in Transactions on Machine Learning Research (02/2024)
0.75
accuracy
MC-LayerNorm f=0.9
0.50
MC-LayerNorm f=0.8
0.25
MC-LayerNorm
MC-LayerNorm f=0.6
0.00
Test
1.0
brier
0.5
Test
0OD
Mix
0.3
ece
0.2
0.1
0.0
Test
OOD
Figure 6: TinylmageNet--C with ConvNext: Comparison of calibration between different MC- LayerNorm (Onc-Shot approximation) subsampling fractions for in-distribution (Test), out-of-disttibutiio (OOD) and the zero-shol prediction-lime domain adaptation (Mix) setting.
16

---

Published in Transactions on Machine Learning Research (02/2024)
A.2
Comparison of MC- MC-LayerNorm drop rates for Vision Transformers
1.0
0.8
0.6
accuracy
LayerNorm
0.4
LayerNorm
MC-LayerNorm 0.9
MC-LayerNorm f=0.9
MC-LayerNorm 0=0.8
0.4
MC-LayenNorm f-0.8
0.2
MC-LayerNorm f=0.7
MC-LayerNorm
MC-LayerNerm f-0.6
0.2
MC-LayerNorm f-0.6
Test
OOD
Test
OOD
1.4
1.2
1.00
1.0
brier
brier
0.8
0.50
0.6
0.25
0.4
0.00
0.2
Test
0OD
0.6
0.6
0.5
0.4
0.4
0.3
0.3
0.2
0.2
0.1
0.1
0.0
Test
OOD
Test
0OD
(a) CLFAR-10-C with Vision Transformer
(b) FinyImageNett.
Figure Comparison of calibration between different MC-LayerNorm (One-Shot approximation) subsam- pling fractions for in-distribution (Test) and out-of-distribution (OOD).
17

---

Published in Transactions on Machine Learning Research (02/2024)
A.2.1
ImageNet Results for Vision Transformers
0.7
0.7
0.6
0.6
accuracy
LayerNorm
0.4
MC-LayerNorm f=
0.3
LayerNor
0.3
MC-LayerNomm f=0.8
MC-Dropout
MC-LayerNorm f-0.7
0.2
MC-LN One-Shot (Ours)
0.2
MC-LayerNom f=0
Test
OOD
Test
OOD
1.0
1.0
0.8
0.8
brier
0.6
0.6
0.4
Test
0OD
Test
OOD
0.30
0.25
0.20
ece
0.15
0.15
0.10
0.10
0.05
(b)
Figure 8: ImageNet-C with Vision Transformers: (a) We compare LayerNorm, MC-Dropoul (I=0.9, 10 MC iterations), and our novel method MC-LayerNorm (f=0.7, One-Shot approximation) on the expected calibration error (ECE), brier score, and accuracy. We do not include results for Masksembles this figure due to bad performance. Error bars are constructed from evaluation runs for the 15 corruptions applied the original test set and over three training runs. (b) Comparison calibration between different MC-LayerNorm (One-Shot approximation) subsampling fractions for in-distribution (Test and out-of-distribution (OOD).
18

---

Published in Transactions on Machine Learning Research (02/2024)
A.2.2 Result Tabels
Scenario
Method
ECE
Brier
Accuracy f
BatchNorm
0.044 + 0.001
0.120 + 0.001
0.925 + 0.001
LayerNorm
0.033 + 0.002
0.087 0.004
0.948 + 0.001
Masksemble
0.041 0.001
0.120 0.001
0.924 0.001
Prediction-Bi
0,042 0.002
0.119 0.001
=025:00000
MC-BN
0.044 0.001
0.120 0.001
1.925=0.00
Test
MC-Dropout
0.035 0.001
0.089 = 0.002
0.947
T-Scaling
0.003
0.083 = 0.003
0.948
MC-LN MC Approx (Ours)
0.031 0.001
0.098 + 0.002
+
MC-LN One-Shot (Ours)
0.002
0.093 = 0.003
0.941 + 0.003
MC-LN + T-Scaling (Ours)
0.019 0.002
0.090 + 0.003
0.941 +
0.003 0.003
BatchNorm
0.276 0.030
0.655 + 0.043
150220000
LayerNorm
0,206 0.012
0.479 = 0.026
0.712 0.015
Masksemble
0.237 0.009
0.587 + 0.014
0.627 0.007
Prediction-BA.
0.122 = 0.003
0.308 0.008
0.809 + 0.005
MC-BN
0.276 + 0.030
0.655 = 0.043
0.592
OOD
MC-Dropoul
0.200 + 0.004
0.468 + 0.007
0.716 7716:0.00
T-Scaling
#
0.454 0.025
MC-LN MC Approx (Ours)
0.166 + 0.013
0.454
0.703
MC-LN One-Shot (Ours)
0.158 0.014
0.419 0.032
0.731
MC-LN + T-Scaling (Ours)
0.125 0.012
0.399 0.030
0.731 0.022
BatchNorm
0.277 + 0.032
0.651 0.046
0.592 - 0.017
LayerNorm
0.212 0.013
0.484 0.029
0.709 + 0.018
Masksemble
0.253 + 0.024
0.610 + 0.035
0.614 + 0.012
Prediction-BA.
0.288 + 0.027
0.681 = 0.048
0.576 + 0.024
MC-BN
0.276 + 0.025
0.649 = 0.036
0.596 + 0.012
Mix
MC-Dropoul
0.202 0.004
0.467 + 0.007
0.718 + 0.007
T-Scaling
0.173 0.009
0.446 + 0.023
0.718 +
MC-LN MC Approx (Ours)
0.166 0.009
0.445 + 0.026
0.708 + 0.021
MC-LN One-Shot (Ours)
0.159 0.015
0.414 0.035
0.734 + 0.023
MC-LN + T-Scaling (Ours)
0.133 + 0.009
0.400 + 0.030
0.734 + 0.022
Comparison of calibration and accuracy for in-distribution (Test), out-of-distribution (OOD) and the zero-shot prediction-time domain adaptation (Mix) setting. In contrast to the figures, the standard deviation is computed distinctively by first averaging across multiple corruption levels, followed by calculating the standard deviation across various trained models.
Table 1: CIFAR-10-C with ConvNext:
19

---

Published in Transactions on Machine Learning Research (02/2024)
Scenario
Method
ECE
Brier
Accuracy t
BatchNorm
0.177 + 0.006
0.461 0.015
0.705 + 0.007
LayerNorm
0.114 0.003
0.280 0.006
0.828 + 0.003
Masksemble
0.156 0.004
0.478 + 0.003
0.675 + 0.002
Prediction-BN
0.172 0.002
0.450 0.001
0.713 #
MC-BN
0.177
0.461 + 0.015
0.705 0.007
Test
MC-Dropout
0.102 0.001
0.262 0.001
0.836 = 0.002
T-Scaling
0.056 + 0.001
0.252 0.006
0.828 + 0.003
MC-LN MC Approx (Our)
0.080 0.001
0.245 0.002
0.837 1: 0.003
MC-LN One-Shot (Ours)
0.000
0.252 0.003
0.834 0.003
MC-LN + T-Scaling (Ours)
0.027
0.234 0.001
0.834 + 0.003
BatchNorm
0.478 +.011
1.177 1 0.016
0.219 + 0.005
LayerNorm
0.443 0.020
1.049 = 0.032
0.337 = 0.013
Masksemble
0.426 + 0.004
1.144 + 0.004
0.195 + 0.002
Prediction-BN
0.373 + 0.001
0.956 = 0.002
0.362 + 0.004
MC-BN
0.478 + 0.013
1.177 + 0.016
0.219 + 0.005
OOD
MC-Dropout
0.403 + 0.015
0.982 + 0.021
0.367 + 0.006
T-Scaling
0.276 0.021
0.904 0.027
0.337 + 0.013
MC-LN MC Approx (Ours)
0.369 0.001
0.970 0.002
0.339 + 0.003
MC-LN One-Shot (Ours)
0.381 0.004
0.980 + 0.002
0.345 + 0.003
MC-LN + T-Scaling (Ours)
0.861 0.003
0.345 0.003
BatchNorm
0.480 0.014
1.182 + 0.016
0.219 0.003
LayerNorm
0.437 0.023
1.039 = 0.037
0.340 0.016
Masksemble
0.420 + 0.018
1.133 - 0.026
0.200 0.008
Prediction-BN
0.472 + 0.005
1.175 0.009
0.215 + 0.001
MC-BN
0.478 0.019
1.175 + 0.025
0.221 + 0.009
Mix
MC-Dropout
#.016
0.985 = 0.020
0.367 + 0.005
T-Scaling
0.278 0.023
0.910 0.032
0.335 + 0.014
MC-LN MC Approx (Ours)
0.370 1 0.004
0.970 0.006
0.339 + 0.007
MC-LN One-Shot (Ours)
0.381 + 0.008
0.984 + 0.008
0.343 + 0.008
MC-LN + T-Scaling (Ours)
0.229 + 0.005
0.857 + 0.008
0.345 + 0.009
Comparison of calibration and accuracy for in-distribution (Test), out-of-distribution (OOD) and the zero-shot prediction-time domain adaptation (Mix) setting. In contrast to the figures, the standard deviation is computed distinctively by first averaging across multiple corruption levels, followed by calculating the standard deviation across various trained models.
Table 2: Tmay-TmageNet with ConvNext:
20

---

Published in Transactions on Machine Learning Research (02/2024)
Scenario
Method
ECE
Brier
Accuracy f
LayerNorm
0.016 0.003
0.051 + 0.003
0.968 + 0.003
MC-Dropout
0.019 - 0.003
0.054 0.001
0.966 + 0.001
Masksemble
0.034 + 0.003
0.114 0.006
0.927 + 0.005
Test
T-Scaling
0.010
0.050 = 0.004
0.968 + 0.003
MC-LN Approx (Ours)
0.015 0.004
0.060 = 0.005
0.962 0.004
MC-LN One-Shot (Ours)
0.015 0.003
0.069 + 0.013
0.955 + 0.010
MC-LN + T-Scaling (Ours)
0.019 + 0.015
0.067 + 0.016
0.958 + 0.010
LayerNorm
0.194 0.020
0.490 - 0.025
0.690 + 0.011
MC-Dropoul
0.206 +: 0.020
0.499 + 0.024
0.688 + 0.006
Masksemble
0,259 = 0.020
0.621 0.029
0.011
OOD
T-Scaling
0,154 + 0.022
0.461 - 0.022
0.690 + 0.011
MC-LN Approx (Ours)
0.211 0.023
0.522 0.037
0.673 = 0.020
MC-LN One-Shot (Ours)
0.180 + 0.018
0.499 + 0.015
0.679 + 0.018
MC-LN + T-Scaling (Ours)
0.146 + 0.024
0.469 + 0.011
0.686 + 0.016
Table 3: CIFAR-10 for Vision Transformers:
Comparison of calibration and accuracy for in-distribution (Test), out-of-distribution (OOD) and the zero-shot prediction-time domain adaptation (Mix) setting. In contrast to the figures, the standard deviation is computed distinctively by first averaging across multiple corruption levels, followed by calculating the standard deviation across various trained models.
Scenario
Method
ECE
Brier
Accuracy t
LayerNorm
0.100 0.003
0.274 0.006
0.826 + 0.004
MC-Dropout
0.083 0.002
0.276 0.003
0.819 1.:900.000 0.003
Masksemble
0.209 0.003
0.568 + 0.004
0.625 + 0.002
Test
T-Scaling
0.029 + 0.000
0.253 0.005
0.826 + 0.004
MC-LN Approx. (Ours)
0.084 0.002
0.268 = 0.002
0.823 + 0.003
MC-LN One-Shot (Ours)
0.045 0.003
0.255 + 0.005
0.824 + 0.005
MC-LN + T-Scaling (Ours)
0.019 - 0.000
0.256 = 0.003
0.820 + 0.001
LayerNorm
0.325 + 0.010
0.905 = 0.017
0.377 + 0.012
MC-Dropout
0.272 + 0.005
0.864 + 0.004
0.375 + 0.003
Masksemble
0.483 0.008
1.201 0.010
0.197
OOD
T-Scaling
0.152 0.008
0.796 + 0.014
0.377 0.011
MC-LN Approx. (Ours)
0.289 + 0.007
0.875 0.005
0.373 + 0.006
MC-LN Onc-Shot (Ours)
0.226 + 0.011
0.834 + 0.015
0.372 + 0.011
MC-LN + TTScaling (Ours)
0.141 + 0.005
0.792 + 0.005
0.372 + 0.005
Table 4: Tiny- ImtgeNet for Vision Transformors:
Comparison of calibration and accuracy for in-distribution (Test), out-of-distribution (OOD) and the zero-shot prediction-time domain adaptation (Mix) setting. In contrast to the figures, the standard deviation is computed distinctively by first averaging across multiple corruption levels, followed by calculating the standard deviation across various trained models.
21

---

Published in Transactions on Machine Learning Research (02/2024)
A.3 Proof of theorem 3.1
Lemma A.1. Denole V
Ni [ as |Si| M-
Ni (at n!)2 and Va= 1 N. D ((ab)² (x1)² -07)²
Assume
(1) Assume that 1 (a))) is bounded, then we have that, as |Si oc,
(at)
S n²) +o NO,
(8)
(2) Assume that -> is bounded, then we have that, as | + oc,
|S|| |/ (10))2 (o')²) ta NO, Vaa).
(9)
Proof. The proof follows from Hájek 1960 under standard regularity conditions.
Proposition A.2. Under the assumption that f(.) in eq. (1) and eq. 4 is differentiable almost everywhere, theorem A.1 implies that there erists variance V such that
|/ 1/2
[T++ h6th N(.
(10)
Proof.
The proof follows from the delta method in Vaart 1998
Proposition A.2 means that the effect of MC Layer Normalization is to approximately add Gaussian noise to the activations of a corresponding model with Layer Normalization.
A.4 Hyperparameters
All models are trained with the AdamW optimizer provided by the timm library Wightman et al. (2023) using default parameters for $1, $2 and e. We use grid-search to tune the learning rate from values between [e 3, le - 5] and for weight decay form values between [0.1, 1e 4).
For CIFAR-10 and TinylmageNet, we run fine-tuning for 20 epochs with a batch size 512 and with a constant learning rate. Meanwhile, for ImageNet, we run fine-tuning for epochs with a batch size of 256 and constant learning rate. For Masksemble models extend learning to 20 epochs the needed changes to the network are much bigger than for the other methods. Finally, run training for 100 epochs with a batch size of 1000 reducing the learning rate by 0.1 on a plateau with patience 2 checked on the validation set for Criteo.
All vision models are trained with absolutely basic data augmentation: We use padding + random crop for CIFAR-10 and TinyImagee ImageNet. while only laking a random crop for ImageNet. Additionally, we use a horizontal flip (p=0.5) augmentation for all vision datasets. No training augmentations are used for the Critco models.
For Temperature scaling (Guo el al. 2017
parameters proposed by Pleiss 2024 NLL Loss, LBFGS
as an Optimizer with learning rate of 0.01 and rumning for a maximum of 50 iterations. For each dataset, a fixed set of samples from the training set is set aside as a calibration set.
All models were trained on an internal cluster consisting of NVIDIA A100 GPUs.
22


Reviews:
[{"review_id": "MZNnl7k3WT", "reviewer": "Reviewer_NxVm", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- Estimating/calibrating the uncertainty/confidence is a major challenge in neural network predictions, and this paper addresses it with an intuitive and simple-to-implement  approach.\n- Overall, the paper is well-written and clear to follow.\n- The related work is well written and it highlights both existing work in the scope of the paper and the existing parallel works to estimating uncertainty of neural networks.\n- This work shows empirically a relative enhancement in the performance of the proposed method in different settings, challenging existing the baselines. However, the major advantage of this method is that it can be applied to several settings unlike the shown baselines.\n\nWeakness:\n- This paper claims that comparing to Prediction-Time Batch Normalization is sufficient, however, a work such as (Monte Carlo Batch Normalization (Teye et al. (2018)) or the Temperature scaling should be present in the experimental section in all settings. Alternatively, a thorough discussion of why MC-LayerNorm is better suited for uncertainty estimation should be provided.\n- In settings 2 and 3 of the experimental section, the paper does not clarify well why \"MC-LayerNorm is the only norm layer based uncertainty calibration method that is applicable to this architecture”,  or why it is only compared to “the baseline of the traditional LayerNorm” in the Criteo challenge? \n- What is the effect of the relative $S_l$ set size to $N_l$ on the final performance on both In-distribution (Test) and Out-of-distribution (OOD) test when it is closer to 0? When the ratio is larger than 70%, is there an explanation why the performance is very stable? Also why is this ablation not present on the Criteo dataset experiment?\n- Temperature scaling showed better performance on Vision Transformer, how about ConvNext? The higher performance of the  combination of MC-LayerNorm and post-hoc Temperature scaling seems to be caused by the temperature scaling itself.\n\nMinor comment:\nCitations should be inside parentheses when they are not to be read as part of the text.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Summary:\n\nThis work proposes a new normalization module, the MC Layer Normalization module, with the goal of calibrating the uncertainty estimation in deep learning models. Differing from traditional Layer Normalization, MC-LayerNorm samples randomly from the units in a given layer and computes the normalization statistics over the sampled units. This allows for a more effective estimation of prediction uncertainty, particularly in scenarios of distribution shift. The module's theoretical design relies on combining Bayesian modeling and Monte Carlo integration. To show empirical evidence of the effectiveness of the proposed layer, this work tests the MC-LayerNorm in different settings where out-of-distribution examples are present, challenging the confidence of the model over these examples.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See questions listed as part of the weaknesses.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "hfWLIqGnj4", "reviewer": "Reviewer_RYYS", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**:\n\n* The paper proposes a new module, MC-LayerNorm. \n* The new module's asymptotic behavior is theoretically analyzed in Theorem 3.1.\n\n**Weaknesses**:\n\n* The motivation on why MC-LayerNorm could be better than other baselines is missing: to me MC-LayerNorm is still another type of approximated Bayesian methods, but it is unclear why it should be better than others. Can you provide more intuition?\n* Related to the above concern, this paper's claim is contradictory to the existing observation in [R1] --- [R1] claims that \"approximate Bayesian methods are not enough for calibration and we need the post-hoc calibration methods\". This makes sense to me as the \"approximation\" in Bayesian methods manifests suboptimal results in calibration and we need the post-hoc calibration methods for the right uncertainty (in terms of the known perfect-calibration definition --- see Guo et al. (2017) for details). Can you discuss whether approximate Bayesian methods are required for calibration and discuss pros and cons of approximate Bayesian methods and post-hoc calibration methods for calibrated uncertainty (in and out-of-distribution settings)?\n* An important baseline, temperature scaling, is missing in Figure 1. Can you add temperature scaling results in Figure 1 as did for Figure 2? It is even better if the same set of baselines are used in Figure 1 and Figure 2.\n* It would be important to discuss on the performance gap between T-Scaling and MC-LN+T-Scaling in Figure 2, but it is missing: to me temperature scaling could be the most effective method for ECE as it directly minimizes the related quantity (i.e., I expect T-Scaling and MC-LN+T-Scaling ECE could be very similar at least in \"Test\"), but the results are counter-intuitive, which is interesting and important to highlight the proposed approach. Can you discuss the reasons for this performance difference? Regarding this, can you share the experiment setup for temperature scaling (e.g., loss, optimization iterations, calibration set size...). \n\n[R1] https://proceedings.mlr.press/v80/kuleshov18a/kuleshov18a.pdf", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Bayesian methods (e.g., Bayesian Deep Neural Networks) offers a principled way to estimate uncertainty, but due to its intractable computational cost, researchers have been proposed approximate Bayesian methods, e.g., Monte Carlo Dropout or Prediction-Time Batch Norm. However, these two methods are not technically applicable on some settings and lead to \"suboptimal\" performance. To address this issue, the paper proposes a new module, called MC Layer Normalization, which is a drop-in replacement for Layer Normalization to seamlessly add uncertainty calibration capability. The proposed module is theoretically motivated from an approximate Bayesian perspective, and empirically demonstrated its efficacy over two architectures and three datasets.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* Provide intuition on why MC-LayerNorm should be better than other type of approximated Bayesian methods.\n* Discuss whether approximate Bayesian methods are required for calibration and discuss pros and cons of approximate Bayesian methods and post-hoc calibration methods for calibrated uncertainty (in and out-of-distribution settings).\n* Add temperature scaling results in Figure 1 as did for Figure 2.\n* (if possible) Use the same set of baselines in Figure 1 and Figure 2.\n* Discuss the reasons for the performance gap between T-Scaling and MC-LN+T-Scaling in Figure 2. \n* Share the experiment setup for temperature scaling (e.g., loss, optimization iterations, calibration set size...). \n\n**Additional editorial comments:**\n* \"conformal prediction methods Zhu & Rigotti (2021); Bates et al. (2021); Angelopoulos et al. (2022).\": citations are wrong. Please acknowledge the original work, i.e., \"Vladimir Vovk, Alex Gammerman, and Glenn Shafer. Algorithmic learning in a random world.\nSpringer Science & Business Media, 2005.\" or related papers. \n* \"(lower is better) quantifies the difference between the confidence of a model and its accuracy computed on bins of samples sorted by confidence Guo et al. (2017).\": same here, Guo et al. (2017) adopts ECE from the earlier work.", "broader_impact_concerns": "no concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "kHZZS5dayI", "reviewer": "Reviewer_ugqw", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n1. **Innovative yet Simple Approach**: Combines successfully ideas from MC batch normalization and layer normalization.\n2. **Practical Application**: Demonstrates some effectiveness across several tasks, particularly in OOD uncertainty estimation.\n3. **Clear Introduction and Method Presentation**: Well-written with a lucid explanation of the method and related work.\n\n**Weaknesses:**\n1. **Limited Comparative Analysis**: Needs more comparison with recent methods like Function-Space Regularization (Rudner et al., 2023).\n2. **Unclear Significance of Impact**: High error bars obscure the method's effectiveness; detailed result tables are recommended for clarity. Also, it can be useful to see the numbers and figures for the different corrupted  tasks and not only the average.\n3. **Outdated Training Routines**: Lacks advanced training techniques like data augmentation and learning rate schedules.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces MC Layer Normalization, a module enhancing neural networks' prediction uncertainty estimation. It serves as a replacement for Layer Normalization, targeting improved calibration in OOD scenarios. The approach is validated through various benchmarks and baselines.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "### Requested Changes\n1. Provide detailed result tables.\n2. Broaden comparative analysis.\n3. Run with recent training routines.\n\n\n### Final Comments\nThe paper offers an innovative approach for uncertainty estimation in neural networks, particularly for OOD scenarios. However, its impact seems incremental and empirical results are a bit weak. Enhancing training routines and providing detailed comparative analyses could strengthen the findings.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "LVoaWkQ8x3", "reviewer": "Reviewer_fKNX", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I think that this is a great and insightful paper that would be interesting to the TMLR audience. Its strengths include:\n- The problem addressed (i.e., neural network calibration) is important and well-motivated.\n- The manuscript is clearly written and easy to follow. The approach is placed in the context of relevant prior work.\n- The proposed method is sound and intuitively sensible. Furthermore, being a drop-in replacement for a common neural network building block (LayerNorm), the approach is conceptually simple and straightforward to implement and can be used to obtain calibrated uncertainty estimates for a range of popular model architectures (especially models to which previous calibration methods were not applicable).\n- The connection to approximate inference / Bayesian model averaging is derived clearly and makes the approach more principled.\n- The empirical evaluation convincingly demonstrates the practical benefits of MC-LayerNorm in improving model calibration (especially on out-of-distribution inputs / under distribution shift) on relevant benchmarks as compared to previous methods. I especially commend the experiments on Vision Transformers due to their prevalence as one of the most important modern neural network architectures.\n\nOverall, I thus recommend acceptance of the paper.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This manuscript develops a new method for calibrating neural network predictions via uncertainty estimation. The proposed approach builds upon Layer Normalization, introducing stochasticity into the computation of the normalization statistics by subsampling the preactivation features. This stochasticity is then exploited to obtain a Monte Carlo estimate of the predictive distribution (hence the name MC-LayerNorm). In addition to a sampling-based MC estimate, the paper also introduces an efficient one-shot approximation that leverages the asymptotic Gaussianity of the induced parameter posterior (i.e., when viewing it as a probabilistic model with a distribution over model parameters with deterministic LayerNorm, instead of stochastic LayerNorm with deterministic model parameters). Finally, the paper demonstrate the empirical effectiveness of the proposed MC-LayerNorm method on a set of standard calibration benchmarks (image classification on CIFAR-10-C, TinyImageNet-C, and ImageNet-C, Criteo Display Advertising Challenge) as compared to relevant baselines.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "N/A", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a drop-in replacement for layer normalization blocks, referred to as "MC layer normalization" so that the neural networks are endowed with uncertainty quantification capabilities. Transformers contain layer normalization blocks, so that the proposed MC layer normalization could be very useful for many practitioners. The proposed method is not the best UQ technique, but adds one easy-to-use entry to the available toolbox of approximate Bayesian methods. One of shortcomings here is the necessity of enhancing the MC layer normalization  with post-hoc calibration, since it not always possible to have a sufficient amount of held-out in-distribution calibration dataset. Despite this limitation, the method here is complementary and synergistic. Thus, the paper contains sufficient interesting contributions.
Decision: Accept (poster)