Amortized Conditional Normalized Maximum Likelihood
Authors: Aurick Zhou
Sergey Levine
Venue:  (2021)
Keywords: Uncertainty Estimation
Calibration

Abstract: While deep neural networks provide good performance for a range of challenging tasks, calibration and uncertainty estimation remain major challenges. In this paper, we propose the amortized conditional normalized maximum likelihood (ACNML)
method as a scalable general-purpose approach for uncertainty estimation, calibration, and out-of-distribution robustness with deep networks. Our algorithm builds on the conditional normalized maximum likelihood (CNML) coding scheme, which has minimax optimal properties according to the minimum description length principle, but is computationally intractable to evaluate exactly for all but the simplest of model classes. We propose to use approximate Bayesian inference technqiues to produce a tractable approximation to the CNML distribution. Our approach can be combined with any approximate inference algorithm that provides tractable posterior densities over model parameters. We demonstrate that ACNML compares favorably to a number of prior techniques for uncertainty estimation in terms of and calibration on out-of-distribution inputs.

Under review as a conference paper at ICLR 2020
AMORTIZED CONDITIONAL NORMALIZED MAXIMUM LIKELIHOOD
Anonymous authors Paper under double-blind review
ABSTRACT
While deep neural networks provide good performance for a range of challenging tasks, calibration and uncertainty estimation remain major challenges. In this paper, we propose the amortized conditional normalized maximum likelihood (ACNML) method as scalable general-purposs -purpose approach for uncertainty estimation, cal- ibration, and -iitr---iittitttioonn robustness with deep networks. Our algorithm builds on the conditional normalized maximum likelihood (CNML) coding scheme, which has minimax optimal properties according to the minimum description length principle, but is computationally intractable to evaluate exactly for all but the simplest of model classes. We propose to use approximate Bayesian inference technqiues to produce tractable approximation to the CNML distribution. Our approach can be combined with any approximate inference algorithm that provides tractable posterior densities over model parameters. We demonstrate that ACNML compares favorably to a number of prior techniques for uncertainty estimation terms of calibration out-of-distribution inputs.
INTRODUCTION
Current machine learning methods provide unprecedented accuracy across a range of domains, from computer vision to natural language processing. However, in many high-stakes applications, such as medical diagnosis or autonomous driving, rare mistakes can be extremely costly, and thus effective deployment of learned models requires not only high expected accuracy, but also way to measure the certainty in a model's predictions in order to assess risk and allow the model to abstain from making decisions when there is low confidence in the prediction. While deep networks offer excellent prediction accuracy, they generally do not provide the means to accurately quantify their uncertainty. out-of-distribution ddistrribttton inputs, where deep networks tend to make overconfident incorrect predictions (Ovadia et al. 2019). In this paper, tackle the problem of obtaining reliable uncertainty estimates under distribution shift.
This is especially true on out-of
Most prior work approaches the problem of uncertainty estimation from the standpoint of Bayesian inference. By treating parameters as random variables with some prior distribution, Bayesian inference can compute posterior distributions that capture notion of epistemic uncertainty and allow us to quantitatively reason about uncertainty in model predictions. However, computing accurate posterior distributions becomes intractable as use very complex models like deep neural nets, and current approaches require highly approximate inference methods that fall short of the promise of full Bayesian modeling practice.
Bayesian methods also have deep connection with the minimum description length (MDL) principle, a formalization of Occam's razor that recasts learning as performing efficient lossless data compres- sion and has been widely used as motivation for model selection techniques. Codes corresponding to
estimators and Bayes marginal likelihoods have been commonly used within the MDL framework. However, other coding schemes have been proposed in MDL centered around achieving different notions minimax optimality. Interpreting coding schemes as predictive distributions, such methods can directly inspire prediction strategies that give conservative predictions and do not suffer from excessive overconfidence due to their minimax formulation.
One such predictive distribution is the conditional normalized maximum likelihood (CNML) (Grün wald] 2007, Rissanen and Roos, 2007: Roos et al. 2008) model, also known as sequential NML or predictive NMI. (Fogel and Feder 2018h). To make prediction on new input, CNML considers

---

Under review as a conference paper at ICLR 2020
every possible label and tries to find the model that best explains that label for the query point together with the training set. It then uses that corresponding model to assign probabilities for each input and normalizes to obtain a valid probability distribution. Intuitively, instead of relying learned model to extrapolate from the training set to the new (potentially out-of-distriuution) input, CNML obtain more reasonable predictive distributions by asking "given the training data, which labels would make sense for this input?"
While CNML provides compelling minimax regret guarantees, practical instantiations have been exceptionally difficult, because computing predictions for test point requires retraining the model on the test point concatenated with the entire training set. With large models like deep neural networks, this can potentially require hours of training for every prediction.
In this paper, we proposed amortized CNML (ACNML), a tractable and practical algorithm for approximating CNML utilizing approximate Bayesian inference. ACNML avoids the need to optimize over large datasets during inference by using an approximate posterior in place of the training set. We demonstrate that our proposed approach is substantially more feasible and computationally efficient than prior techniques for using CNML predictions with deep neural networks and compares favorably to number of prior techniques for uncertainty estimation outoof-distribution inputs.
2 MINIMUM DESCRIPTION LENGTH: BACKGROUND AND PRELIMINARIES
ACNML is motivated from the minimum description length (MDL) principle, which can be used to derive connection between optimal codes and prediction. We begin with a review of the MDL principle and discuss the challenges in implementing minimax codes that motivate our method. For more comprehensive treatments of MDL, we refer the readers to (Grünwald 2007 Rissanen 1989).
Minimum description length. The MDL principle states that any regularities in a dataset can be exploited compress it, and hence learning is reformulated as losslessly transmitting the data with the fewest number of bits (Rissanen, 1989; Grünwald 2007]. Simplicity is thus formalized as the length of the resulting description. MDL was originally formulated in a generative setting where the goal is to code arbitrary data, and we will present a brief overview in this setting. We can translate the results to a supervised learning setting, which corresponds to transmitting the labels after assuming either fixed coding scheme for the inputs or that the inputs are known beforehand. While MDL typically described in terms of code lengths, in general, we can associate codes with probability distributions, with the code length of an object corresponding to the negative log- likellhodood under that probability distribution (Cover and Thomas 2006).
Normalized Maximum Likelihood. Let (x1:)) denote the maximum likelihood estimator for a sequence of data Fi.n over all A E θ. For any Fi:z x and distribution q over X", define a regret relative to the model class θ as
R(g. (1:n) def log Pō(z)... (11m) log (
(1) This regret corresponds to the excess number of bits q uses to encode Zi: compared the best dis- tribution in Q denoted (.)1.)). We can then define the normalized maximum likelihood distribution (NML) with respect to θ as
Pô(εl..) (1)m))
NML (11m) Eein X 2 Pb(a (T
p
when the denominator is finite. (Shtarkov 1987 Rissanen
The NML distribution be shown achieve minimax regret 1996)
p NML argmin
max R(q, O 1in)) EM
(3)
This corresponds, in sense, to an optimal coding scheme for sequences of known fixed length.
Conditional NML. Instead of making predictions across entire sequences at once, we can adapt NML to the setting where we make predictions about the next data point based on the previously seen
data, resulting conditional NML (CNML) (Rissanen and Roos 2007 Grünwald] 2007 Fogel and Feder 2018a). While several variations on CNML exist, we consider the following:
CNML (In 11) 1) x Po(ml..)) (3..)
2

---

Under review as a conference paper at ICLR 2020
For any fixed sequence 21.n-1, CMM. solves the minimax regret problem
p CNMI. argmin max log Po(as.)) (ax) log ((xn),
(5)
Zn
where the inner maximization is only over the last data point Zn.-
We can extend this approach to the supervised classification setting, where our models represent con- ditional distributions (y]æz). The CNML distribution, given sequence of already seen datapoints (I1m-1, 11m-1) and the next imput Tn, then takes the form
CNML (yn aan 21m-1: 11m-1:3tmm1) Polyrin (m) (zn)
(6)
and solves the minimax problem
CCML= argmin max log P8gn (m)...) (ym (rn) log g(yn)-
Ve
We see that this conditional distribution is amenable to our usual inductive learning procedure, where (Xim--1 1 M1:n-1) is training set, and want to output predictive distribution over labels y/, for a new test input In-
CNML provides conservative predictions. For each query point, CNML considers each potential label and finds the model that would be most consistent with that label and with the training set. If that model assigns high probability to the label, then minimizing the worst-case regret forces CNML to assign relatively high probability to it. In particular, compared to simply letting a model trained only on the training set extrapolate blindly, we expect CNML to give more conservative predictions explicitly considers what would have happened if the new data point had been included in the training dataset with each particular label.
out-of-distribution inputs, since it
We use a 2D logistic regression example to illustrate CNML's con- servative predictions, showing heatmap of CNML probabilities in Figure T CNML provides uniform predictions on most of the input space away from the training samples. In Figure 2 illustrate how CNML arrives at these predictions, showing the predictions for the parameters 6, and 0, corresponding to labeling the test point 1.
Figure 1: CNML probabilities with logistic regression model. Note that CNML provides uni- form predictions (indicated by the white color) on most of the input space away from the training set (shown in blue and orange dots).
(shown in pink in Figure 2 left) with either the label
However, CNML may be too conservative when the model class is very expressive. Naively applying CNML with large model classes can result in the per-label models fitting their labels for the query point arbitrarily well, such that CNML gives unhelpful uniform predictions even on inputs we would hope to reasonably extrapolate on. We see this in the 2D logistic regression example in Figure [ Thus, the model class ④ would need to be restricted some form, for example by only considering only parameters within certain distance from the training set solution hard constraint.
i train
i
-34
Figure 2: Given the labeled training set (blue and orange dots). want to predict the label at the query input (shown in pink in the left image), which the training set MLE drain confidently classifies as the blue class. However, CNML assigns near-uniform prediction on the query point, it computes new MLEs d and 9 (center and right images) by assigning different labels to the query point, and finds both labels consistent with the training data.

---

Under review as a conference paper at ICLR 2020
A 0.1
(b)
(0) = 10
Figure 3: CNMAP probability heatmaps with different levels of 1.2 regularization X ||ull22 We see predictions are less conservative as regularization increases.
Another approach for controlling the expressivity of the model class is to generalize CNML to regularized estimators instead maximum likelihood, resulting in normalized maximum a posterion (NMAP) (Kakade al., 2006) codes. Instead of using maximum likelihood parameters, NMAP selects Øs be the parameter that maximizes both data likelihood and regularization term, or prior, over parameters, and we can define slightly altered notions of regret using these MAP estimators in all the previous equations to get conditional normalized maximum a posteriori distribution instead. See Appendix D completeness.
Going back to the logistic regression example, we plot heatmaps of CNMAP predictions in Figure 3 adding different amounts of 1.2 regularization to the logistic regression weights. As we add more regularization, the model class becomes effectively less expressive, and the CNMAP predictions become less conservative.
Computational Costs of CNML. major practical issue with actually utilizing CNML or CNMAP with neural networks is the prohibitive computational costs of computing the maximum likelihood estimators for each new input and label combination. To evaluate the distribution on a new test point, one must solve nonconvex optimization problem for each possible label, with each problem involving the entire training dataset along with the new test point. This direct evaluation of CNML. therefore becomes computationally infeasible with large datasets and high-capacity models, and further requires that the model carry around the entire training set even when it is deployed. In settings where critical decisions must be made in real time, even running a single epoch of additional training would be infeasible. For this reason, NML-based methods have not gained much traction as a practical tool for improving the predictive performance of high-capacity models.
AMORTIZED CNML
In this section, we derive method, amortized conditional normalized maximum likelihood (ACNML). ACNML provides a tractable approximation for CNML and CNMAP via approximate Bayesian inference. Instead of directly computing maximum likelihood parameters over the query point and training set, our method uses an approximate posterior distribution over parameters to capture the necessary information about the training set, and thus reduces the maximization only the single new point. The computational cost at test-time therefore does not increase with training size. We specialize our notation the supervised learning setting, where our aim to obtain predictive distribution p(gn (x) after observing a training set (I1:n-1) 31m--1) 1) and a test input zn..
3.1 ALGORITHM DERIVATION
Incorporating an exact posterior into CNML. Given a prior distribution p(0), the Bayesian poste- rior likelihood conditioned on the training data is given by
p(0 [10 /1mn-1) 1) xx p(0) 1
(8)
We can write the MAP estimators in the CNMAP distribution for a fixed query input Tn as
0, argmax log po (2) + log pa (111--1/1mmm111111 + (())
(9)
#ce
log p(0 1i (imm-1))
4

---

Under review as a conference paper at ICLR 2020
Algorithm 1 Amortized CNML (ACNML)
Input: Model class θ Training Data (Iim-11, /1m Test Point: Zni Classes .l....kk k)
Output: Predictive distribution pay|a
Training: Run approximate inference algorithm training data (21mn-1 W1:n-1) get posterior
density ((0)
for all possible labels i (L...., k) do
Compute O, = argmax, log po + log (())
end for
Pdy (v)
Return p(yl Σ Pa, (K)mm)
We can thus replace the training data log-likelihodd pe density log (1:///11n111 (innn-1-nnn)) 11:nn1) when computing 0y. We can also recover CNML as special case of CNMAP by using a uniform prior, but as discussed previously, CNML with highly expressive model classes can lead to overly conservative predictions, so we will opt to use non-uniform priors that help control model complexity instead. For example, with deep neural networks, we may elect to use zero-mean Gaussian prior p(0) the network weights, corresponding to 1.2 regularization.
with the Bayesian posterior
ACNML with an approximate posterior. Of course, the exact Bayesian likelihood no easier to compute than the original training log likelihood. However, we can derive tractable approximation by replacing the exact posterior We can obtain an approximate posteriors via standard approximate Bayesian techniques such as variational inference or Laplace approximations. We focus on Gaussian posterior approximations for computational efficiency, and discuss in Section 3.2 why this class of distributions provides reasonable approximation.
-1,11n11) with an approximate posterior g(0) instead.
For practical purposes, we expect the approximate posterior log-likelihood to ensure the optimal Ó selected for each label retains good performance the training set. By replacing the likelihood over the training data with the probability under approximate posterior, it becomes unnecessary to retain the training data at test time, only the parameters of the approximate distribution. Optimization also becomes much simpler, as it no longer requires stochastic gradients, and the Gaussian posterior log density log (()) can serve as a strong convex regularizer.
ACNML algorithm summary A summary of the ACNML algorithm is presented in Algorithm 1 The training process for obtaining (0) only needs be performed once on the training set, whereas the inference step is performed for each test point. However, this inference step only requires optimizing the model parameters on single data point, with the regularizer provided by log (0).
3.2 ANALYSIS OF GAUSSIAN APPROXIMATIONS IN ACNML
In this section, argue that using a Gaussian approximate posteriors in ACNML, which correspond to second-order approximations to the training set log-likelihood, suffice for accurately computing the CNML distributions when the training set large. The intuition is that for large training sets, the combined likelihoods of all the training points dominate over the single new test point, so the perturbed MLEs O, remains close to the original training set MLE ô letting us rely on local approximations the training loss.
Under some simplifying assumptions, we can formalize this argument using the concept of influence functions, which measure how maximum likelihood parameters (and more general M Messtmmators) for a dataset would change if the dataset were perturbed by reweighting inputs an infinitesimal amount.
We recall that maximum likelihood estimators for a dataset with n datapoints (Iiss: Vi..) is given by
i = argmax
log pe (z/12)).
(10)
1=1
Influence functions analyze how ò relates to the MLE of a perturbed dataset
0 argmax
Eloo +
log Do (oooll)))))
(1I)
0
5

---

Under review as a conference paper at ICLR 2020
where oz.y.e is the new MLE if we perturb the training set by adding datapoint (z, y) with weight classical result (Cook and Weisberg, 1982) shows that ggy, is differentiable (under appropriate regularity conditions) with respect to with derivative given by the influence function
d0a.y.e de l==0
-H 6 Vo log Py(v)=),
(12)
where I is the MLE for the original dataset and H3 the Hessian of the mean training set log- segliiklihhodd evaluated at 0. CNML computes the MLE after adding datapoint (a. y) with equal weight to points in the training set, which is precisely given by Ōz.g.e evaluated at e = 1/n. Thus, for sufficiently large n, first order Taylor expansion around & should be accurate and the new parameter can be estimated by
duy 0 H Vo logPo(ylx).
(13)
which is equivalent to solving
1 (0 8) "V'V log + 2100 (0g0))=a(0 -0.
Bxy argmax
(14)
This suggests that, with large training datasets, the perturbed MLE parameters 8, in Equation can be approximated accurately using quadratic approximation to the training log- lieeekkiihood, corresponding to a Gaussian posterior obtained via a Laplace approximation. We can explicitly quantify the accuracy of this approximation in the theorem below, which is based Theorem from Giordano et al. 2019), with full details and proof in Appendix E
Theorem 3.1. (Adapted from Giordano al. (2019)) Consider training set with datapoints and an additional datapoint (2r, y). Assume assumptions 1-5 hold with constants Cops Cu, As as defined
in Appendix Let Ow.y denote the exact MLE if we had appended (x, y) to the training set, and On the parameter obtained via the approximation Equation 13
Let
8=
max sup Vo log log sup A20
(15)
n
n
AE
f S As, then
|| |0.. - dnv < 20 Cuo
(16)
Given bound how accurately we estimate the new parameters for CNML, we also explicitly quantify the accuracy of the resulting normalized distributions, with proof in Appendix E Proposition 3.2. Suppose y E y with |>] = k (classification with k classes). Let 0z.g be the exact MLE after appending the datapoint (z. y) to the training set, and let 0z.y be an approximate MLE with [2..- θ..v for each y. Further suppose log pa x)) is L-Lipschitz in 9.
Denote the exact CNML distribution for the fixed input x to be PCNML(V) x Po... (y|z) and an approximate CNML distribution PACNM((y) x Pū.., (y)a)). We then have
sup |log (V) log PACNML(y)) S 2L8.
(17)
Theorem 3.1 and Proposition 3.2 together suggest that the approximation produced by ACNML will be increasingly close the exact CNML distribution as the training set size n grows. However, this formal theoretical result only holds for sufficiently large datasets and under strong simplifying assumptions including smoothness and strong convexity of the training loss, so does not necessarily hold in practical settings with deep neural networks.
In the context of interpreting how different data points influence the predictions of neural networks, Koh and Liang|showed that influence function approximations were able to provide useful predictions for estimating leave-one-out retraining with deep convolutional neural networks. This closely resembles the conditions we encounter when computing parameters for each label of the query point with ACNML, with the key difference being that ACNML adds a datapoint while leave-one-out retraining removes one. This suggests second-order approximations to the training loss, corresponding to Gaussian approximations ACNML, may suffice to yield useful predictions about how parameters change when the query point is added, despite lacking formal guarantees with deep neural networks.

---

Under review as a conference paper at ICLR 2020
4 RELATED WORK
Minimum description length has been used to motivate neural net methods dating back to Hinton and van Camp (993)). who treat description length regularizer to mitigate overfitting. The Idea of preferring flat minima (Hochreiter and Schmidhuber 1997) also has its origins in the MDL. framework, as allows a coarser discretization of the weights (and thus fewer bits needed).
Bayesian methods typically serve as the starting point for uncertainty estimation in deep networks, and a commonly used approach to use simple tractable distributions to approximate the true posterior (Hoffman et al., 2013 Blundell et al. 2015) Ritter et al.) 2018). Recent work (Maddox et al., 2019; Dusenberry et al. 2020) has shown fairly simple posterior approximations are able achieve well-calibrated predictions with marginalization. Our method builds on top of these approximate posterior methods, but contrast to the Bayesian methods, where the posterior is typically used to efficiently sample models for Bayesian model averaging, our method uses the posterior density to enable efficient optimization for computing the CNML, without needing to retain the training data.
Ovadia et al. (2019) evaluate various proposed methods for uncertainty estimates in deep learning under different types of distribution shift. They found that good calibration on in-distribution points did not necessarily indicate good calibration under distribution shift, and that methods relying on marginalizing predictions over multiple models (Lakshminarayaman et al. 2016 Srivastava et al., 2014) gave better uncertainty estimates under distribution shift than other techniques. We show that our method ACNML maintains much better calibration under distribution shift than prior methods.
Perhaps most closely related distribution in the context of supervised learning (under the name predictive NML), citing its minimax properties. [Bibas et al. (2019b) estimate the CNML distribution with deep networks by finetuning the last layers of the network on every test input and label combination appended to the training set. Since this finetuning procedure trains for several epochs, it very computationally intensive at test-time and requires continued access to the entire training set when evaluating. In contrast, our method amortizes this procedure by condensing the information in the training data into distribution over parameters, allowing for much faster test-time inference without needing the training data.
work,[Fogel and Feder (2018b) advocate for the of the CNML
In the analysis for our approximation, we use influence functions (Cook and Weisberg 1982). which have been studied as asymptotic approximations to how M-estimators change when perturbing dataset. deep learning, Koh and Liang advocated for using influence functions to interpret neural nets, generate adversarial examples, and diagnose errors in datasets. We use a theorem from Giordano jet al. (2019), which broadened the necessary assumptions for these infinitisemal approximations to be accurate and provides explicit guarantees for fixed datasets rather than asymptotic results.
EXPERIMENTS
To instantiate ACNML, we must select method for obtaining the approximate posterior. In principle, any technique for computing a tractable posterior over parameters can be used, and we demonstrate this fiexibility by implementing ACNML on top of Stochastic Weight Averaging Gaussian (SWAG) (Maddox et al. (209), KFAC-Laplace (Ritter et al. 2018), and 2015). SWAG computes posterior by fitting Gaussian distribution to the trajectory of SGD iterates. For simplicity and computational efficiency, we instantiate ACNML with the SWAG-D variant, which uses Gaussian posterior with only a diagonal covariance. KFAC-Laplace uses Gaussian posterior approximation with the MAP solution as the mean and the inverse Hessian of the negative log likelihood as covariance, approximating the Hessian using KFAC (Martens and Grosse 2013) to allow for tractable inversion and storage. Bayes-by- bybckkppp (Blundell et al.) 2015) uses the reparameterization trick to learn a diagonal Gaussian posterior via the variational lower bound.
(Blundell et al.
For each model, we report results across 3 seeds. We compare negative log likelihood (NLL), accuracy, and expected calibration error (ECE) (Naeini et al., 2015) as well as showing reliability diagrams (Guo et al. 2017) to further assess calibration. For reliability diagrams, we sort data points by confidence and divide them into twenty equal sized buckets, plotting the mean accuracy against the mean confidence for each bucket. This allows the prediction relates to the actual accuracy, as well as showing how the confidences are distributed for each method.
qualitatively see how well the confidence of
7

---

Under review as a conference paper at ICLR 2020
10-
SWAGD+ACNML (ours)
SWA
0.9
0.9-
0.8
Aauurayy
0.8
SWAGD
ACcuraRy
0.8
Accurscy
0.7
0.7
0.6
0.7-
0.3
Confidence
0.8
1.0
0.6 0.7
Confidence
0.8
0/9
1.0
0.3
0.4
0.8 Confidence
0.7
1.0
(a) CIFARIO Test
(b) CIFARIO-C Corupttinn Level 3 (c) CIFARIO-C Corpution Lovel 5
Figure 4: Reliability diagrams plotting confidence vs. accuracy for VGGI6 on in in-distribution and out- of-distribution data. ACNML provides more conservative predictions than other methods, resulting in better calibration on out-of-disrribution inputs. For the OOD task, we show results for the Gaussian hlur comuptions at levels 3 and 5, with level 5 corresponding to a higher amount of corruption. Fach point shows the mean confidence and mean accuracy within bucket, so the spread of points along the r-axis shows that ACNML makes more low confidence predictions than other methods.
MNIST. We start with a simple il- lustrative task based on the MNIST dataset, where we construct out-of- distribution inputs by randomly rotat- ing the images in the MNIST test set. Here, ACNML is implemented on top of Bayes-by-backprog (Blundell et al.) 2013). and we compare to the MAP estimate and the marginal over models obtained from the same Bayes-by-backprop posterior. The results in Table [ show that all methods perform well on the in-distribution MNIST test set, though ACNML likelihoods are somewhat worse due to the more conservative CNML distribution. On OOD rotated digits, we see that ACNML exhibits substantial improvements in calibration measured by the ECE metric, as well as slightly better NLL value. In general. this agrees with what we expect from ACNML: the predictions are more conservative across the board, which does not necessarily improve results in-distribution, particularly for easy domains like MNIST, but offer considerable improvements in calibration for -ut-o--iistriuution inputs where errors are prevalent. We additionally compared to a much more computationally expensive instantiation of CNML used by Bibas et al. (2019a) (denoted naive CNML in Table D which di- rectly finetunes for several epochs using the training set to obtain the optimal parameters for each query point and label, rather than using the approximate posterior like ACNML does. This direct instantiation of CNML performs the best in terms of accuracy and NLL the in-distribution test set, while also improving over the MAP solution in terms of NLL and calibration on the OOD inputs. However, we find that ACNMI. is overall more conservative when using this particular posterior approximation, resulting in better NLL and calibration on the OOD inputs (see Appendix C for more detailed comparisons between ACNML and naive CNML).
MNIST Accuracy
Retated MNIST (OOD)
NLL
ECE
NLL
Accuracy
ECB
ACNML (ours)
0.1818
91,258
2.766
37.34
0.1540
MAP
0.08644
97.28
0.0047
3,991
37.2
0.4371
Marginal
1069
97.22
0.0313
3.017
37.63
0,2928
naive CNML
0.0774
98.05
0.0231
3,100
37.33
0.2197
Table 1: Comparative results for ACNML MNIST using a posterior obtained via Bayes by Backprop.
CIFAR and Corruptions We evaluate all methods using the VGG16 (Simonyan and Zisserman 2014) network architecture. Focusing on the most direct comparisons, we compare against the MAP solution for the given posterior, which is equivalent to Stochastic Weight Averaging (SWA) (Izmailov) jet al. 2018), and Bayes model averaging with SWAGD and KFAC-Laplace, which provide an apples-to-apples comparison to the two versions of our method that directly utilize the posteriors from these prior approaches. We use CIFAR10 (Krizhevsky 2012) for training and in-distribution testing. Following (Ovadia et al. 2019). we evaluate predictive uncertainty in out-of-distribution settings using the CIFARIO-CO CoFruto-Corrupeed (Hendrycks and Dietterich| 2019) datasets, which apply different severities of 15 common corruptions to the test set images. With this, we can assess performance over a wide range of distribution shifts, as well as how performance degrades as shifts become more extreme. We include additional comparisons across other methods and architectures in Appendix B
Examining the reliability diagrams in Figure 4 we see that ACNML provides more conservative (less confident) predictions than other methods, the point of being underconfident on the in-distribution CIFAR1O test set, while other methods tend toward being overconfident. On out-of-distribution datasets, where accuracy degrades, we see that ACNML's conservative predictions lead many better calibrated low-confidence predictions, while other methods drastically overestimate confidence.
8

---

Under review as a conference paper at ICLR 2020
Method
Methcd
0.30
SWAGD+ACNNL (ours)
SWAGD BANDDD+ACNLLoooa
KFAC+ACNNL (ours)
2.0
KFAC+ ACNML (ours)
0.25
KFAC Laplace
SWAG3
SWAGD
0.20
SWA
1.5
0.15
1.0
0.10
0.05
0.5
0.00
Corruptioneeell
Comrutiionevel
(a) CIFAR IOC VGG16 ECEs (lower is better)
(b) CIFARIOC VGGI6 NLLs (lower betier)
Figure 5: ACNML compared against their Bayesian counterparts and the deterministic MAP baseline cutoof-distributinn CIFARIO-Corrus Corrupted datsets. plot medians and 95% confidence intervals across all corruptions. We see that ACNML methods (solid lines) achieve much lower ECE at higher corruption values, and ACNML with SWAGD also achieves better NLL than other methods.
All methods perform similarly in terms of accuracy in all domains, and we find that ACNML's more conservative estimates perform competitively with Bayesian methods in NLL and calibration on in-distribution datasets, with all evaluated methods performing reasonably well in in-distribution (see Table in Appendix B However, differences calibration are much more pronounced for the OOD results in Figure 5 We see that as the corruption strength increases, ACNML variants provide much better calibration while performing similarly to or slightly better than other methods in terms of NLL.
Timing Comparison vs. standard CNML: In Table 2] we examine the computational costs of our method. We compare against a naïve imple- mentation of CNML that fine-tunes for N epochs on each test point and la- bel, similarly the method proposed by Bibas et al. (2019)). In total, predicting single input with & possible labels involves rumning k.N epochs of training. While ACNML is over two orders of magnitude faster than naïve CNML even with just single epoch of training (our experiments with naive CNML on MNIST used epochs), it still slower than standard inference. The computational requirements of our method scale linearly with the number of classes, but are constant with respect to dataset size. It is also not easily amenable to data batching, as new copies the model parameters are needed for each data point. Timing experiments are run using single NVIDIA 1080Ti, using MNIST for the MNIST MLP timing reselts and using CIFARIO for VGG16 and WideResNee28x 10, with no parallelization over data points.
MNIST MLP
VGG16
WRN28x10
ACNML (ours) naive CNMI. (per epoch)
0.08s
0.37s
1.1s
13.83s
102.0s
359.
feedforward inference
0.0001s
0.0013s
0.004s
Table 2: Inference time per input (in seconds).
DISCUSSION
In this paper, we present amortized CNML (ACNML) as an alternative to Bayesian marginalization for obtaining uncertainty estimates and calibrated predictions with high-capacity models, such as deep neural networks. The CNML distribution is a theoretically well-motivated strategy derived from the MDL principle with strong minimax optimality properties, but actually evaluating this distribution is computationally daunting. ACNML utilizes approximate Bayesian posteriors to tractably approximate and be instantiated top of wide range of approximate Bayesian methods. We view ACNML as step towards practical uncertainty aware predictions that would be essential for real-world decision making. Future work could further improve on our proposed method, for example by combining ACNML with more complex and expressive posterior approximations. In particular, training losses are highly non-convex and have many local minima, so incorporating local approximations around multiple diverse minima could allow for even more reliable uncertainty estimation. More broadly, tractable algorithms inspired by ACNML could in the future provide for substantial improvement in our ability produce accurate and reliable confidence estimates on out-of- dutto---istriuutionn inputs, improving the reliability and safety of reamiing-enabled enabled systems.

---

Under review as a conference paper at ICLR 2020
REFERENCES
K. Bibas, Y. Fogel, and M. Feder. A New Look at an Old Problem: A Universal Learning Approach to Linear Regression. IEEE International Symposium on Information Theory Proceedings, 2019-
July:2304 2308 5 2019a. doi: 10.1 10.1009/IST12200998844999 URL htt.tttp://axx...oggg/.
1905. /tt.a//0.0...a.aaoo....o........
201 8849398
K. Bibas, Y. Fogel, and M. Feder. Deep poml: Predictive normalized maximum likelihood for deep neural networks. arXiv preprint arXiv:/904. /2286, 2019b.
C. Blundell, Cornebise, K. Kavukcuoglu, and D. Wierstra. Weight Uncertainty Neural Networks. 32nd Infernational Conference on Machine Learning, ICMI, 2015, 2:1613-1622, 2015.
R. D. Cook and S. Weisberg. Residuals and influence regression. New York: Chapman and Hall, 1982.
T. M. Cover and J. A. Thomas. Elements of Information Theory (Wiley Series in Telecommunications and Signal Processing). Wiley-Interscience,, USA, 2006. ISBN 0471241954.
M. W. Dusenberry, G. Jerfel, Y. Wen, Y.a. Ma, J. Snoek, K. Heller, B. Lakshminarayanan, and D. Tran. Efficient and Scalable Bayesian Neural Nets with Rank-1 Factors. arXiv preprint arXiv:2005.07186 2020.
Y. Fogel and M. Feder. Universal Supervised Learning for Individual Data. 12 2018a. URL http:/ http://awxxxx....aaa/a/
09520
Y. Fogel and M. Feder. Universal batch learning with log-loss. In 2018 IEEE International Symposium on Information Theory (ISIT), pages 21-25, 2018b.
Y. Gal and Z. Ghahramani. Dropout as Bayesian Approximation: Representing Model Uncertainty in Deep Learning. 33rd International Conference on Machine Learning, ICML 2016, 3:1651-1660, 2015.
R. Giordano, W. Stephenson, R. Liu, M. Jordan, and T. Broderick. A swiss army infinitesimal jackknife. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 1139-1 1147, 2019.
P. Grunwald. A tutorial introduction to the minimum description length principle. 6 2004. URL /ttp:///axivy org/ abs//ab///aa math/4406077
P. Grünwald, T. Van Ommen, and others. Inconsistency of Bayesian inference for misspecified linear models, and a proposal for repairing Bayesian Analysis, 12(4): 1009-10091..3 2017.
D. Grünwald. The Minimum Description Length Principle (Adaptive Computation and Machine Learning). The MIT Press, 2007. ISBN 0262072815.
C. Guo, G. Pleiss, Y. Sun, and Q. Weinberger. On calibration of modern neural networks. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pages 1321-1330, 2017.
D. Hendrycks and T. Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. arXiv preprint arXiv: 2019.
G. E. Hinton and D. van Camp. Keeping neural networks simple by minimizing the description length of the weights. In Proceedings of the Sixth Annual Conference Computational Learning Theory, pages 5-13, 1993.
S. Hochreiter and J. Schmidhuber, Flat minima. Neural Computation, 911::1442, 1997.
M. D. Hoffman, D. M. Blei, C. Wang, and J. Paisley. Stochastic variational inference. The Journal of Machine Learning Research, 14(1):1303-1347, 2013.
Izmailov, D. Podoprikhin. T. Garipov, D. P. Vetrov, and A. G. Wilson. Averaging Weights Leads to Wider Optima and Better Generalization. In UAI, 2018.
10

---

Under review as a conference paper at ICLR 2020
S. M. Kakade, M. W. Seeger, and D. Foster. Worst-case bounds for Gaussian process models. In Advances in neural information processing systems, pages 619-626666666666666 2006.
P. W. Koh and P. Liang. Understanding black-box predictions via influence functions.
A. Krizhevsky. Learning Multiple Layers of Features from Tiny Images. University of Toronto, 2012.
B. Lakshminarayanan, A. Pritzel, and C. Blundell. Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles. Advances in Neural Information Processing Systems, 2016.
W.J. Maddox, P. Izmailov, T. Garipov, D. P. Vetrov, and A. G. Wilson. A simple baseline for bayesian uncertainty in deep learning. In Advances in Neural Information Processing Systems, 2019.
J. Martens and R. Grosse. Optimizing neural networks with troneccer-factored approximate curvature. In International conference on machine learning, pages 2408 2417, 2015.
M. P. Naeini, Cooper, and M. Hauskrecht. Obtaining well calibrated probabilities using bayesian binning. In Twenty-Ninth AAAI Conference on Artificial Intelligence, 2015.
Y. Ovadia, E. Fertig, J. Ren, Z. Nado, D. Sculley, S. Nowozin, J. V. Dillon, B. Lakshminarayanan. and J. Snoek. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift, 2019.
Rissanen, Stochastic Complexity in Statistical Inquiry Theory. World Scientific Publishing Co., Inc., USA, 1989. ISBN 9971508591.
J. Rissanen and T. Roos. Conditional NML universal models. In 2007 Information Theory and Applications Workshop, pages 337-341, 2007.
J. Rissanen. Fisher information and stochastic complexity. IEEE Transactions on Information Theory, 42(1):40 47, 1996. ISSN 00189448. doi: 10. 109/18.481776.
H. Ritter, A. Botev, and D. Barber. A scalable laplace approximation for neural networks. In 6th International Conference on Learning Representations, ICLR 2018- 9018-Conference Track Proceedings, volume 6. International Conference on Representation Learning, 2018
Roos, Silander, Kontkanen, and P. Myllymaki. Bayesian network structure learning using factorized NML universal models. In 2008 Information Theory and Applications Workshop, pages 272-276, 2008
Y. Shtarkov. Universal sequential coding of single messages. Problems of Information Transmission, 23(3):186, 1987.
K. Simonyan and A. Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409. 1556, 2014.
N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever, and R. Salakhutdinov. Dropout: A Simple Way to Prevent Neural Networks from Overfitting. Journal of Machine Learning Research, 2014.
V. G. Vovk. Aggregating Strategies. In Proceedings of the Third Annual Workshop on Computational Learning Theory, 1990.

---

Under review as a conference paper at ICLR 2020
A EXPERIMENTAL DETAILS
For obtaining approximate posteriors with SWAG and KFACCLLppace,, we follow the exact training procedures given in Maddox et al. (2019). We then implement ACNML on top of the diagonal SWAG posterior and the KFAC-Laplace posterior.
The variance of the SWAG posterior depends in a complex way on the learning rate and gradient covariances. To account for this, we introduce an additional temperature hyperparameter a and solve for the ACNML approximation using
0" argmax log pe (2n) + a
log (()).
(18)
To calibrate a, we can calculate the CNML distribution using a validation set, by training on the entire training set and the validation point, and then selecting a such that our ACNML procedure produces similar likelihoods. We can also treat a as tunable hyperparameter and select using a validation set, similarly to how temperature scaling (Guo et al. (2117) is used to achieve better calibration for prediction, or how the relative weighting of priors and likelihoods are used in generalized Bayesian inference (Vovk) 1990) or safe Bayesian inference (Grünwald et al. (2017) as way to deal with model misspecification. For our experiments using the SWAGD posterior, we heuristically tune a to be as large as possible without degrading the accuracy compared to the MAP solution. Note, however, that this procedure specific to the particular way which SWAG estimates the parameter distribution, and any posterior inference procedure that explicitly approximates the posterior likelihood (e.g.. Blundell et al. (2015) would not require this step. To select a for each model class, we swept over values [0.25, 0.5, 1, 1.5, and selected the highest value such that accuracy and NLL the validation set did not degrade significantly compared to SWA. For VGG16, we use a = 0.5 and for WideResNet28x 10, we used a L.5.
With our posterior q(0) being Gaussian with covariance Σ, we approximately compute the MAP solution for each label y as per Algorithm
by initializing On to be the posterior mean and iterating
0++1 0, + e D ( lo (v)=,) + V log ((e)),
(19)
using the covariance as a preconditioner. For our experiments, we run steps of gradient ascent on this objective, with a constant step size 0.5. We empirically find that 5 steps often enough to find an approximate stationary point with the SWAG-D posterior, and 10 steps for the KFAC-Laplace posterior.
For the reliability diagrams in Figure 4 we again follow the procedure used by Maddox et al. (2019). We first divide the points into twenty bins uniformly based on confidence (each bin has the same number of points), then plot the mean accuracy vs mean confidence within each bin. This differs from the reliability diagrams used by Guo al. (2017), where they divide the range of confidence values into bins uniformly. resulting in unevenly filled bins.
For our expected calibration error (ECE) numbers, we use the same bins as computed for our reliability diagrams, and compute
K ECE Σ P(i. Pa) -ei-
(20)
1=1
where P(i) is the empirical probability a randomly chosen point lies in bin i, o; is the accuracy within bin i, and e; is the average confidence in bin i. al/iiss/ggihhu..oo//jjaaaaa/vv// com/wjmaddoxx//wa Lawsssaanaaainnnnnaaa to include the ACNML procedure for test time evaluation, and include a copy of the modified codebase in the supplementary materials with instructions on how to reproduce our experiments. We additionally include pretrained models that were used for our experiments. Experiments were conducted using mix of local GPU servers and Google Cloud Program compute resources.
We adapted the SWAG authors' implementation at /tttts:///gihuhbb..ttt.
For the MNIST experiments, we used feedforward network with 2 hidden layers of size 1200, with no data augmentation. The posterior is factored as independent Gaussians for each parameter, with the prior for each parameter being zero-mean Gaussian with standard deviation 0.
12

---

Under review as a conference paper at ICLR 2020
CIFARTO Results
VGG16
WedeResNer8e 10
NLL
Accuracy
NEL
Accuracy
ECE
ACNML-SWAGD (ours)
0,2180
0.XI41
83.23
0.09
0.0246 0.0010
0.1130
0.0012
96.38 0.03
0,0122 = 0,0006
ACNML-KFAC (ours)
0.2329
0.0028
93.14+0.08
0.0016
MAP (SWA)
0.2694
0.0056
93.23 0.
0.0430 0.0010
0.1128
0.0014
96.41
0.0009 + 0.0004
SWAGD
0.2257
0.001T
93.31 0.00
0.0281
0.1125 0.0012
98
0.0012 0.0003
SWAG
0.2016 0.0031
03.00++0.110
0.0158 0.0080
0.1122 0.0009
96.32 +
0.0085 + 0.0006
Laplace
0.0013
0.02
SWA-Dropoi
0.2562 0.0025
92.85 0.0.4
0.0380
0.0007
0.1111 + 0.004
96.36 0.09
0.0107 + 0.000S
SWA. Temp
0.0245
93.61
0.0366
0.0063
SGD
0.3285 0.0139
93.17 0.14
0.0283 2 0.0022
0.1294 0.0022
96.41 08.11::.111 0.10
0.0166 1: 0,0007
Table 3: In-distribution comparative results We see that for in-distribution performance, ACNML variants perform comparably to other methods, without large separations between most methods. Results SWA-Temp and SGD are taken from Maddox al. (2019).
B FURTHER EXPERIMENTAL RESULTS AND COMPARISONS ON CIFARIO
In addition to the comparisons in the main paper, we additionally compare to SWA-Gaussian (SWAG). which uses a more expressive posterior than SWAG-D, and SWA with Monte Carlo Dropout Gal and Ghahramani 2015) (SWA-Drop). For reference, show in-distribution performance of all methods in Table 3 Overall, performance differences between all methods are quite small, and ACNML's conservative predictions do not improve NIL ECE over some baselines on in-distribution performance, which is to be expected, since the main aim our method is produce more calibrated predictions on ontoo-ddistributin tasks.
For completeness, we show expanded results on IT-ARIO-Corrppted in Figures 6 and g With the
same architecture, all methods generally have very similar accuracy. ACNML consistently achieves significantly better ECE on the more severe corruptions, and generally comparable or slightly better NLL.
While evaluating MC-Dropout, we found that adding dropout before each layer in VGGI6 (labelled VGG16Drop in 7 significantly improved performance on 1OFARIO-C. For fair comparisons, we reran all methods with the VGG16Drop architecture as well.
13

---

Under review as a conference paper at ICLR 2020
1552
SUNNGEI ACCAL (ouns)
KHNC ACMMOL (ouis)
KRHC-1 ASNC-Laplace
0.3
0.7
Lovd
(a) CIFAR VGG16 Accuracies (higher betler)
2.5
SWANGD AWNGD-AACMLLL teurs)
SWANG
SAMA
0.2
CIFAR IOC VGG16 ECEs (lower better)
SWAND AHAND-ACCMML
3.0
2.5
TWAR
SVIA
2.0
3.5
Comuptioneoel
(e) CIFARIOC VGG16 NLLLs (lower is better)
Figure 6: CIPAR10-C performance with the VGG16 architecture. Instantations of methods shown stripes. Boxplots show quartiles of statistic over different corruption types of the given intensity. with the mean indicated by circle. The accuracy (a) and NLL (c) for most methods are similar, but both ACNML variants attain significantly better ECE (b) on the more severe corruptions, the images move further out of distribution.
14

---

Under review as a conference paper at ICLR 2020
0.7
0555
SHANG-AACNA lours)
0.4
0552
KENC-ACNUL (our)
SWAKK
RUAN
Lagioce
Comv
CIFARIOC VGG16Dmp Accuracies (higher is better)
SWIND-ACCMLL leurs)
0.6
SWAG
SWA
NI MHCC-Laple
SWA MCCH spout
0.7
0.1-
(b) CIFARIOC VGGI VGG16Drop ECEs (lower better)
AWAND-ACHM
3.3
3.
TWA
KFAC Japlace
2.5
SWN--MCDrooour
2.0-
15-
133
Comuptionooell
CIFARIDC VGG16Drop NLLS (lower better)
Figure 7: CIFAR10-C performance with the VGG16Drop architecture. Instantations of methods shown in stripes. Boxplots show quartiles of each statistic over all different corruption types given intensity. with the mean indicated by a circle. Again, the accuracy (a) and NLL (c) for most methods are similar, but both ACNML variants attain significantly better ECE (b) the more severe corruptions, as the images move further out of distribution.
15

---

Under review as a conference paper at ICLR 2020
1.0
0.8
2.5
0.6
ACVML teurs)
SWANG
SWA
10 LinC-Lappace
0.2
5U1A MDN-MCCrrg opoutt
Corns sptiintnnnn
(a) CIFARIOC WRN28x10 Accuracies (higher better)
ACXMI (ours]
RUAA
KPAC Luplexe
MCDropout
Comupthnnt
(b) CIFARIOC WRN28<10 ECEs (lower betler)
SWAN
SVIA
LPpccLaplace
SUU-MEDDrspou
2.9
Commppiicniuual
(e) CIFARIOC WRN28x1O NLLS (lower is better)
Figure 8: CIFARIO-C performance with the WideResNer28x 10 architecture. Instantations of our methods are shown in stripes. Boxplots show quartiles of each statistic over all different corruption types of the given intensity, with the mean indicated by circle. Again, comparable methods the heavier corruptions (b). Note that the best performing prior method, SWAG, uses a substantially more expressive postcrior than diagonal approximation used by SWAGD+ACNML, whereas the comparable SWAGD method attains worse ECE.
that ACNMI. attains better ECE values than
16

---

Under review as a conference paper at ICLR 2020
1324
25
35
189
26-
175
125
05
100
311
20
3.4
30 0.1
(a) CNML Normalizers Eg P8y (M]x)
(b) NLLS
(c) Confidences
Figure 9: In Distribution Comparisons between ACNML and naive CNML. We plot scatter plots of the values of cach statistic for naive CNML (x-axis) ACNML (y-axis), with the red line indicating Looking the CNML normalizers, we that the ACNML adaptation procedure using the approximate posterior much less constraining than using the training set, resulting the normalizers being higher ACNML than naive CNML for almost all inputs. This leads to excess conservatism, with ACNML almost always having lower confidence its predictions, and many inputs with close to 0 NLL with naive CNML having higher NLL with ACNML.
C COMPARISONS BETWEEN ACNML AND NAIVE CNML ON MNIST
In this section, we include expanded comparisons between ACNML and naive implementation of CNML from Bibas et al. 2019b) that computes the MLE/MAP 0, for each label by appending the query point and label to the dataset and finetuning for N epochs. Both ACNML and naive CNML. initialized from the same MAP solution, with ACNMI. taking 5 gradient steps on the query point and posterior and naive CNML finetuning with the query point and training set for 5 epochs.
This naive implementation differs slightly from Bibas et al. (019b) in that we finetune the entire network, while Bibas et al. (2019b) proposed only tuning the last few layers. During the finetuning, we also append the query point and label to every batch in optimization, and downweighting that portion of the loss accordingly to get unbiased gradient estimates. We found this led to more efficient optimization than randomly sampling
We first examine how closely ACNML and naive CNML's predictions match on the same datapoint. To assess this, we compare the CNML normalization terms Σ, Pa, (v/zz), NLLs, and the confidences of the two methods. The CNML normalization term captures how much each procedure was able to adapt to different labels for that input. A higher normalization term for an input means that we flexible enough fit multiple different labels well together with the training set (or approximate posterior in the case of ACNML), and typically means a less confident prediction on that input.
In Figures and 10
We show scatter plots over 1000 randomly selected test points (from the in- distribution test set and the rotated OOD images respectively) comparing the CNML normalizers, NLLs, and confidences of ACNML and naive CNML. In each scatter plot, include a diagonal red line to illustrate where points would lie if predictions of ACNML and naive CNML matched exactly.
We additionally plot reliability diagrams for MNIST experiments in Figure
For the in-distribution test set, we see from the CNML normalizer plot that the ACNML adaptation procedure using the approximate posterior is much less constraining than using the training set, resulting in the normalizers being higher for ACNML than naive CNML for almost all inputs. This leads to excess conservatism, with ACNML almost always having lower confidence its predictions. As a result, we see that on many points where naive CNML outputted confident correct answers and achieved close to 0 NLL loss, ACNMI still incurs some higher losses due to its less confident predictions.
On the OOD rotated images, we again see that ACNML typically adapts more than CNML as measured by the CNML normalizers, though the difference is much less extreme compared to the nn-distribution dataset. In the confidence scatter plot, we again see that ACNML tends to make lower confidence predictions than naive CNML (especially when naive CNML's predictions confident), and as seen in Table and Figure []] result in ACNML having better NLL and calibration on the OOD inputs.
17

---

Under review as a conference paper at ICLR 2020
75
20.0
12.2
150
125
CNML Normalizers Ey pa. (ulxe)
NLLs
(c) Confidences
Figure 10: OOD Comparisons between ACNML and naive CNML. We plot scatter plots of the values of each statistic for naive CNML (x-axis) ACNML (y-axis). Looking the CNML normalizers, again see that the ACNML adaptation procedure using the approximate posterior less constraining than using the training set, with the normalizers being higher for ACNML than naive CNML for most inputs (though to lesser extent than the in distribution data). ACNML again outputs more conservative predictions with lower confidence many inputs, which leads to better NI.I and calibration the OOD dataset, unlike with the in-distribution test
1.00
ACVEL (oural
ACVEL (oura)
Naw CHOLL
0.9
Have CARL
095
0.8
HSP
930
8.0
8.00
0.6
0.75
03
0.5
0.6
033
1.0
02
0.3
0.1
0.8
09
1.0
MNIST Tost Sct
Randomly Retated MNIST (OOD data)
Figure Reliability diagrams plotting confidence vs. accuracy the MNIST test set and the randomly rotated MNIST test set (OOD). ACNML's conservative predictions provided better calibrated predictions on the OOD test set.
experiments on
Strictly speaking, the CNML distribution is not well defined when there exist multiple potential MLEs d,, that can output different predictions (prior references to CNML typically assume such MLEs are unique). However, the non-convexity of the objective for deep neural networks means multiple MLEs can exist, and to properly define CNML in this case, we would need to select a particular MLE to use when assigning probabilities in CNML. line with the min- max formulation of CNML, we propose to select the MLE 0, that maximizes the likelihood Pa, (ylx) of the query point and proposed label, this is the choice that maximizes the regret for that particular label over all MLEs.
Handling multiple MILEs in CNML:
With our naive CNML instantiation, we observe that during the finetuning for each query point x and label y. the predicted probability of that label (z)])) does not monotonically increase over iterations as we might hope (since we initialize to be the MLE of the training set, then update it to maximize likelihood of the training set with the query point and label), but can potentially oscillate substantially throughout the finetuning process. We suspect this is due to the stochasticity in the optimization procedure from to sampling minibatches of the training data causing the trajectory of parameters can potentially visit several different (approximate) local optima that output different predictions on the query point. While our instantiation of naive CNML simply used the parameter found at the end of epochs, we additionally compare against variant that explicitly tries to select the MLE that maximizes the likelihood of the proposed label. This variant heuristically uses the bset value of (v(c)) over all 0 encountered in the last epoch of finetuning. We in Table 4 that this variant,
18

---

Under review as a conference paper at ICLR 2020
MNIST
Rotated MNIST (OOD)
Accuracy
Accuracy
ACNML (ours)
0.1818 + 0.002
91.28 0.21
0.1013 = 0.0000
1.766 + 0.019,
3r.34 +: 0.000
0,1540 = 0.002
0.02
Marginal
0.1069 + 0.0067
97.22
0.24
0.0313 + 0.0010
3.017
0.022
37.63 0.31
0.2028 + 0.0032
CNML.
0.0774
0.34
naive CNML (mix over itrs)
0 1 0.0018
97.90
2: 0.23
0.03555 = 0.0005
2,991 1: 0.021
37.34 1: 0.003
0.1858 =: 0.0075
Table 4: Expanded comparative results for ACNML on MNIST using posterior obtained via Bayes by Backprop.
denoted naive CNML (max over itrs), gives more conservative predictions than naive CNML and improves in NLL and calibration on the OOD dataset.
D NMAP AND ACNML
NML type methods can be extended with prior-like regularization term on the selected parameter, resulting in Normalized Maximum a Posteriori (Kaade et al.) (006). also referred as Luckiness NML (Grunwald, 2004). For regularizer given by log p(O), NMAP assigns probabilities according to
p NMAP (x²) Pāçw=))(") (x")
0(x") argmax pogp(("") log p(0).
Similarly to CNML, there are several variations on NMAP or LNML that predict slightly different distributions, but we adopt the one of the same form as our CNML. Similarly to how NML was extended CNML, NMAP chosen via MAP rather than MLE. As mentioned Section 3.1, with a non-uniform prior, ACNML actually approximates version of conditional NMAP, with the Bayesian prior term the parameters corresponding to the additional regularizer.
be extended to conditional version, again with the 0's being
We also note that with the calculations in section 3.1, NMAP on the new test point, with regularizer corresponding to the likelihoods on the training data. In this perspective, ACNML approximates CNML by using an approximation that training loss regularizer.
that CNML be viewed as performing
E DETAILS OF ANALYSIS IN SECTION 3.2
E. BOUNDING ERROR IN PARAMETER ESTIMATION
Here we state the primary theorem Giordano al. (2019) along with the necessary definitions and assumptions.
Here, we attempt to estimate an unknown parameter 0 E S4 R.D. where Qg compact. Suppose we have a dataset N datapoints and weight vector Wi..... WNN Let g:))) denote the gradient of the loss datapoint evaluated at 0, and (0) the Hessian. We can then define
1
G(8, w)
(.):(0)
(21)
(0,ww)
Σ w; (0).
(22)
The MLE ((w) for the dataset weighted by w is given by solving for G(î(w), w) 0. Let 1 denote the vector of weights consisting of all 1s. We define O to be the MLE for the whole unweighted dataset, which is equivalent to evaluating 0(1.)) and also define the corresponding Hessian H1 H (0,, 1. We now wish to estimate @(w) using first order approximation around 0 given by
uu(w) ê - H Aw).
(23)
19

---

Under review as a conference paper at ICLR 2020
where we define A = w Tw- The theorem will proceed to bound
|| [((w)
8u|| for suitable
weights w.
Now we further define g(0) E RNXD to be the concatenation of all gi (0)s and similarly for h(0) e xDDDD We let [l9(0)],, and ||А(0)]| to refer the p-norms when treating those vector quantities.
Assumption (Smoothness): For all 0 E Ωg each gu (0) is continuously differentiable. (, 1w) is nonsingular and
Assumption 2 (Non-degeneracy):: For all 0 E Ωg, H(,,,
sup ( ) < Cap S
(24)
Assumption 3 (Bounded averages): There exist finite constants C and Cn such that suPseRt VN [1g(0)l22 Cg and supsens v [Ih(0)]22 & Ch.
Assumption (Local Smoothness): There exists a Ag > and finite constant La such that | < Ln
0 0- 01 A implies
0 - 81
V/
Assumption 5 (Bounded weight averages). finite constant Cw.
VN || is uniformly bounded for all e W by
We note that assumption is equivalent to H1 being strongly positive definite. Assumption 5 is not relevant for our use cases, but stated for completeness.
Condition 1 (Set Complexity): There exists a 5 2 0 and corresponding set Ws C W such that
N
max sup W'W AEfta
Σ
w; - ())
I
(25)
N
1=1
1
max sup wEW Aef
Σ
we
1)A+(0)
8.
(26)
N
Condition 1 essentially describes the set of weight vectors for which O. will be accurate approxi- mation within order 6.
Definition 1: Given assumptions 1-5, define
Cμ = 1 + DC
(27)
As min C=
(28)
n
We now state the main theorem of Giordano al. 2019
Theorem (Error Bound for the approximation). Under assumptions 1-5 and condition 1,
[du(w) B(w) ||, < 2C C C
8 < As + max Wss
(29)
We can now apply the above theorem to provide error bounds for setting where we have a training set of n datapoints and wish to consider the MLE after adding a new datapoint z. The issue is that the theorem as stated bounds the error of the approximation when the approximation is centered around the uniform weighting over the datapoints, which would be appropriate for considering the impact of removing datapoints from the dataset.
To apply the theorem to bound the effects of adding a datapoint, we have to do some slight manip- ulation. We apply the previous theorem with N = n + 2, where (0) correspond to the gradients of training data point for in (1...., n), 9n+1 = V log pe(z), and 9n+2 = V log pe (z), and similarly for the Hessians (0). We have thus added the query point to the dataset, well as another fake point that serves to cancel out the contribution of the query point under a uniform weighting, so G(O, 1. and H H(,, 1w) are the mean gradients and Hessians for just the training set. Now supposing
20

---

Under review as a conference paper at ICLR 2020
assumptions 1-5 met for this problem, then need to check condition for the particular W; that contains the vector 0 of all 1s, except for a 2 in the last entry. We can then find the smallest that satisfies
sup de0,
+ 2 In+2(0)
(30)
sup 0E
N + h 2
Σ 66
(31)
and long As. applying the theorem bounds
|
((w)
The above theorem gives explicit conditions for the accuracy of the approximation that we can verify for a particular training set and query point. Under assumptions that we have some limiting procedure for growing the training set such that the constants defined hold uniformly, we can extend this to an asymptotic statement to explicitly say that the approximation error decays as O(n" ?)
Commentary:
E.2 BOUNDING ERROR THE RESULTING CNML DISTRIBUTION
We now provide the proof for Proposition 3.2 which we restate here. For notational simplicity, we ignore any dependence on the input z, which we consider fixed.
Proposition E.1 (3.2). Suppose E Z with |2] le (for example classification with & classes). Let 0_ be the exact MLE after appending to the training set, and let 0g be an approximate MLE with || 0- 0. for all z. Further suppose log po(z) is L-Lipschitz in 0.
Denote the exact CNML distribution PCNML (z) & Pá, (2) and an approximate CNML distribution PACNML(2) x Pú. (2). Then, we have the bound
sup flog (2)
log PACNML(2)) 2L6.
(32)
Proof. The assumed bound 0- 0_
< 5 combined with I-LLipschitzness implies bound
differences of logits of each class
log Po, (z) - log Pa. (z) s L8.
(33)
We note that the log probabilities of the exact CNML distribution PCNMIL (PACNML is given by a similar expression using 0g instead of 0.) is given by
PCNML (z) log (z) log Σ Po.
(34)
z'EZ
For any Z, we then expand, apply the triangle inequality and then Equation 33 obtain
llog PCNML (z) log PACNML(2)
logP8, (z) log Pš. (z) log Σ Pa.,
(z') + log Σ Ps (2')
z'EZ
z'EZ
(35)
< log Pa. (2) log Po, (2) log Σ
Pa - loog Σ
Po., (2')
z'EZ
2'EZ
(36)
< L8 log Σ pb., (x') log Σ Pa_ (2')
(37)
z'EZ
We now bound the difference between the log-normalizers
|log Σx' Pa_ (z') - -log E.- Po,,
Pú, (z1)
21

---

Under review as a conference paper at ICLR 2020
We first let Puni (2) = min (pi (2), (2)) and (max(2) max (z), (2)}. and note that Equation 33 implies log Pinax (2) S log Punin (z) Lø for all We then bound the difference in log-normalizers
log Σ pa. (2') - -log Σ Po. (2)
S log
Σ
Pmax (2') log
Σ #'EZ
Punin (2')
(38)
2'EZ
2'EZ
Z'EZ
Evez Pinax (2') Σ
log
(39)
E'EZ Pmin (z')
EZEZ exp (log Pmax (2'))
log
(40)
Exez Paaiimmn (=minzz))
Evez exp(log Pmin (z') + Lo)
I log
(41)
Evez Pmii (2')
exp(L5) Exez Pmin (z')
log
(42)
ΣvEZ (m)n(
T.S.
(43)
Plugging back into Equation 37, we have the following bound for all
llo PCNML (2)
log PACNML (z) - 2L6.
(44)
22


Reviews:
[{"review_id": "Yb3QJseQ4e", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper presents an approach based on conditional normalized maximum likelihood (CNML) for uncertainty estimation, calibration, and out-of-distribution robustness with deep networks. CNML is intractable to compute in general and therefore the authors propose a tractable approximation which uses approximate Bayesian inference techniques. Experimentally, the authors show that their new approach is competitive and sometimes better than existing approaches for uncertainty estimation and calibration on out-of-distribution test data points.\n\nThe paper is well written, easy to understand and the authors are well versed with the domain (the writing clearly demonstrates scholarship, knowledge and expertise). As far as I can tell, related work is appropriately cited and the authors have done an excellent job in describing how their current work advances and relates to prior work on this topic. My criticisms/suggestions which I outline next are purely constructive and my hope is that they will greatly improve the paper:\n\n1. The algorithmic advances (Algorithm 1) are not that novel. I am not sure why the equation for the regret and definitions of unconditional normalized maximum likelihood are presented in the paper. The authors can simply start with the non-regret based definition of CNML given in Equation 6 and derive their approximation (the minimax optimization is not relevant to the algorithm). \n2. It would be informative to include comparison of your method with the method of Bibas et al. 2019 in terms of accuracy, NLL and ECE scores (at least for the MNIST dataset). In a way, what you are proposing is a faster, more scalable approach and it would be informative to know how far its estimates are from a much slower (but potentially more accurate) approach by Bibas et al.\n3. Another important comparison missing in the paper is how various methods perform on synthetic datasets where CNML can be exactly/accurately computed. As such, it is unclear whether smoothing/averaging effects of CNML are responsible for performance improvements on out-of-distribution samples or whether it is highly dependent on the quality of the approximate Bayesian inference method used.\n4. Latex tricks are used to somehow squeeze Figure 1 in the paper. Perhaps, what you could have done is removed definitions of (unconditional) NML.\n5. Section 3.2, although important from a theoretical point of view does not advance the main thesis/point of the paper (a faster approach for uncertainty estimation/calibration that uses approximate Bayesian inference). Also, most of the results described in section 3.2 are already known (specifically, Theorem 3.1). \n\nOverall, a good paper which advances the state-of-the-art on an important problem (uncertainty estimation and calibration) in neural networks. However, the paper lacks novelty and its presentation and experimental evaluation can be improved. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "TYWfMBohol8", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes ACML method which provides a tractable approximation to the CNML distribution and reduces the computational complexity of CNML method.\n\nStrength:\n1. Provide a tractable approximation for CNML which makes training easier and faster.\n2. Theoretically prove the approximation error is bounded, but the proof part has too many typos and cannot be confirmed.\n \t\nWeakness:\n1. The author didn’t use the experiment to show the accuracy of CNML distribution and ACNML distribution. The author only compared the training time of these two methods. Since in equation 18, the given bound of the approximation error increases with k, we don’t know whether this approximation is well bounded.\n2. In Table 1, MAP is a better method than ACNML for the MNIST dataset.\n3. In Table 3, it misses the results of ACNML-KFAC for WideResNet28x10 dataset, and it seems that the proposed methods don’t outperform other existing methods.\n\nTypos:\n1. In Equation 34, there should be a θ ̃_z instead of θ ̂_z\n2. In Equation 36, the third and fourth terms sum over z’, but there is no such parameter, and the symbol of absolute value disappears, due to this reason, I cannot verify if this equation holds.\n3. Looks like Equation 39 misses a minus symbol.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "Hg1lF5baxiX", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose the amortized conditional normalized maximum likelihood (ACNML) method for scalable, general-purpose uncertainty estimation, calibration and out-of-distribution robustness for deep neural networks. ACNML is a tractable approximation to CNML that leverages Bayesian inference by using an approximate posterior in place of optimizing over the dataset during inference.\n\nBased in the formulation provided CNML does not seem to be intractable but prohibitive.\n\nThe proposed approach is very simple and intuitive, even after the modification in (19). The analysis of the second-order approximation in Section 3.2 though interesting does not seem to impact the proposed method (as implemented) and the experimental results.\n\nIt is difficult to evaluate the advantage of the proposed method because from the results: 1) all compared methods perform about the same in terms of accuracy, 2) the predictions from the proposed models do not necessarily result in better NLL values (as a proxy for uncertainty estimation performance), or calibration on in-distribution datasets. It is only in terms of calibration of OOD calibration where the proposed approach seems to have an edge over existing methods.\n\nMinor:\n- $p_\\theta(x)$ is not defined.\n- there is a missing reference in the Appendix, above (20).", "extra_scores": {}, "extra_text": {}}, {"review_id": "4RGFljF0XQj", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper explores the application of an alternative inference scheme for learning well-calibrated predictive models especially target to ut-of-distribution samples. The paper is based on the so-called conditional normalized maximum likelihood, which is an inference approach based on min-max regret principles. Experimental evaluation shows promising performance to robust predictions under distributional shifts. \n\nStrong points:\n\n- The paper explores the use of alternative and well-founded inference approaches. \n\n- The presented solution is well-motivated and scales to large data samples and large models. \n\n- The empirical evaluation shows promising results on ut-of-distribution samples. \n\nWeak points:\n- The adaptation of the conditional normalized maximum likelihood approach presented in this work lacks the min-max regret guarantees provided by the original framework. \n\n- Empirical evidence about the advantage of the method is limited. \n\n- The presented method is too cautious when making predictions which are not out-of-distribution samples, and its accuracy is not competitive in this regime. \n\nI support the recommendation for acceptance of this work because it explores the use of a well-founded inference framework, the conditional normalized maximum likelihood, which optimal min-max regret guarantees but which is computationally intractable in its exact form. This work proposes an approximate inference scheme which makes the application of this inference framework tractable for deep neural networks. \n\n\nQuestions:\n\n- Does  the provided analysis of second-order approximations provide any relevant insight? It is based on very strong assumptions which do not hold on reality.\n\n- What happens when there are multiple maximum-likelihood estimates like happens in deep ensembles? Every member of the ensemble achieves zero training error. \n\n- Why don't you compare with Bibas et al. (2019)?\n\nMinor comments:\n-  Space below Figure 1. \n\n- Y label Figure 1 b) and c) is not readable. \n\n- Missing Eq. reference in Pag. 11\n\n- Figure 2 is hard to read. Please consider another format (bar based plot?). \n\nPost-rebuttal: I thanks the authors for their answers to my concerns. After looking the reviews of the other reviewers and the ongoing discussions, I have decided to keep my score. I think this paper makes a novel proposal which deserves to be published. ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper had four borderline reviews with none enthusiastic about championing the merits of the paper. While it was felt that the extension of an existing technique to deep learning via amortization is a useful procedure, it is also not very novel and the experiments didn't demonstrate a significant leap in performance.
Decision: Reject