On the Pitfalls of Heteroscedastic Uncertainty Estimation with Probabilistic Neural Networks
Authors: Maximilian Seitzer
Arash Tavakoli
Dimitrije Antic
Georg Martius
Venue: ICLR 2022 Poster (2022)
Keywords: Uncertainty Estimation
Probabilistic Neural Networks
Aleatoric Uncertainty
Heteroscedastic Uncertainty
Analysis

Abstract: Capturing aleatoric uncertainty is a critical part of many machine learning systems. In deep learning, a common approach to this end is to train a neural network to estimate the parameters of a heteroscedastic Gaussian distribution by maximizing the logarithm of the likelihood function under the observed data. In this work, we examine this approach and identify potential hazards associated with the use of log-likelihood in conjunction with gradient-based optimizers. First, we present a synthetic example illustrating how this approach can lead to very poor but stable parameter estimates. Second, we identify the culprit to be the log-likelihood loss, along with certain conditions that exacerbate the issue. Third, we present an alternative formulation, termed $\beta$-NLL, in which each data point's contribution to the loss is weighted by the $\beta$-exponentiated variance estimate. We show that using an appropriate $\beta$ largely mitigates the issue in our illustrative example. Fourth, we evaluate this approach on a range of domains and tasks and show that it achieves considerable improvements and performs more robustly concerning hyperparameters, both in predictive RMSE and log-likelihood criteria.

Published as a conference paper at ICLR 2022
ON THE PITFALLS OF HETEROSCEDASTIC UNCERTAINTY ESTIMATION WITH PROBABILISTIC NEURAL NETWORKS
Maximilian Seitzer', Arash Tavakoli', Dimitrije Antic Georg Martius
1 Max Planck Institute for Intelligent Systems, Tübingen, Germany
2 University of Tübingen, Tübingen, Germany ma x iiimmil aiainnn.. weetzeertten www.eeeeeeeeneeeeeeeeeee mpg.
ABSTRACT
Capturing aleatoric uncertainty is a critical part of many machine learning systems. In deep learning, a common approach this end is to train a neural network estimate the parameters of heteroscedastic Gaussian distribution by maximizing the logarithm of the likelihood function under the observed data. In this work, we examine this approach and identify potential hazards associated with the use of log- ligliikll conjunction with gradient- optimizers. First, we present a synthetic example illustrating how this approach can lead to very poor but stable parameter estimates. Second, we identify the culprit to be the log-likelihood loss, along with certain conditions that exacerbate the issue. Third, we present an alternative formulation, termed 3-NLL, in which each data point's contribution the loss is weighted by the B-exponentiated variance estimate. We show that using an appropriate .3 largely mitigates the issue in our illustrative example. Fourth, we evaluate this approach on range of domains and tasks and show that it achieves considerable improvements and performs more robustly concerning hyperparameters, both in predictive RMSE and log-likelihood criteria.
INTRODUCTION
Endowing models with the ability to capture uncertainty is of crucial importance in machine learning. Uncertainty can be categorized into two main types: epistemic uncertainty and aleatoric uncertainty (Kiureghian & Ditlevsen, 2009). Epistemic uncertainty accounts for subjective uncertainty in the model, one that is reducible given sufficient data. By contrast, aleatoric uncertainty captures the stochasticity inherent in the observations and can itself be subdivided into homoscedastic and heteroscedastic uncertainty. Homoscedastic uncertainty corresponds to noise that is constant across the input space, whereas heteroscedastic uncertainty corresponds to noise that varies with the input.
There are well-established benefits for modeling each type of uncertainty. For instance, capturing epistemic uncertainty enables effective budgeted data collection in active learning (Gal et al., 2017), allows for efficient exploration in reinforcement learning (Osband et al., 2016), and is indispensable in cost-sensitive decision making (Amodei et al., 2016). On the other hand, quantifying alcatoric uncertainty enables learning of dynamics models of stochastic processes (e.g. for model-based or offline reinforcement learning) (Chua al., 2018; Yu al., 2020), improves performance in semantic segmentation, depth regression and object detection (Kendall & Gal, 2017; Harakeh & Waslander, 2021), and allows for tisk-sensitive decision making (Dabney al., 2018; Vlastelica et al., 2021).
We examine a common approach for quantifying aleatoric uncertainty in neural network regression. By assuming that the regression targets follow a particular distribution, we can use neural network to predict the parameters of that distribution, typically the input-dependent mean and variance when assuming a heteroscedastic Gaussian distribution. Then, the parameters the network can be learned using maximum likelihood estimation (MLE), i.e. by minimizing the negative log-likelihood (NLL) criterion using stochastic gradient descent. This simple procedure, which the de- de-facto standard (Nix & Weigend, 1994; Lakshminarayanal et al., 2017; Kendall & Gal, 2017; Chua al., 2018), is known to be subject to overconfident variance estimates. Whereas strategies have been proposed to alleviate this specific issue (Detlefsen al., 2019; Stirn & Knowles, 2020), argue that equally important

---

Published as a conference paper at ICLR 2022
0.3
0.3
MSE Loss Optimal
Target
RMSE
0.2
0.0
0.1
-0.3
0.0
0
5
10
0.00
0.25
0.50
0.75
1.00
Input X
Update Steps
X 10
Figure 1: Training probabilistic neural network to fit simple sinusoidal fails. Left: Learned predictions (orange line) after 102 updates, with the shaded region showing the predicted standard deviation. The target function is given by y(z) 0.4 sin(2 + E, where is Gaussian noise with standard deviation of 0.01. Right: Root mean squared error (RMSE) over training, mean and standard deviation over 10 random seeds. For comparison, we plot the training curve when using the mean squared error as the training objective - achieving an optimal mean fit (dashed line) in 10a updates. This behavior stable across different optimizers, hyperparameters, and architectures (see Sec. B.2).
issue is that this procedure can additionally lead to subpar mean fiss. In this work, we analyze and propose a simple modification mitigate this issue.
Summary of contributions We demonstrate pitfall of optimizing the NLL loss for neural network regression, one that hinders the training of accurate mean predictors (see Fig. for an illustrative example). The primary culprit is the high dependence of the gradients While such dependence is generally known to be responsible for instabilities in joint optimization of mean and variance estimators (Takahashi et al., 2018 8; Stirn & Knowles, 2020), we identify fresh perspective on how this dependence can further be problematic. Namely, we hypothesize that the issue arises due the NLL loss scaling down the gradient of poorly-predicted data points relative to the well-predicted ones, leading to effectively undersampling the poorly-precd poreiy-peedicted data points.
the predictive variance.
We then introduce an alternative loss formulation, termed 3 3-NLL, that counteracts this by weighting the contribution of each data point to the overall loss by its 3-exponentiated variance estimate, where controls the extent of dependency of gradients on predictive variance. This formulation subsumes the standard NIL. loss for 3 = 0 and allows to lessen the dependency of gradients on the variance estimates for < 3 S 1. Interestingly, using 3 = completely removes such dependency for training the mean estimator, yielding the standard mean squared error (MSE) loss but with the additional capacity uncertainty estimation. Finally, we empirically show that our modified loss formulation largely mitigates the issue of poor fits, achieving considerable improvements range of domains and tasks while exhibiting more robustness to hyperparameter configurations.
2 PRELIMINARIES
Let X,Y be two random variables describing the input and target, following the joint distribution Y). We assume that Y is conditionally independent given X and that it follows some probability distribution P(Y X). In the following, we use the common assumption that Y is normally distributed given X; i.e. P(Y X) = A (j((X, ox))) where μ:: TR + and σ²: R R+ are respectively the true input-dependent mean and variance functions. Equivalently, we can write == ===((X)<(X)) with e(X) ~ A (0, o²(X)); i.c. is generated from X by ((X) plus a zero-mean Gaussian noise with variance x² (X)). This input-dependent variance quantifies the heteroscedastic uncertainty or input-dependent aleatoric uncertainty.
To learn estimates ((X), (X) of the true mean and variance functions, it is common to use neural network fe parameterized by 0. Here, ((X) and ox can be outputs of the final layer (Nix & Weigend, 1994) or use two completely separate networks (Detlefsen et al., 2019). The variance output hereby constrained to the positive region using a suitable activation function, e.g. softplus.
For notational convenience, we focus on univariate regression but point out that the work extends to the
multivariate as well.
2

---

Published as a conference paper at ICLR 2022
festure granuiarity
feature granularity
anco waighted gradients P(x) xx
Early
(NLL)
M(- 3
factor 1000
0.5 M(e)
Sample weighting
Sample weighting
ampiticatin
I fatsorr 32
P(a) A
factor 10
Pta) d
factor 1000
(MBE for mean)
Figure 2: Illustration of the pitfall when training with NLL (negative log- (og-likelihodd) versus our solution. An initial inhomogeneous feature space granularity (see Sec. 3.1) results early on in different fitting quality. The implicit weighting of the squared error in NLL can be seen as biased data-sampling with p(a) x cج() (see Eq. 6). Badly fit parts are increasingly ignored during training. On the right, the effect of our solution (Eq. 7) on the relative importance data points shown.
The optimal parameters liLL can then be found using maximum likelihood estimation (MLE) by minimizing the negative log-likelihood (NLL) criterion LNLL under the distribution P(X, Y):
E [ log (X) (Y (x)) + const 0 X,Y
NL. = arg min LNLL (0) arg min
(1)
26²(()
In contrast, standard regression minimizes the mean squared error (MSE) LMSE:
(Y 2
MASE arg min (MSE(0) = arg min
(2)
X.Y
2
In practice Eq. and Eq. 2 are optimized using stochastic gradient descent (SGD) with mini-batches
of samples drawn from P(X, Y). The gradients of LNL. w.r.t. (with respect to) ((X), 6² (X) are given by
((X) -Y VR X.Y 62(X)
62(X) - (Y A(X)) -Xx)
(3,4)
X.Y
2(02(X))²
3 ANALYSIS
We now return to the example of trying to fit sinusoidal function from Sec. 1. Recall from Fig. I that using the Gaussian NLL the objective resulted in a suboptimal fit. contrast, using MSE as the objective the model converged to the optimal mean fit in a reasonable time. We now analyze the reasons behind this surprising result.
From Eq. 3, we see that the true mean (xX) is the minimizer of the NLL loss. It thus becomes clear that a) the solution found in Fig. is the optimal one, and b) the NLL objective should, in principle, drive jæX) to the optimal solution (xX). So, why does the model not converge to the optimal solution? We identify two main culprits for this behavior the Gaussian NLL objective:
1. Initial flatness of the feature space create undercomplex but locally stable mean fit. This fit results from local symmetries and requires a form of symmetry breaking to escape.
2. The NLL loss scales the gradient of badly-predicted points down relative to well-predicted points, effectively undersampling those points. This effect worsens training progresses.
These culprits and their effect on training are illustrated in Fig. 2 (left). If the network cannot fit certain region yet because its feature space (spanned by the last hidden layer) is too coarse, would then perceive high effective data variance. This leads to down-weighting the data from such regions, fueling a vicious cycle of aelf-amplifying the increasingly imbalanced weighting. In the following, we analyze these effects and their reasons in more detail.
3.1 SYMMETRY AND FEATURE NON-LINEARITY
is instructive to see how the model evolves during training as shown in Fig. 3. The network first learns essentially the best linear fit while adapting the variance to match the residuals. The situation
3

---

Published as a conference paper at ICLR 2022
(a) After 000 updates.
(b) After 10000 updates.
(c) After 500000 updates.
0.4
0.4
Target
Target Y
Target
0.0
0.0
0.0
-0.4
-0.4
-0.4
Input
Inpur X
Input
Figure 3: Model fit using the NLL loss different stages of training shown in orange with to uncertainty band. Black dots mark training data. Fitting the function begins from the left and is visibly slow.
is locally stable. That is, due to the symmetries of errors below and above the mean fit, there is no incentive to change the situation. Symmetry breaking is required for further progress. One form of symmetry breaking comes with the inherent stochasticity of mini-batch sampling in SGD, or the natural asymmetries contained in the dataset due to, e.g., outliers. Moreover, we hypothesize that the local non-linearity of the feature space plays an important role in creating the necessary non-linear fit.
Let us consider the non-linearity of the feature space. This quantity not easy to capture. To approximate it for a dataset D, we compute how much the Jacobian Jy of the features f(x) w.r.t. the input varies in an L2-ball with radius around point x, denoted as the Jacobian variance:
V(x) BI
1
Σ E'EB.
1
Σ Z"EB,
B, (a' E D: ||
x' 12122s Sr.. (5)
B3
Figure 4 visualizes the Jacobian variance over the input space as a function of the training progress. Although initially relatively flat, it becomes more granular in parts of the input space, the parts which later well fit. The region with low Jacobian variance remains stuck in this configuration (see Fig. 1). This provides evidence that the non-linearity of the feature space is important for success or failure of learning on this dataset. However, why does gradient descent not break out of this situation?
Update Steps x
3
Update Steps x10%
3
Sampling Probability
1100
Jacobian Variance
-10
2
2
-10
-10
0
8 12
0
8
12
Input X
Input X
Figure 4: Jacobian variance over training time,
Figure Probability of sampling data point
using the mean of matrix V (z) (see Eq. 5).
input over training time.
3.2 VARRIANV WEIGHTING EFFECTIVELY UNDERSAMPLES
The answer lies in an imbalanced weighting of data points across the input space, Recall that the gradient NNLLL of the NLL w.r.t. the mean scales the error fi(X) - Y by a²(x) (Eq. 3). As symmetry is broken and the true function starts to be fit locally, the variance quickly shrinks in these areas to match the reduced MSF. If the variance is well-calibrated, the gradient becomes R(X)-Y _(X)-Y
8²(X) (a(X)-Y)² - Data points with already low error will get their contribution in the batch gradient scaled up relatively to high error data points "rich get richer" self-amplification. Thus, NLL acts contrary to MSE which focuses on high-error samples. If the true variance o² on the well-fit regions is much smaller than the errors on the badly-fit regions, or there are much more well-fit than badly-fit points, then learning progress is completely hindered on the badly-fit regions.
This is form of approximative second-order derivative computed numerically, which also gives non-zero results for networks with relu activation (in contrast to, for example, the Hessian).
4

---

Published as a conference paper at ICLR 2022
Another way view this is to interpret the different weighting of points changing the training distribution P(X,Y) to a modified distribution P(X, Y) in which points with high error have lower probability of getting sampled. This can be shown by defining P(X,Y) Z-1 P(X.)) x²(X)'
where Z = = Piz.v) drdy is a normalizing constant, and recognizing that the gradient of the NLL is
d=(w)
proportional to the gradient of the MSE loss in Eq. 2 under the modified data distribution P(X, Y):
(Y ((x))²
V (p) = Z PXX,Y - Y & VxEx,y~P(x..)
(6)
2
In Fig. 5, we plot P(X, Y) over training time for our sinusoidal example. It can be seen that the virtual probability of sampling a point from the high-error region drops over time until is highly unlikely to sample points from this region (10-9 as opposed to 10-3 for uniform sampling). We show that this behavior also carries over to a real-world dataset Sec. B.3.
Sometimes, "inverse-varianc" weighting" is seen as a feature of the Gaussian NLL (Kendall & Gal, 2017) which introduces self-regularizing property by allowing the network to "ignore" outlier points with high error. This can be desirable if the predicted variance corresponds to data-inherent unpredictability (noise), but it undesirable if it causes premature convergence and ignorance of hard- to-fit regions, shown above. In our method, we enable control over the extent of seff-regularization.
4 METHOD
In this section, we develop a solution method to mitigate these issues with NLL training. Our approach, which we term 3-NLL, allows choosing arbitrary loss- loss-interpolation between NLL and MSE while keeping calibrated uncertainty estimates.
4.1 VARIANCE- WEIGHTING THE GRADIENTS OF THE NLL
The problem want address is the premature convergence of NLL training to highly suboptimal mean fits. In Sec. 3, we identified the relative down-weighting of badly-fit data points in the NLL loss together with its self-amplifying characteristic as the main culprit. Effectively, NLL weights the mean-squart reaan-squred-erroo per data point with 2/ which can be interpreted as sampling data points with P(x) x Z.. Consequently, we propose modifying this distribution by introducing a parameter allowing to interpolate between NLL's and a completely uniform data point importance. The
resulting sampling distribution is given by P(x) o²
o² and illustrated in Fig. 2 (right).
How could this weighting be achieved? We simply introduce the suriance-weightingg term 20 to the LNLL loss such that it acts as factor on the gradient. We denote the resulting loss as B-NLL:
log 02) (Y - A(X))² + const) ]
CL-NLL= = E X.Y
(X)
(7)
2x²
where I.] denotes the stop gradient operation. By stopping the gradient, the variance-weighting term acts as an adaptive, input-dependes learning rate. In this way, the gradients of C.3 are:
[x(X) 02 (x))
Tx2(X) - ( -xx))
V VaLa VaLa-NLL(O) = E X.Y
(o = E X.Y
(8, 9)
28 28 (X)
Naturally, for 3 = O, we recover the original NLL loss. For B = 1 the gradient w.r.t. μ in Eq. 8 is equivalent the one of MSE. However, for the variance, the gradient in Eq. 9 a new quantity with 202 in the denominator. For values 0 < 3 < 1, we get different loss interpolations. Particularly interesting the case of B = 0.5, where the data points are weighted with (inverse standard deviation instead of inverse variance). In our experiments (Sec. 5), find that 8 = 0.5 generally achieves the best trade-off between accuracy and log-likelihood A Pytorch implementation of the loss function is provided in Sec. D.5.
Note that the new loss C3-NLL is not meant for performance evaluation, rather it is designed to result in meaningful gradients. Due to the weighting term, the loss value does not reflect the model's 's quality. The model performance during training should be monitored with the original negative likeliiidoood objective and, optionally, with RMSE for testing the quality of the mean fit.
5

---

Published as a conference paper at ICLR 2022
(a) LNLL
(b) CASE:
(c) Lg NLL with 3 0.5
1500
1500
1500
Count
1000
Count
1000
Count
1000
500
500
500
0
10-6 10
10
10
10
10 5
10"h
10P
Residual Error
Residual Error
Residual Error
Figure 6: Distribution of residual prediction errors depending on the loss function for the ObjectSlide dataset (see Sec. 5.2). Dashed lines show predictive RMSE. (a) The NLL loss (LNLL) yields multimodal residuals. There is a long tail of difficult data points that are ignored, while easy ones are fit to high accuracy. (b) The MSE loss (CMse) results in log-normal residual distribution. (c) Our 3-NLL loss (NI.) yields highly accurate fits on easy data points without ignoring difficult ones.
4.2 ALLOCATION OF FUNCTION APPROXIMATOR CAPACITY
Even though LMSE. LNLL. and C5-NLL all have the same optima w.r.t. the mean (and also the variance in the case of LNLL and C3-NLL), optimizing them leads to very different solutions. In particular, because these losses weight data points differently, they assign the capacity of the function approximator differently. Whereas the MSE loss gives the same weighting to all data points, the NLL loss gives high weight to data points with low predicted variance and low weight to those with high variance. 3 NLLLL interpolates between the two. The behavior of the NLL loss is appropriate if these variances are caused by true aleatoric uncertainty in the data. However, due to the use of function approximation, there also the case where data points cannot be well predicted (maybe only transiently). This would result in high predicted variance, although the ground truth is corrupted by little noise. The different loss functions thus vary in how they handle these difficult data points.
An example of how the differences between the losses manifest in practice is illustrated in Fig. 6. Here we show the distribution of the residuals for a dynamics prediction dataset containing easy and hurd to model areas. The NLL loss essentially ignores a fraction of the data by predicting high uncertainty. By analyzing those data points, we found that they were actually the most important data points to model correctly (because they captured non-trie trivial interactions in the physical world).
How important are the data points with high uncertainty? Are they outliers (i.e. do they stem from truly noisy regions), to which we would be willing to allocate less of the function approximator's s capacity? Or are they just difficult samples that are important to fit correctly? The answer is task- dependent and, as such, there is one-loss fits-all solution. Rather, the modeler should choose which behavior is desired. Our NLLLL loss makes this choice available through the / parameter.
5 EXPERIMENTS
In our experiments, we ask the following questions and draw the following conclusions:
Sec. 5.1: Does 3-NLL fix the pitfall with NLL's convergence?
Yes, B-NLL converges to good mean and uncertainty estimates across a range of B values.
Sec. 5.2: Does B B-NLL improve over NLJ in practical settings? How sensitive to hyperparameters 3-NLL? We investigate a diverse set of real-world domains: regression on the UCI datasets, dynamics model learning, generative modeling on MNIST and MaShion-MNIST. and depth-map prediction from natural images.
Yes, B-NTL generally performs better than NLL and is considerably easier to tune.
Sec. 5.3: How does 8-NLL compare to other loss functions for distributional regression? We compare with a range of approaches: learning to match the moments of Gaussian (termed "moment matching" (MM); see Sec. A), using Student's 's t-distribution instead of a Gaussian (Detlefsen et al., 2019), or putting different priors on the variance and using variational inference (XAMPP, x xVAMP* VBEM, VBEM*) (Stimn & Knowles, 2020).
It depends. Different losses make different trade-offs, which we discuss in Sec. 5.3.

---

Published as a conference paper at ICLR 2022
| = 0 (NLL)
(b) B 0.5
(c)
Learning Rate
le-03
0.2
Learning Rate
le-03
0.2
Learning Rate
le-03
0.2
3e-M4
RAMSE
3e-04
RMSE
3e-04
RMSE
1o-04
0.1
1o-04
0.1
1o-04
35.05
3c.05
3e-05
0
2
Architecture
Architecture
Architecture
Figure Convergence properties analyzed on the sinusoidal regression problem. RMSE after 200 000 epochs, averaged over 3 independent trials, is displayed by color codes (lighter is better) function of learning rate and model architecture (see See. D.I). The original NLL (3 0) does not obtain
good RMSE fits for most hyperparameter settings. Figure S3 shows results for the NLL metrie.
(a) NLL
(b) 3 = 0.25
3= 0.5
(d) B =
(e) MM
(f) std. dev.
10
W
gt datts shd
Figure 8: Fits for the heterosccdastic sine example from Detlefsen et (2019) (a-c). Dotted lines show the ground truth mean and =oo, respectively. (f) The predicted standard deviations (with shaded std. over 10 independent trials) with the same color code. Note that / = 0.5 and / graphs lie on top of one another. Inside the training regime, all 3- B-NLL variants (a-d) yield well-calibrated uncertainty estimates. Moment matching (e) significantly underestimates the variance everywhere.
We refer the reader to Sec. and Sec. D for a description of datasets and training settings. We evaluate the quality of predictions quantitatively in terms of the root mean squared error (RMSE) and the negative log-likelihood (NLL).
5.1 SYNTHETIC DATASETS
Sinusoidal without heteroscedastic noise We first perform an extended investigation of our il- lustrative example from Fig. 1 sine curve with small additive noise: y/ 0.4 sin(2mz) with & being Gaussian noise with standard deviation a = 0.01. One would expect that a network with sufficient capacity can easily learn to fit this function. Figure inspects this over range of architectures and learning rates.
We find that for the standard NLL loss (3 = O), the networks do not converge reasonable mean fit. There is a trend that larger networks and learning rates show better results, but when comparing this to 8-NLL with B 0.5 see that the networks are indeed able to fit the function without any issues. As expected, the same holds for the mean squared error loss (MSE) and 3-NLL with 3 = The quality of the w.r.t. NLL is shown in Fig. $3.
Sinusoidal with heteroscedastic noise We sanity-check that 3-NLL I is still delivering good uncertainty estimates on the illustrative example from Detlefsen et al. (2019) a sine curve with increasing amplitude and noise: y = x sin(x) xf1 + ε2. with y and &a being Gaussian noise with standard deviation o 0.3. Figure 8 (a-e) displays the predictions of the best models (w.r.t. NLL validation loss) and (f) compares the predicted uncertainties over 10 independent trials. Fitting the mean is achieved with all losses. On the training range, 8-NLL with 3 > 0 learss virtually the same uncertainties as the NLL loss.
5.2
WORLD DATASETS
UCI Regression Datasets As standard real-world benchmark in predictive uncertainty estimation, we consider the UCI datasets (Hemmádeez-Lobato & Adams, 2015). Table gives an overview comparing different loss variants. We refer Sec. B.4 for the full results on all 12 datasets. The results are encouraging: B-NLL achieves predictive log-likelihoone on par with or better than the NLL loss while clearly improving the predictive accuracy on most datasets.
7

---

Published as a conference paper at ICLR 2022
Table 1: Results for UCI Regression Datasets. We report predictive log-likelihood and RMSE (#: standard deviation). Ties denotes the number of datasets (out of 12) for which the method cannot be statistically distinguished from the best method (see Sec. B.4). We compare with Student-t (Detlefsen et al., 2019) and XVAMP/VBEM (Stirn & Knowles, 2020). Section B.4 lists the full results.
LL
RMSE
Loss
Ties concrete
energy
naval
yacht
Ties
energy
naval
yacht
Es-Nia.
3.22 + 0.31
-3.22
1.41
12.46
1.18
22.86 1 5.18
6,08
0.65
2:25 1 0.34 0,0021
= 0.0006
1.22 0.47
Co-NLL
0.25
33.11
+ 0.51
0
13.78
0.33
-197
+ 1,14
5.79 0.74
1.81
0.30 0.0012
0.0004
1.73
1.00
Ea-Nia.
0.5
2.1
1.12
1.44
Cg-NLL
0.75
-3.27
+ 0.34
-2880
0.59
13663
0.62
-1.87
0.55
8 5.67
0.73
1.31
0:45
0.0004
= 0.0001
1.97
1.03
Ca-Nta. 1.0
13.23
0.33
0.58
13.39
0.30
-2.27
1,07
9 5.55
1: 0.77
1.54
0.54
0.0004
0.0000
2.08
1.13
Cusa
-3.49
= 0.38
-.26
0.50
12.73
0.64
-11.2
31.0
5 6.28
0.82
2.19
0.28 0.0005
0.0001
3,02
1.38
CMSE
12 4.96
0.64
0.92
0.11
0.0004
0.0001
0.78++ 0.22
-2.46
0.34
-1.23
0.55
1.34
XVAMP
-3.06 0.15
-247
0.32
12.44
0.60
-0.99
0.33
8 5
1.87
0.32 0.0023 = 0.0004
0.99 0:43
XYAMP=
-3.03
0.13
2.41
0.32
12,80
0.55
-1.04 0.47
5.35
0.73
2.00
0.26
0.0020
0.0006
1.13 1:13::00 0.66
VBEM
-3.14
0.07
4.29
0.16
8.05
= 0.13
-2.65 0.10
5.21 = 0.58
2.29
0.33 0.0009
0.0004
1 0.84
VDEM
2.99
1 0.13
199
1
0.21
13.10
1 0.47
0.98
0.24
9 5.17
1 0.59
1.0
1 0.17
0.0015
& 0.0005
0.65 1 0,20
(a) ObjectSlide
(b)
CNLL
(B =0.5)
Density
(B- 1.0)
Density
10
0
RMSE
x10
NLL
RMSE
x10
NLL
Figure 9: Sensitivity analysis loss functions hyperparameters on the dynamics model learning tasks: ObjectSlide (a) and Fetch-PickAndPlac (b). The distributions over validation RMSE and NLL are shown as a function of hyperparameters, based on grid search over different model configurations (see Sec. D.2). While the NLL loss is highly sensitive when evaluating RMSE, the 3-NLL loss shows much less sensitivity and yields good results regardless of the exact configuration.
Dynamics models As major application of uncertainty estimation lies in model-based reinforce- ment learning (RI.. we test the different loss functions on two dynamics predictions tasks of varying difficulty, ObjectSlide, and Fetch-Pick PickAndPlace. In both tasks, the goal is predict how an ob- ject will move from the current state and the agent's ' action. Whereas ObjectSlide (Seitzer et al., 2021) is a simple ID-environment, Aetch-PickAndPleae (Plappert al., 2018) complex 3D robotic-manipun monoiiimmnipilltiion environment. The models are trained on trajectories collected by RL agents.
For both datasets, we perform a grid search over different hyperparameter configurations (see Sec. D.2) for a sensitivity analysis to hyperparameters settings, presented in Fig. 9. It reveals that NLL is vulnerable to the choice hyperparameters, whereas NTN.L achieves good results over wide range of configurations. The best performing configurations for each loss are then evaluated on hold-out test set (Table 2). One can see that the NLL loss results in poor predictive performance and also exhibits quite high variance across random seeds. Our method yields high accuracy and log-likelihood fits for range of |3 values, with = 0.5 generally achieving the best trade-off.
Generative modeling and depth-map prediction For generative modeling, we train variational autoencoders (Kingma & Welling, 2014) with probabilistic decoders on MNIST and FaShiOn-MNISST. For the task of depth regression, we modify and test it on the NYUv2 dataset (Silberman et al., 2012) with our loss (Fig. S6). Table 3 presents selected results for both tasks, yielding similar trends as before. We refer Sec. B.5 and Sec. B.6 for more details, including qualitative results.
ofate-oftthe-ae method (AdaBins; Bhat et al. (2021))
5.3 COMPARISON TO OTHER LOSS FUNCTIONS
The previous sections have demonstrated that our 3-NLL loss has clear advantages over the NLL loss. However, the comparison other loss functions requires more nuanced discussion. First,
8

---

Published as a conference paper at ICLR 2022
Table 2: Test results for dynamics models, using best configurations found in a grid search. The reported standard deviations are over 5 random seeds. We compare with Student-t (Detlefsen et al., 2019) and XVAMP/VBEM (Stim & Knowles, 2020).
1D-Sidee
Feich-Pick.AndPlce
RMSE 4
1
RMSE
f
0.0192 + 0.006
7.97
3.62
0,00163 0.00008
18.72 + 7.32
CD-NLL
0.25
0.0107 1 0.004
9.03
0.47
0,00102 0.00004
24.43 1 1.64
Gg NLL
0.5
0.0064 = 0.002
9.28
+ 0.75
0.00996 + 0.00002
24.68 +: 0.08
Ca-NLL
0.75
0,0087 0.003
6.61 1.83
0,00098 0.00001
22.77
0.17
Cg
1.0
0.0074 = 0.001
6.58 0.29
0.00102 + 0.00001
21.12 + 0.07
CMM
0.0078 = 0.001
diverges
0.00104 0.00003
19.33
1.31
EMSR
0,0068 = 0.001
0.00103 = 0.00000
Student-4
0.0155 + 0.006
11.30 1: 0.03
0.00117 1 0.00001
30.44 1: 0.08
XAMP
0.0118 + 0.002
10.58 0.19
0.00128 0.0000S
29.02 1:
X
0.0199 = 0.006
10.89 + 0.10
0.001 = 0.00001
29.19 + 0.08
VBEM
0.0039 0.000
3.79 0.00
0.00104 0.00003
17.39 0.29
VBEM
0,0280 = 0.011
10.13 0.49
0.001 18 + 0.00003
28.62 0 0.15
Table 3: Selected results for generative modeling and mepth-map prediction. Left: Training variational autoencoders on MNIST and Mashion-MNIST. Right: Depth-map prediction on NYUv2. Full results can be found in Table S3 and Table $4.
MNIST
Fushion-MNIST
NYUv2
Loss
RMSE
LL
RMSE
LL
Loss
RMSE
LL f
NLL
0
0.237 0.002
2116
0.170 + 0.001
1940 104
Ca-NLL
0.3854
-4522
NL-NLL
0.5
0.151 = 0.003
2220
+ 25
0.125 + 0.003
1639 = 52
CA-NLL
0.5
0,3789
-7.50
Na-NII
1.0
0,152 = 0.001
1706
+: 30
0,138 +: 0.002
1142 1:
Ca-NLL
1.0
0.3845
-5.10
Student-t
0.273 0.002
4291
1: 103
0.182 1: 0,002
2857 1.9
EMSE
0.3776
XVAMP*
0.225 = 0.001
3062
+ 215
0,160 + 0.002
2150 + 131
Ci
0.3850
VBEM*
0.176 0.008
3213
238
0.150 0.003
2244
So
we also test an alternative loss function based on matching the moments of the Gaussian (CMM: see Sec. A). While this loss results in high accuracy, is unstable to train and exhibits poor likelihoods.
Second, we test several loss functions based on a Student's t-distribution, including x VAMP and VBEM (Stirn& & Knowles, 2020). These approaches generally achieve better likelihoods than the 8-NLL; we conjecture this is because their ability to maintain uncertainty about the variance results in a better fit (in terms of KL divergence) when the variance is wrongly estimated. In terms of predictive accuracy, 3-NLL outperforms the Student's t-based approaches. Exceptions are some of the UCI datasets with limited data (e.g. "concrete" and "yacht") where VAMP and VBEM are on par or better than N-LL This is likely because these methods can mitigate overfitting by placing a prior on the variance. However, x VAMP and VBEM are also non-trivial to implement and computationally heavy: both need MC samples evaluate the prior; for x VAMP, training time roughly doubles as evaluating the prior also requires second forward pass through the network. In contrast, B-NLL is simple to implement (see Sec. D.5) and introduces no additional computational costs.
CONCLUSION
We highlight a problem frequently occurring when optimizing probabilistic neural networks using the common NLL loss: training gets stuck in suboptimal function fits. With our analysis. we reveal the underlying reason: initially badly-fit regions receive increasingly less weight in the loss which results in premature convergence. We propose a simple solution by introducing family of loss functions called N-NLL. Effectively, the gradient of the original NLL loss is scaled by the B- exponentiated per-sample variance. This allows for a meaningful interpolation between the NLL. and MSE loss functions while providing well-behaved uncertainty estimates. The hyperparameter gives practitioners the choice to control the reguaaauuaii strength of NLL: how important should high-noise regions or lufficult-to-predii data points be in the fitting process. In most cases, 3 = 0.5 will be a good starting point. We think the problem discussed in this paper is primarily why practitioners using the Gaussian distribution in regression or generative modeling tasks often opt for a constant homoscedastic (global) variance, as opposed the more general heteroscedastic (data-dependent) variance. We hope that our simple solution contributes to changing this situation by improving the usability and performance of modeling data uncertainty with deep neural networks.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGMENTS
The authors thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Maximilian Seitzer. Georg Martius is member of the Machine Learning Cluster of Excellence, EXC number 2064/1 - Project number 390727645, We acknowledge the financial support from the German Federal Ministry of Education and Research (BMBF) through the Tübingen Al Center (FKZ: 011S SUIS180392 18039B).
REPRODUCIBILITY STATEMENT
All settings are described in detail in Sec. C and Sec. D. We make full code and data available under https: / / thub. com/ ma us- 1 Lab/ beta-nl1.
REFERENCES
Dario Amodei, Chris Olah, Jacob Steinhardt, Paul Christiano, John Schulman, and Dan Mané. Concrete problems in Al safety. ArXiv, fass/606.065655 2016.
Shariq Farooq Bhat, Ibraheem Alhashim, and Peter Wonka. AdaBins: Depth estimation using adaptive bins. In 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 1008-4017 Los Alamitos, CA, USA, jun 2021. IEEE Computer Society. doi: 10.1109/CVVR446332222221 .00400. URL http:://o...eeeccmmmp.. ieeecomptteroocietyyy.oo org/ 19..109/CVPRRCR
2021. 00400.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. OpenAl Gym. ArXiv, abs//606.011544 2016.
Xiaotian Chen, Xuejin Chen, and Zheng-Jun Zha. Structure-aware residual pyramid network for monocular depth estimation. In Proceedings of the Twenty-Fighth International Joint Conference on Artificial Intelligence, IJCAI-19, pp. 694-700. International Joint Conferences on Artificial Intelligence Organization, 2019. doi: 10. 20924998/jcaii.2119998. URL https::/doi.org/ 24963/1
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep rein- forcement learning a handful of trials using probabilistic dynamics models. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Asso- ciates, Inc., 2018. URL
neurips. ec cc/paaprr2018/ELL www.:eesssssc-ea..ppppee pdf.
Will Dabney, Georg Ostrovski, David Silver, and Rémi Munos. Implicit quantile networks for distributional reinforcement learning. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1096-1 PMLR, 10-15 Jul 2018. URL https: ptttps:///proceedingss / v80 dabney htmll
Nicki S. Detlefsen, Martin Jørgensen, and Søren Hauberg. Reliable training and estimation of variance networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran As- sociates, Inc., 2019. URL http::::ap/p/poppepeeeeeeeeeeeee / neurips wwtpscccpppppp..e.eeeeeee 9/file/ pdf.
Yarin Gal and Zoubin Ghahramani. Dropout Bayesian approximation: Representing model uncertainty in deep learning. In Maria Florina Balcan and Kilian Q. Weinberger (eds.), Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pp. 1050-1059, New York, New York, USA, 20-22 Jun 2016. PMLR. URL https: /ttps:/pproccees pttpp/eeeiiiieec/cc ngs ml press v48 ggggll1 html.
Yarin Gal, Riashat Islam, and Zoubin Ghahramani. Deep Bayesian active learning with image data. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference
10

---

Published as a conference paper at ICLR 2022
Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1183- 1192. PMLR, 06-11 Aug 2017. URL htw.t:tp///// http:s::/ppccceedeieee mlr wttppppes////////11171/. html.
Ali Harakeh and Steven L. Waslander. Estimating and evaluating regression predictive uncertainty in deep object detectors. In International Conference on Learning Representations (ICLR), 2021. URL wwwps:tpp/.......... hwwp///eeeeeeeee..eeeee/e..eeeer net.i/ewweeeeeeeeeeeeeeeeeeeeeee wtt./orumtrrme id=YLewtnvkgR..
José Miguel Hernández-Lobate and Ryan P. Adams. Probabilistic backpropagation for scalable learning of Bayesian neural networks. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 1861-1869, Lille, France, 07-09 Jul 2015. PMLR. URL https:/ / html.
proceedings mlr.. www.prrsssseeee /vsss//v3// hww...pres//33///ereerre sewernandez-lbbattoo
Alex Kendall and Yarin Gal. What uncertainties do we need in Bayesian deep learning for computer vision? In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran As- sociates, Inc., 2017. URL https: http::////ccceeoeeeeoeooooo ngs. neurips httpppapp/////////ppe .dd.
2650d608 a 6d 64 00w.e882288265622222----
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2015.
Diederik P. Kingma and Max Welling. Auto-encoding variational Bayes. International Conference
on Learning Representations (ICLR), 2014.
Armen Der Kiureghian and Ove Ditlevsen. Aleatory or epistemic? does it matter? Struc- tural Safety, 122 2009. ISSN 0167 doi: https://doi.org/I/ http://ddb..gg/0000000f.aaaaaa s0 2008.06.020 URL ht https: /www. rett. com/ www.///ccenne/ertiieeee//eee/eee $0167473008000555 Risk Acceptance and Risk Communication.
Balaji Lakshminarayanann, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https: https:: http://ppocceeee..eeeeeposccsses cc cc/paper 2017 /fle//il/////////////eee
2ee444ffffaaaaeeeeee www.riiin.n................ paper.pds
Jin Han Lee, Myung-Ky Kyu Han, Dong Wook Ko, and II Hong Suh. From big to small: Multi-scale local planar guidance for monocular depth estimation. ArXiv, abs/1907. 10326, 2019.
David A. Nix and Andreas S. Weigend. Estimating the mean and variance of the target probability dis- tribution. In Proceedings of 1994 IEEE International Conference on Neural Networks (ICNN"94), volume 1, pp. 55-60, 1994, doi: 10.1 10.1109//CN.1190 1994.744199
Ian Osband, Charles Blundell, Alexander Pritzel, and Benjamin Van Roy. Deep exploration via bootstrapped DQN. In D. Lee. M. Sugiyama. U. Luxburg, I. Guyon, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 29, Curran Asso- ciates, Inc., 2016. URL https: htttpss//pprceeedii ngs. wwtprriss.cpppppp//eee 6/ file/ .dd.
Cristina Pinneri, Shambhuraj Sawant, Sebastian Blaes, and Georg Martius. Extracting strong policies for robotics tasks from zero-order trajectory optimizers. International Conference on Learning Representations (ICLR), 2021. URL https: atttps://ppenrviiie http:///ppppeeeeeeeeeeeeee wet.ewwwwwweeeeeeee.eee..e.eee forumm?m idd Nc3TJqbe13.
Matthias Plappert, Marcin Andrychowiez, Alex Ray, Bob McGrew, Bowen Baker, Glenn Powell, Jonas Schneider, Josh Tobin, Maciek Chociej, Peter Welinder, Vikash Kumar, and Wojeiech Zaremba. Multi-goal reinforcement learning: Challenging robotics environments and request for research. ArXiv, abss/880.00944. 2018.
Maximilian Seitzer, Bernhard Schölkopf, and Georg Martius. Causal influence detection for im- proving efficiency in reinforcement learning. In M. Ranzato, A. Beygelzimer, K. Nguyen, P. S. Liang, J. W. Vaughan, and Y. Dauphin (eds.), Advances in Neural Information Processing Systems,
11

---

Published as a conference paper at ICLR 2022
volume 34. Curran Associates, Inc., 2021. URL httpss//prroeeiinns.ee...... cc/
Nathan Silberman, Derek Hoiem, Pushmeet Kohli, and Rob Fergus. Indoor segmentation and support inference from RGBD images. In Andrew Fitzgibbon, Svetlana Lazebnik, Pietro Perona, Yoichi Sato, and Cordelia Schmid (eds.), Computer Vision ECCV 2012, pp. 746-760, Berlin, Heidelberg, 2012. Springer Berlin Heidelberg. ISBN 978-3-642-33715-4
Andrew Stirn and David Knowles. Variational variance: Simple and reliable predictive variance parameterization. ArXiv, abs/2006.04910, 2020.
Hiroshi Takahashi, Tomoharu Iwata, Yuki Yamanaka, Masanori Yamada, and Satoshi Yagi. Student-t variational autoencoder for robust density estimation. In Proceedings of the Twenty-Seventh Artificial Intelligence, IJCAI-18, pp. doi: 10.24963/ijcai201 URL hwtps:ttp//cmm /ttt:p:://dd.... http:/oi.oo//ooo/oooo org / 10. 0......244..cc.e. /01..........1.......
International Joint Conference
International
Joint Conferences on Artificial Intelligence Organization, 7 2018.
Marin Vlastelica, Sebastian Blaes, Cristina Pinneri, and Georg Martius. Risk-averse zero-order trajec- tory optimization. In Aleksandra Faust, David Hsu, and Gerhard Neumann (eds.), Proceedings of the 5th Conference on Robot Learning, volume 164 of Proceedings of Machine Learning Research, pp. 444-454. PMLR, 08-11 Nov 2021. URL https: http//poooeeengge////e/////. mlr. wwl..press//v6.. /
wlastelica22a.her html.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y. Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. MOPO: Model-based offline policy optimization. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 14129-14142. Curran Asso- ciates, Inc., 2020. URL ht httpst//pprceeeeing neurips cc/ httpp/ppppr////// file/ 322392ec007/800007777cccbeedddd Papaaaaappppppp... pdf.
12

---

Published as a conference paper at ICLR 2022
APPENDIX
A THE MOMENT MATCHING LOSS
We also investigated an alternative loss function designed to counter the imbalanced weighting of data points by the NLL loss, which call "moment matching" (MM). It uses standard squared error losses to estimate the moments of the target distribution, i.e. directly estimating the sufficient statistics of the target distribution. In our experiments, found that this loss function generally fixes the problems with premature convergence when using LNLL. but it also leads to underestimation of variances and exhibits considerable training instabilities.
To fully describe Gaussian distribution, only the first two moments jr, o² need to be estimated. So, why not directly define losses based on the moment estimators? Starting from the conditional mean of the targets, E|y X], we can define the squared deviation as loss and see that the MSE is an upper bound:
E
X - ((X))) B(Y ((x)) x1 S E (Y ((X))? X] = L
(SI)
Notice that CAM is the standard MSE loss. Thus, learning a mean fit ((xx) can follow the standard mion-distributionaa) regression procedures. Interestingly, due the moment matching viewpoint, can analogously define a loss for fitting the variance (or the second central moment). Variance is defined as E [(Y - xxx))² X As such, by analogy, we can define the following loss: C ==
E [[( x()))² x²) 2 [
To use the same physical unit AM we reformulate it in
terms of the standard deviation as:
S.M = E (VXX- x x
(S2)
Thus, the moment matching loss can be expressed simply by the sum of the two losses: LMM C + 2. We found that using fxM instead of C makes it easier to balance the losses.
Interestingly, our B NLLL subsumes the moment matching loss if we allow for different values of to be used for mean and variance estimation. In particular, using 3-NLL with 8 1 for the mean and B = 2 for the variance results in the same gradients using 2 f + ā f the loss. Note that we did not investigate this connection further, and also did not perform any experiments with different values of 3 for mean and variance.
B ADDITIONAL RESULTS
B. FURTHER CONSEQUENCES OF INVERSE- VARIANCE WEIGHTING
In See. 3.2, we interpreted the inverse-varianc weighting of the NLL training under different data distribution P(X, Y) in which points with high error have low probability of being sampled. A consequence of this is that the distribution P(X, Y) continuously shifts while the model improving its fit. For datasets where the noise level is low on parts of the input space and the underlying function can be modeled to high accuracy. an interesting phenomenon can be observed: even though the training curve for RMSE looks like it indicates convergence (i.c. it flattens out), the model can actually still be adapting with the changing training distribution. on data points that already have such low prediction error relative to the average error that further improvements are virtually invisible in the average error. The actual training progress can revealed using a histogram of the prediction errors a log-scale, in the manner of Fig. 6. Eventually, the training distribution stabilizes when for all data points, either the prediction error reaches the noise level at that point, the variance can not decrease further at that point because it reaches manually set lower bound. If no lower bound on the variance is set, the distribution may never stabilize, leading to training instabilities when the variance reaches close zero
just that the changes happen
13

---

Published as a conference paper at ICLR 2022
B.2 SYNTHETIC DATASET
In Fig. S1, we replicate the experiment from Fig. (training the NLL loss on sinusoidal for 10? updates) on several more random seeds. In order test the dependence of reported issue the optimizer, we repeat the experiment from Fig. 1 but use different optimizers than Adam with 31 0.9, 32 0.999. The results are shown Fig. S2. We find that none of the configurations reaches below an RMSE of 0.1 (the optimal value corresponds to an RMSE of 0.01), indicating that the issue occurring stably across optimization settings. In Fig. $3, we provide the results for the NLL metric on the sinusoidal dataset, complementing the results of Fig. (see discussion in Sec. 5.1).
Target
Target y
Target Y
0.0
0.0
.044
-04
0.4
Input X
Input X
Input X
Figure S1: Repeating the experiment from Fig. I, i.. training with NLL loss for 107 update steps.
The observed behavior stable across different independent trials.
Leaming Rate
0.05
0.25
Leamiing Rate
0.001
0.25
RMSE
0.01
0.20
0.0005
0.20
0.15
0.15
0.005
0.0001
SGD. No Mom.
SGD, Mom Mom 20-9
SGD. Mom. Mom.s00.99 .0.99
RMSProp. 0.9
RMSPppp a 0.99
Adam, Apeis
Adam, Pridd
Figure S2: Using different optimizers train the sinusoidal from Fig. with the NLL loss. The color code indicates the mean RMSE over independent trials per optimizer setting. Black indicates that all trials diverged. Training was done for 10 update steps and used architecture 2 from Table S5, The observed behavior is stable across optimization settings.
(a) B 0 (NLL)
(b) 0.5
(c) 8
Learning Rate
1e-03
Learning Rate
1e-03
Leaming Rate
1e-03
le.04
le.04
le.04
36-05
2 Architecture
Architecture
Architecture
Figure S3: Convergence properties analyzed on the sinusoidal toy regression problem. Same as Fig. 7 but for the negative log-likeliin likelihood (NLL) criterion. Due to the bad mean fit, the original NLL. loss (3 = 0) is also bad for most hyperparameter settings. With > 0.5 good fits are obtained for many settings. Also for 1, which corresponds to MSE for fitting the mean, good uncertainty predictions are obtained with our 3-NLL as testified by the low NLL scores.
B.3 ANALYSIS OF SAMPLING PROBABILITIES ON FETCH-
In Fig. S4, we show how the analysis from Sec. 3.2 transfers to a real world dataset, namely Fetch- PickAndPlace. The figure shows how the distribution of effective sampling probability evolves during training (over a fixed set of training points) and compares that against a proxy "oracle": the distribution of effective sampling probabilities when using the squared residuals from model trained
14

---

Published as a conference paper at ICLR 2022
with the MSE loss. The mismatch between the two distributions demonstrates that optimizing the NLL loss drastically undersamples in comparison to the reference, effectively never sampling some data points. This further corroborates analysis in Sec. 3.2.
Initial
Epoch
Epoch 10
600
600
600
CNLL
Count
400
CMSE (trained)
Count
400
Count
400-
200
200
200
0
10
10
10
10
10
10
10"
10
10
Sampling Probability
Sampling Probability
Sampling Probubility
Epoch 50
Epoch 100
Epoch 200
600
600
600
Count
401
Count
800
Count
400
200
200
200
10
10
10-2
10-
10
10
10
10
10-2
Sampling Probability
Sampling Probability
Sampling Probability
Figure S4: Undersampling behavior of LNLL on Picch-iikAndPlace. The plot shows how the distribution of effective sampling probability evolves over training time, taken over 200 fixed training points sampled the initial epoch. The dashed blue histogram shows the distribution of effective sampling probabilities when using the squared residuals from model trained with MSE loss fMSE- This gives a reference distribution that CNLL should roughly match, taking into account the relative hardness of prediction different samples. The LNLL drastically undersamples compared to the reference (note the log-scale), effectively never sampling some points.
B.4 UCI DATASETS
In this section, we include results for predictive log-likelihood (Table S1) and RMSE (Table S2) for all UCI datasets we evaluated on.
We find that baselines based on the Student's t-distribution (x VAMP and VBEM (Stim & Knowles, 2020)) tend to have better predictive likeliool than 3 3-NLL, although there is also dataset where NI. performs better ("naval"") or where there is no statistically significant improvement ("housing", "kin8m", "wine-red", "wine-white"). For RMSE, LMSE unsurprisingly performs best. Our N-NLL is often on par with LMSE and on par or better than the other baselines, except for "yacht". At the same time, our method is very simple implement (see Sec. D.5) and computationally lightweight compared x VAMP and VBEM, which require the costly evaluation of prior and Monte Carlo sampling at each training step.
B.5 GENERATIVE MODELING WITH VARIATIONAL AUTOENCODERS
We test different loss functions on the task generative modeling using variational autoencoders (VAEs) (Kingma & Welling, 2014). To this end, we parameterize the decoder distribution p(æ: with (p(z), (z)), where the mean μx(2) and variance o² (2) are outputs of neural network. We train the VAE by maximizing the ELBO Eg(z(e) [logp(x p( 2)] (ql( x)) l| p(2)), plugging in different loss functions for log p(z 2). Following Stirn & Knowles (2020). we evaluate the log- posterior predictive likelihood log Ey(a)a) [p(2 =)]. We approximate the expectation using finite mixture of 20 Monte-Carie Carlo samples from q(= 1). To compute the RMSE, we take the mean of that mixture. We compare 3-NN NLL against LMM. LNLL with a fixed variance 1, Student-t (Takahashi al., 2018; Detlefsen et al., 2019), XVAMP and VBEM (Stirn & Knowles, 2020).
We evaluate on MNIST and FashionMNIST.. Table S3 presents quantitative results. VBEM achieves the best reconstruction error the expense of poor log- likelihood. Vice versa, Student-t, VAMMP x VAMP* and * achieve strong log- likelihoods but worse reconstruction errors. B-NLL with 0 provides good compromise between ligliiell and reconstruction error. Figure $5
15

---

Published as a conference paper at ICLR 2022
Table S1: Results for UCI Regression Datasets. Predictive log-likelihood (higher is better) and standard deviation, together with dataset size, imput and output dimensions. Best mean value in bold. Results that are not statistically distinguishable from the best result marked with f.
carbon (10721, 5, 3)
concrete
energy
housing
(1030, 8, 1)
(768, 8, 2)
(506, 13, 1)
(3 0)
11.36 = 2.11
-3.25 = 0.31
-3.22 = 1.41
-2.66 = 0.50
L -.NNLI (3
(3- = 0.25)
10.91 + 2.42
-3.31 +: 0.51
-2.82 +: 0.82
-2.75 1: 0.42
L (5 =0.5)
10.22 4.00
-3.29 + 0.36
-2.41 0.72
-2.64 0.36
L3 3- NLL =0.75)
10.82 1.37
-3,27 1 0.34
-2800 0.59
-2.72 0.421
Cg (3 = =1.0)
3.31 20.88
-3.23 = 0.33
-3.37 0.58
-2.85 = 0.87
LMM
4.72 + 5.80
-3.49 + 0.38
-4.26 0.50
-3.42 + 1.01
MSEE
Student-t
15.59 + 0.43
-3.07 0.14
-2.46 0.34
-2.47 + 0.24
XVAMP
13.28 1 0.19
3.1511..15
-2.47 0.32
-2.43 1 0.21
XVAMP*
13.17 + 0.26
-2.41 + 0.32
-2.45 + 0.221
VBEM
5.68 + 0.70
+..14
-4.99 = 0.16
-2556 + 0.15
VBEM
13.23 + 0.36
-2.99 1 0.13
-1.91 0.21
22.4 0.22
kin8m (8192, 8,
naval (11934, 16, 2)
power
protein (45730, 1)
(9568, 4, 1)
C3 NLL (3 0)
1.140 + 0.039
12.46 + 1.18
-2.807 + 0.057
-2.80 + 0.05
C NLL (3 = 0.25)
1.142 + 0.026
13.78 + 0.33f
-28001 + 0.053
-2.99 + 0.05
L NLL (5 = 0.5))
1.141 = 0.046
13,99 = 0.40
-2.805 0.052
-278
L. NLL (3 = =0.75)
NLL
1.137 + 0.041
13.63 0.62f
-28006 0.053
+
C. (5 =1.0)
1.126 0.041
13.59 0.30
-2.110 + 0.051
2
LMM
0.999 + 0.063
12.73 1 0.64
-2.918 + 0.092
-2.98 1 0.06
MSE
Student-L
1.155 0.037
12.47 0.48
-2.738
0.026
2256 0.02
XVAMP
1.147 + 0.037
12.44 + 0.60
-2.788
+ 0.032
-2.74 + 0.03
XVAMP*
1.147 = 0.036
12.80 + 0.55
-2.785 =
-2.71 + 0.02
VBEM
1.019 1 0.054
8.05 0.13
-2.119 0.018
-2.85 0.01
VBEM*
1.138 + 0.036
13.10 0.47
2.783 0.031
-2.33 0.01
superconductivity (21263, 81, 1)
wine-red (1599, 11, 1)
wine-white
yacht (308, 6, 1)
(4898, 11, 1)
L3 - NL.
-3.60 1 0.23
-1.03 1: 0.24t
-1.059
0.074
-2.86 1. 5.18
C9 - NLL (3 0.25)
-3.56 0.14
-0.98 = 0.12
-1.041 + 0.064
-1.97 + 1.14
C9 NLL (3 = -..5)
-3.60
-1.036 0.065
-2.47 1.68
C. NII = =0.75)
-.11
-1.02
-1.039 1: 0.060
-1.87 1: 0.55
L (B = 1.0)
3.83
0.97
-1.040 + 0.0671
-2.27 1.07
LMM
4.45 0.39
-1.2 + 0.43
-1.135
0.093
-11.24 31.03
AASE
Student-t
-3.8 1 0.04
-0.94 1 0.10t
-1.034 1: 0.062
-1.23 1: 0.55t
xVAMP
-3.40 + 0.03
-0.55 + 0.06
-1.038 + 0.050
-0.99 + 0.33t
XVAMP*
33.4 0.04
-1.029 + 0.048
-1.04 + 0.471
VBEM
-3.63 1 0.09
-0.95 1 0.07
-1.028
0.048
-..5 0.10
VBEM*
-3.40 + 0.04
-0.94 + 0.07t
-1.031 0.0571
-0.98 0.24
16

---

Published as a conference paper at ICLR 2022
Table S2: Results for UCI Regression Datasets. RMSE (lower is better) and standard deviation, together with dataset size, input and output dimensions. Best mean value in bold. Results that are not statistically distinguishable from the best result are marked with f.
carbon (10721, 3)
concrete (1030, 8, 1)
energy (768, 8, 2)
housing (506, 13, 1)
C3-NLL(8 0)
0.0068 = 0.0029t
6.08 = 0.65
2.25 + 0.34
3.56 = 1.07
L3 (B = 0.25)
0.0069 + 0.00281
5.79 + 0.74
1.81 + 0.30
3.48 + 1.15t
(8 = 0.5)
0.0068 1: 0.00291
5.61 0.65
1.2:::025
3,42 1: 1.04
NB = =0.55)
0.0069 + 0.00281
+
1.31
3.43 1.07t
(S = -.0)
0.0073 = 0
5.77
1.54
3.50
AMM
0.0097 = 0.0034
6.28
2.19
4.02
LSSE
0.0068 = 0.0028
4.96 0.64
0
3.24 1.08
Student-t
0.0067 + 0.0029t
5.82
2.26 - 0.34
3.48 + 1.171
X VVMMP
0.0067 + 0.00291
5.44
1.87 -.32 0.32
3,23 1.00
VVMPP*
0.0067 1: 0.00291
5.35 1: 0.73
2.00 1 0.26
3.38 1 1.15
VBEM
0.0074 = 0.0026
5.21 = 0.581
1.29 = 0.33
3.32 = 1.06
VBEM
0.0067 + 0.0029
5.17 0.59t
1.08 + 0.17
3.19 1.02
kin8m (8192, 1)
naval (11934, 16, 2)
power (9568, 1)
protein (45730, 1)
Lg NLL (8 = 0)
0.087 + 0.004
0.0021 + 0.0006
4.06 + 0,18
4.49 : 0.11
L. NLL = 0.25)
0.083 0.003
0.0012 = 0.0004
4.04 + 0.18
4.35 - 0.05t
L NLL = -0.5)
0.082 0.003
0.0006 = 0.0002
4.04
4.31 0.02
L. 3- (3 -0.75)
0.081 0.004
0.0004 0.0001
4.04 0.15
4.28 0.02
C9 (8 1.0)
0.081 0.003
0.0004 + 0.0000
4.06
131
LMM
0.082 0.003
0.0005 1 0.0001
4.07
4.32 0.071
LMSE
0.081 = 0.003
0.0004 = 0.0001
4.01
4.28 0.07
Student-t
0.085 + 0.005
0.0026 + 0.0009
4.76 + 0.24
XVAMP
0.081 + 0.003t
0.0023 + 0.0004
4.03 0.17
4.38 + 0.05t
XVAMP*
0.082 0.003
0.0020 + 0.0006
4.03 1 0.18
4.31 4 0.02
VBEM
0.082 0.003
0.0009 = 0.0004
+ 0.15
4.31 0.01
VBEM
0.082 0.004
0.0015 +1 0.0005
4.02 + 0.18
4.35 0.09
superconductivity (21263, 81.
wine-red (1599, 11,
wine-white
yacht
(4898, 11, 1)
(308, 6, 1)
(B = 0)
13.87 + 0.50
0.636 + 0.038
0.691 + 0.032
1.22 + 0.47
L9 (8 = 0.25)
13.50 + 0.49
0,638 1 0.036
0.687 + 0.039
1.73 + 1.00
L. (8 = 0.5)
13.02 0.47
0.635 0.037
0.685 = 0.035
2.35 1.44
L3 (3 = 0.75)
$32000..44
0.638 = 0.035t
0.689 = 0.034
1.97 1.03
(3 1.0)
13.42 0.41
0.639 +
0.684 0.031
2.08 1.13
LMM
13.68 0.79
0.652 0.044t
0.692 0.032 0.032
3.02 1.38
LMSE
12.48 + 0.40
0.633 + 0.036
0.684 0.038
0.78 0.25t
Student-t
13.52 0.60
0.636 0.038
0,688 0.036
1.34 0.63
XVAMP
13.33 + 0.52
0.635 0.035
0.691 + 0.032
0.99 0.43
x XVAMP
13,42 + 0.59
0.633 + 0.035
0.685 + 0.032
1.13 - 0.66
VBEM
12.72 0.57t
0.639 + 0.041
0.685 + 0.035
1.66 0.84
*BEM
13.15 0.43
0.633 1 0.040
0.686 0.036
0.65 0.20
17

---

Published as a conference paper at ICLR 2022
shows qualitative examples. Our 3-NLL loss with 3 0 allows to learn good reconstructions and meaningful uncertainties. Moreover, images produced by sampling latents from the prior are semantically meaningful, indicating that adding loss function does not break the disentangling properties of VAFs. Compare that to Student-t, x x and VBEM(*), which do not produce similarly clear images when sampling from the prior.
Table S3: Results for generative modeling with variational autoencoders on MNIST and Fashion- MNIST. We report RMSE and posterior predictive log-likelihood (LL) with standard deviation over independent trials.
MNIST
Fashion-MNIST
Loss
RMSE 4
LL
RMSE
LL
LNLL (a² =1)
0.153 + 0.002
-730 +
0.143 + 0.002
-729 + 0
C3-NLL(3
0.237 + 0.002
2116 + 55
0.170 + 0.001
1940 + 104
Ca CA-NLI =0.25)
0.181 + 0.004
2511 88
0.140 0.002
1
CALN =0.5)
0.151 1 0.003
2220 1 25
0.125 1 0.003
1639 1 52
(3 = =0.75)
1.0001
1
0,131 0.000
1331 32
(B 1.0)
0.132 0.001
1
0.138 0.002
1142 26
LMM
385
0.000
-151
Student-t
4291 103
0.182 0.002
2857 9
xVAMP
0 0.002
2989 + 268
0.161 0.001
2158 + 100
XVAMP*
0.225 + 0.001
215
0.160 0.002
2150 + 131
VBEM
0.114 + 0.001
719
0.108 0.000
660 +
VBEM
0.176 + 0.008
3213
2244 0.003
B.6 DEPTH REGRESSION
We evaluate 3-NLL on the task depth regression on the NYUv2 dataset (Silberman et al., 2012). For this purpose, we use a state-of-the-art method for depth regression, AdaBins (Bhat et al., 2021). and train it with different loss functions. Note that we remove the Mini- Mini-ViT Transformer module from the model, thus our results are not directly comparable with those reported by Bhat et al. (2021).
Table $4 presents quantitative results. N-NLL with > achieves better RMSE than the NLL loss. B-NLL with 3 = 0.5 again provides a good trade-off, achieving similar RMSE as CMSE and performing better on some the other metrics. Figure S6 shows qualitative examples. Compared to LNLL, the depth maps predicted by N-NLL with B > O are noticeably sharper.
Table $4: Results for depth regression on the NYUv2 dataset (Silberman et al., 2012). We adapt a state- of-the-art network for depth regression from AdaBins (Bhat al., 2021) and train it with different loss functions. In contrast to AdaBins, our network does not include the Mini-ViT Transformer module and thus our results are not directly comparable with those originally reported. For reference, we also report numbers from other recent literature on this task. Notably, the Gaussian NLL in all variants clearly outperforms the Scale Invariant (SI) loss, despite the latter being loss function specifically designed for the task of depth regression. We refer the reader to Bhat et al. (2021) for description of the metrics.
Method
82
63 f
REL 4
RMSE |
log1o
LL 1
La NLA. (3 0)
0.8855
0.9796
0.9959
0.1094
0.3854
0.0462
-4.52
Ca NL1. (8- = 0.25)
0.8887
0.9812
0,9956
0.1081
0.3818
0.0458
-8.15
C BBL 0.5)
0.8885
0.9813
0.9956
0.1093
0.3789
0.0458
-7.50
(B =0.75)
0,8902
0.9804
0.9952
0.1095
0.3800
0.0462
-7.35
(B 1.0)
0.8872
0,9813
0.9958
0.1088
0.3845
0.0467
-5.10
LMSE L1
0,8890
0,9806
0,9960
0.1086
0.3776
0.0461
0.8877
0.9798
0.9955
0.1073
0.3850
0.0459
SI Loss (Bhat et al., 2021)
0.881
0,980
0,996
0.111
0.419
AdaBins (Bhat et al., 2021)
0.903
0.984
0.997
0.103
0.364
0.044
BTS (Lee et al., 2019)
0.885
0.978
0.994
0.110
0.392
0.047
DAY (Chen et al., 2019)
0.882
0.980
0.996
0.108
0.412
18

---

Published as a conference paper at ICLR 2022
MNIST
Fashion-MNIST
dataset
A
2
3
mean
7
2 2
7
9 9 7 7
(3-NLI(3=0)
9
7
B R
Ø
pusterior
f 2
8-
A
E
&
(0,
0
9
6
4
@
3.
2
/
C3-NLL = 0.25)
std
7
2.
7
/
9
posterior
7
2.
7
/ 9
1 8
9 6
(0,
0
5
6
Ø
mean
a
2
/
9
(3-NLL = 0.55)
std posterior
7
7
/
9
9 6
(0, 1)
7 O
7
6
Q
mean
2
2
/
4
C3-NLI 0.75)
std posterior
2
7
2 3
/
(O, 1)
7 8
5
J
mean
7
2
7
2
/
9
9
3
C3 1.0)
std posterior
2
2
2
2.
2
/
(O,
5
7
8
9
9
mean
7 3 7
/
4
9
LNLL (0²=1)
std posseriir
~ NO 1)
5
0
mean
LMM
pusterior A' 1)
19
6: 3/4 18
87
0
2
9
1
Student-t
26
7
6
posterior
7 3
2
7
AO 1)
9
9
3
9
5
5
mean std posterior
0
2
8
X XVAMP
7
26
f
2
(0, 1)
5 0
2
4
meam std posterior
7 3.
2
/
3
XVAMP*
3
I
AO 1)
5 5
Z
mean
2 7
2
/
R
9
3
VBEM
std posterior
1 7
3
2
/
(( 1)
f
C
R
3
mean
7 2
7
/
1
9
3
VBEM
std
2
posterior
2
/
(O
5 0
0
9
Figure S5: Generative modeling with variational autoencoders on MNIST and Rashion-MNIST The overall first row shows inputs from the test set. For each method, we present posterior predictive means (i.e. reconstructions), the posterior predictive standard deviations, samples from the posterior predictive distribution, and finally samples using the prior A (0, 1). (NLL(O² = 1) refers to Gaussian log-likelihood with a fixed variance of 1. Values are clipped to the interval [o, 1]. Examples are not cherry-picked.
19

---

Published as a conference paper at ICLR 2022
Input
Ground Truth
pred.
C3-NLL (8 = 0.0)
std.
pred.
L3-NLL (3= 0.5)
std.
pred.
CB3NNN NL.L
(3 = 0.75)
std.
pred.
C NIL.
(3 1.0)
std.
LMSE
pred.
Figure S6: Example results for depth regression on the NYUv2 dataset (Silberman et al., 2012) First rows show input image and ground truth depth map, where black values in the depth map represent missing values. For each method, we present the predicted depth map (pred.), and the aleatoric uncertainty form of the predicted standard deviation (std.). Results for L3-NLL with > are noticeably sharper. The last column shows magnified view on the previous picture.
20

---

Published as a conference paper at ICLR 2022
C DATASETS AND TRAINING SETTINGS
Sinusoidal without heteroscedastic noise This dataset is created by taking 1 000 uniformly spaced points on the interval [0, 12] as inputs x and applying the function y(a) = sin sin (2Tx) + 5 to them to create the targets y. where & is Gaussian noise with a standard deviation of 0.01.
Sinusoidal with heteroscedastic noise We use the synthetic data as introduced in Detlefsen et al. (2019). From the functional form = 2 sin(() + zςi + f2. with Gaussian noise with standard deviation a 0.3 for { and {2. we sample 500 points uniformly spaced in the interval [0, 10). The model is an MLP with one hidden layer of 50 units and tanh activations (as used in Detlefsen et al. (2019) and Stirn & Knowles (2020)).
UCI Datasets We use the UCI datasets suite commonly used to benchmark uncertainty estimation, stemming from the UCI Machine Learning Repository. In particular, we use the training-test protocol from (Hernández-Lobee Lobatto & Adams, 2015; Gal & Ghahramani, 2016), and their data splits", except for "carbon", "energy", "naval", "superconductivity"" and "wine-white", where we generate our own random splits.
Inputs and targets are whitened on the training set. Metrics are reported in the original scale of the data. Each dataset is divided into 20 randomly sampled train-test splits (80%-20%). For each split, we further randomly divide the training set into 80% training data and 20% validation data and search for optimal learning rate from the set [10-+,3- 100 7 10 10-3,3 3 100 7. 10 $} by monitoring log- likeliiodoo on the validation set. We train for maximum of 200 updates, except for the larger "kin8m", "power plant", "protein", and "naval" datasets where we train for a maximum of 100 000 updates. We perform early-stoppin with a patience of 50 epochs, retrain the model with the best found learning rate the full training set, and then evaluate on the test split. The reported performance and standard deviations are taken as averages over all test splits. Note that the performance we report is not comparable with other publications, performance is known to differ strongly over different data splits. Some other works also perform early-stopping on the test set, which distorts the results.
We use a single-layer relu hidden network with 50 neurons, except for "protein" where we use 100 neurons. The batch size is 256.
Following Stirn & Knowles (2020). for each method, we report the number of datasets for which the method is statistically indistinguishable from the respective best method in Table (Ties). For this purpose, we performed a two-sided Kolmogorov-Smirno test with p S 0.05 significance level.
ObjectSlide This environment consists of an agent whose task is to slide an object to target location (Seitzer et al., 2021). The continuous state space consists of dimensions: agent and object positions and velocities, and the continuous action space is one-dimensional movement command. The forward prediction task consists of predicting the change in object position in the next state from the current state and action. The dataset we use consists of 180000 transitions collected using a random policy, which we split into training, validation, and testing sets with 60 000 transitions each. Inputs and targets are whitened on the training set. Metrics are reported in the original scale of the data. We train for maximum of 5 000 epochs with batch size of 256, and evaluate the model with the best validation log-likelihood on the test set afterwards.
Tetch-PickAndPPace We use the OpenAl Gym (Brockman al., 2016) challenging real-world scenario. The task of the agent is to use position-controllee 7 DoF robotic arm to lift an object to a target location in space. The state space is 25-dimensional and the action space is 4-dimensional, As in ObjectSlide, the prediction task is to predict the 3-dimensional change in object position from the current state and action. We use 840 000 transitions collected using the APEX method (Pinneri et al., 2021) as our dataset, which we split into 70% training, 15% validation, and 15% testing data. Inputs and targets are whitened on the training set. Metrics are reported in the original scale of the data. We train for a maximum of 500
environment (Plappert et al., 2018) from
hittps://archiee www.ciiiiie.......eee.... avvailable under httpp:://gtthhib.cmm http:///mmmnitaee//gggg./..o.//o http://htw...oooooooooo///////ta
21

---

Published as a conference paper at ICLR 2022
epochs with batch size of 256, and evaluate the model with the best validation log- logeliieihood on the test set afterwards.
NYUv2 Depth Regression We use the dataset in the variant provided by Lee et al. (2019). Training settings and evaluation protocol were taken from Bhat et al. (2021). We train for 25 epochs using a batch size of 16 and validate the model every 100 updates. We use the model with the best "REL" metric on the validation set for testing.
D HYPERPARAMETER SETTINGS AND IMPLEMENTATION DETAILS
For all experiments, we used the Adam optimizer (Kingma & Ba, 2015) with standard settings 31 = 0.9, B2 0.999. We parameterize the Gaussian distribution using linear layers top of shared features produced by an MLP. The variance o² (a:) is constrained to the positive region using the softplus(z) (z) log(1 exp(x)) activation function. We additionally add small constant of 10- to prevent the variance from collapsing to zero and clamp the maximum variance to 000.
Some baselines a Student's t-distribution (Student's I, VAMP(*), VBEM(()) as their predictive distribution. This distribution results from integrating out the unknown variance of Gaussian with learned Gamma prior on the inverse variance (Detlefsen et al., 2019; Stirn Knowles, 2020). We parametrize the Gamma distribution in terms of data-dependent alpha and beta parameters, i.e. a(x) and (x), which are computed using linear layers on top of the shared features. In this case, the MLP has three outputs: mean, alpha, and beta. Both alpha and beta are constrained to the positive region using the softplus activation. We add a positive constant of 1.001 for alpha and 10- 1008..0.00 for beta. Alpha is clamped to maximum value of 1 000 and beta to 10-8 8 999. These values were chosen such that the resulting variance matches the range (10-5 8 1000] while ensuring that the degrees-of of freedom oreeees-of-rreedm parameter of the Student's is always greater than 2.
0.001
For xVAMP and VBEM, the MLP outputs a fourth term, ग (z). representing the logits a categorical distribution that specifies the mixture weights of the prior. We initialize the prior parameters exactly the same as Stirn & Knowles (2020). For these methods, the objective function additionally contains KL divergence between a Gamma distribution and unixture-of-Gamen distributions. Following Stirn & Knowles (2020), we approximate this KL divergence using 20 Monte-Carlo samples.
D.1 SINUSOIDAL REGRESSION PROBLEM
The sinusoidal fit in Fig. 1 results from a network of two hidden layers with 128 neurons per layer and tanh activations, optimized with a learning rate of 5 10-4 and batch size of 100. For the experiment Sec. 5.1, we sean over learning rates and architectures with different hidden layers and units per layer, detailed in Table S5.
Table S5: Architectures used for the sinusoidal regression task. Fully-connected feed-forward neural networks with tanh activations.
Architecture #
0
1
2
3
# Hidden Layers
2
2
2
3
# Units per Layer
32
64
128
128
256
D.2 OBJECTŠLIDE AND PETCH-PICKANDPPLCC
For each tested loss function, we performed a grid search on the ObjectSlide and Fetch-PickAndPlace datasets. We report the parameters we scanned over in Table S6. Table S7 reports the model configurations with the best validation likelihodo on the grid search. For the results in Table we retrained the best model configuration with five different random seeds and evaluated them on the hold-out test set.
available under hnttps://giie cowccooppppc/cccppccccccc cogap1 ex /bts
22

---

Published as a conference paper at ICLR 2022
Table S6: Hyperparameter settings for our grid search on the ObjectSlide and Fetch-PickAnndPlce datasets. We run 96 configurations per loss function.
Hyperparameter
Set of Values
Learning Rate # Hidden Layers
{3 10 10 3- 10 10--10----0 -10
{2. 22.3,44 3]
# Units per Layer
{128. 256, 386, 512]
Activation
{tanh, relu}
Table S7:
Best hyperparameters found by grid search on ObjectSlide and Fetch-PickAndPlace datasets, measured by best log-like bog-likelhhood on the validation set.
(a) ObjectSlide
(b) Fetch Bicch-PickAndPlcee
Method
LR
Layers
Act.
Method
3
LR
Layers
Act.
LMSE
10
3 x 128
relu
LMSE
10
4 x 128
relu
ENLL
10-3
3 x 128
relu
CNLL
3 10
x 128
relu
C NLL
0.25
10-3
3 x 128
relu
L NLL
0.25
3- 30-4
4 x 128
relu
CB NL..
0.5
10-3
3 x 128
relu
C --NNL
NLI
0.5
3-
4 x 128
relu
Cg NLL
CLL
0.75
10
3 x 128
relu
LLLNLL NLL
0.75
10
x 128
relu
C NLL
1.0
10-3
x 128
relu
CL-NLL
1.0
10-3
4 x 128
relu
LMM
10-3
3 x 128
relu
PMM
10-3
4 x 128
relu
Student-t
10-3
2 x 386
relu
Student-t
3.10-
3 x 256
relu
XVAMP
10
128
relu
XVAMP
10
3 386
relu
XVAMP*
10-4
3 x 256
relu
XVAMP*
10-1
386
relu
VBEM
3 10-4
2 x 256
tanh
VBEM
10-3
3 x 386
relu
VBEM
10-3
2 386
relu
VBEM
10-4
3
relu
D.3 VARIATIONAL AUTOENCODERS
largely follow Stirn Knowles (2020) for their training settings for the VAE experiment. In particular, we use an encoder with three layers of 512, 256, 128 neurons and decoder with three layers of 128, 256, 512 neurons, all with relu activations. The latent space is 10-dimensional for MNIST and 25-dimensional for FaSHi We train the VAEs for maximum of I 000 epochs, using Adam with a learning rate of 0.0003 and a batch size 256. Early-stopping with a patience of 50 epochs is performed the aggllikllhood of the validation set. The validation set consists of 20% of the MNIST/FashionMNISS training set.
D.4 DEPTH REGRESSION
We use the official implementation of AdaBins (Bhat et al., 2021), " thereby reproducing their exact training settings and evaluation protocol. We remove the AdaBins/mini-Viii Transformer from the model. Instead, the feature map output by the U-Net is reduced to two channels using 1 x 1 convolution, where we use the first channel the mean predictor and the second channel as the variance predictor. In this setting, both mean and variance are constrained to positive numbers by a softplus activation. On top of that, we add positive offset to ensure a minimum output value of 10-3 for the mean (the minimum possible depth value) and 10-6 for the variance and clamp both mean and variance to a maximum value of 10.
23/Adanin
23

---

Published as a conference paper at ICLR 2022
D.5 IMPLEMENTATION BETA-NLL PYTORCH
def @ata_nli_loss (mean, variance, target, beta): :
**mmpptt beta-Nll. loss
:param mean: Predicted mean of shape x :param variance: Predicted variance of shape x :param target: Target of shape B D
5
:param beta: Parameter from range [0, 1] controlling relative
weighting between data points, where '0 corresponds
high weight on low error points and '1' to an equal weighting.
10
:returns: Loss per batch element of shape B
loss 0.5 ((target mean) / variance variance. ())
13
14
if beta > 0:
15
loss = loss + variance. detach ** beta
less. (axis=-1)
24


Reviews:
[{"review_id": "ym9GfU701tq", "reviewer": "Reviewer_5Jtp", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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": "### Strengths:\n- The NLL loss scaling \"the gradient of badly-predicted points, effectively under sampling those points\" is an interesting perspective I have not yet seen. The conventional focus, based on my reading of the literature, is the infinite gradient problem (a perfect mean estimate is impossible to obtain using NLL because variance tends towards 0). \n\n### Weakness:\n- Figure 1 attempts to illustrate the problem of fitting a neural network to a noisy sinusoid using a NLL loss. However, a few things could be improved. Can the author(s) please confirm if the left panel is single run? If so, does rerunning the code under different initialization move the region with poor predictive mean estimates around? Lastly, it probably would be more clear to show the MSE predictive means in addition to its low RMSE (yes, RMSE is low enough to suggest a perfect fit, but it reduces a reader's burden).\n- Claiming the identification of \"high dependence of the gradients on the predictive variance as the primary culprit\" for NLL failures as a contribution is quite bold. Takahashi, et. al (2018) and Stirn, et al (2020) both discuss this issue, but admittedly with a different perspective--neither claims it as a contribution.\n- \"We observe a drastic difference between the NLL loss and the MSE/moment matching loss.\" I understand that the the MM loss for the mean is the same as NLL, but the MM loss also has a variance loss term. How then is 6.b both MSE and MM?\n- The second paragraph of section 4.2 mentions not yet seen experimental results. I recommend moving discussion of results to appear after the experiments.\n- The author(s) seemingly contradict themselves several times as to whether MSE is a fundamentally different loss from their $\\beta = 1$.\n\t1) \"For β = 1 the gradient w.r.t. μ in Eq. 10 is equivalent to the one of MSE. However, for the variance, the gradient in Eq. 11 is a new quantity with 2σ2 in the denominator.\"\n\t2) \"As expected, the same holds for the mean squared error loss (MSE) which is recovered for β = 1.\"\n\t3) The results (figure 8 and table 1) cast  ($\\beta = 1$) as MSE, but their loss term has an added variance term.\n- The author(s) should have MSE in their loss table since I don't believe ($\\beta = 1$) is equivalent. Please correct me if I am missing something!\n- The author(s) do not compare to ANY of the cited works, which often outperform them on the UCI regression tasks.", "summary_of_the_paper": "The author(s) attribute the under-fitting of predictive means by neural networks that parameterize heteroscedastic Gaussian likelihoods to 1) initial inability to break symmetry and 2) Gaussian negative log likelihoods tendency to down weight poorly predicted points. The authors note the latter effectively under samples those points with poor predictive means, thus inhibiting their fit in a rich-get-richer scheme. The author(s) propose two alternative loss functions: moment matching (MM) and $\\beta$-exponential variance estimate with $\\beta \\in [0,1]$. The author(s) demonstrate their $\\beta$-exponential generally outperforms NLL ($\\beta = 0$) and MSE.", "summary_of_the_review": "I think this paper offers a fresh perspective on NLL loss failures as I mentioned in my main review. However, their only support for this a contrived toy example over a single run. The paper would be much stronger if:\n- The author(s) could apply their analysis to identify either of the two conditions from section 3 in a real world data set.\n- Their proposed methods offered comparable (using simpler methods) or superior performance to related works.\nUnfortunately, without either of these, the paper's potential significance remains limited.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "J0JPAndCsk", "reviewer": "Reviewer_gjks", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "The problem of optimizing Gaussian likelihoods is ubiquitous and, thus, understanding its pitfalls thoroughly is important. While the phenomenon of variance weighting is known, the authors provide a nice analysis with explanations, visualizations, and derivations of  high quality. The claims they make in Secs. 3 and 4 are all well supported and fleshed out nicely. The presentation is very clear and well-organized. I really enjoyed reading the paper and found it quite insightful. The authors provide a pragmatic and easy-to-implement method to alleviate the identified problems in a situation-specific manner (by adjusting $\\beta$). \n\nUnfortunately, I feel the experimental evaluation is lacking:\n- The results in Tab. 1 and Tab. S1 seem to be not significant according to the reported standard errors? How many seeds were evaluated? Tab. S1 should be moved to the main part (assessing the MSE values is central to judging the effectiveness of the proposed method).\n- Why is the RMSE of MM in Tab. 1 worse than that of $\\beta$-NLL? It is optimizing directly for the mean and thus should yield better results, shouldn't it? Did you use two-headed networks? If so, this might explain the worse MSE of MM. But then, I'm still wondering why the MSE in the dynamics experiments (Tab. 2) is worse than $\\beta$-NLL as here you train on the pure MSE loss (not MM). Why don't you evaluate pure MSE on the experiments in Tab. 1/S1?\n- While I agree that the proposed method visually seems to yield better calibrated uncertainty estimates than MM in Fig. 8, I would also like to see a quantitative evaluation here.\n- The tasks are low-dimensional (1D, 3D). I would be interested in results on more complex, higher dimensional problems.\n\nMinor points not influencing my score:\n\n- Shouldn't Eq. 6 contain a factor of $Z$?\n- Sec. 5.1. misses a reference to Fig. 7.\n- On p. 8, paragraph \"Sinusoidal with heteroscedastic noise\", $\\xi$ should read $\\xi_1$ and $\\xi_2$.\n- Fig. 9: $\\beta$ is not rendered in the labels.\n- Sec. 6: wrong full stop before \"we reveal the underlying reasons\".", "summary_of_the_paper": "The authors analyze the known phenomenon that the optimization of the Gaussian likelihood of a probabilistic neural network with heteroscedastic output variance using gradient ascent can get stuck at solutions with a suboptimal mean fit, compensated by large output variance: the likelihood amplifies non-uniform initial distributions of feature granularity due to variance weighting, i.e., the gradient w.r.t. the mean are scaled by the inverse variance leading to well-fitted samples dominating the gradient. The authors argue that this effect can be undesirable because it can prevent spending model expressivity on hard-to-fit regions, while it can be desirable in other situations because it allows the network to ignore outliers arising due to noise. This is complementary to the MSE objective which is dominated by badly-fitted samples (and thus hard to fit regions) but sensitive to outliers, and which does not allow to fit the output variance. The authors propose a modified loss function, $\\beta$-NLL, which interpolates between the likelihood and the MSE loss functions with the incentive to retain the \"best of both worlds\". They compare $\\beta$-NLL with standard NLL, MSE, and a moment-matching version on several datasets.", "summary_of_the_review": "The authors provide a very nice discussion and analysis of the known phenomenon of variance weighting arising in SGD on NLL loss of Gaussian likelihoods. While I judge this to be an interesting contribution on its own, I would like to see an improved experimental evaluation. Indeed, I am not yet fully convinced that the proposed $\\beta$-NLL is effective. In its current state I rate the submission as borderline but I will raise my score if the authors improve upon the mentioned points during the rebuttal.\n\n------------------------------\n\n**Update after rebuttal:**\nI thank the authors for their effort to improve the experimental evaluation. I raise my score to 6.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "365iZ3xRWWj", "reviewer": "Reviewer_5nZv", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "**Strong Points**\n\n- The discussion of the optimization issue caused by the NLL is very clear with nice synthetic examples.\n- The paper is well written and mostly easy to follow\n- Table 2 shows clear improvements in performance over the standard NLL approach\n\n\n**Weak Points**\n\n- Section 3.1 is difficult to follow. Figure 4 suggests the measurement in Equation (5) is important, but there is a lack of intuition for its significance.\n- No ablation study comparing other approaches to solving this problem.\n\n\n**Questions**\n\n- I wonder how optimizer dependent this issue is. Which optimizers were used and how were they tuned?\n- There is a long line of work about “critically” initializing neural networks so that their curvature grows exponentially with depth. I wonder if this is related to Section 3.1. The authors might consider looking at Poole et al., “Exponential expressivity in deep neural networks through transient chaos” and Raghu et al., “On the expressive power of deep neural networks.”\n- What is the noise distribution assumed by the beta-NLL? I understand gradients are not taken with respect to the beta-NLL, but I’m curious whether part of the issue is that the Gaussian assumption on the noise is inappropriate.\n- How is the bolding chosen in Table 1. It seems like there’s significant overlap between the performance of the different methods.\n\n\n**Additional Feedback**\n\n- The notation in Equation (9) is confusing. The expectation is over the random variable X, but X also appears outside of any expectation in the stop-gradient operation.\n", "summary_of_the_paper": "**Summary**\n\nFor regression tasks with heteroscedastic noise, a Gaussian model whose mean and variance are a function of the input is frequently used. Neural networks are then used as function approximators for the mean and variance and are fit using the negative log-likilhood. The paper identifies a problem with this approach: optimization can get stuck in configurations where the predicted mean is far from the true mean, as this is compensated for by a high predicted variance that also stalls learning. By adjusting the NLL objective, the authors compensate for the problem.", "summary_of_the_review": "**Recommendation**\n\nI like aspects of this paper, but there are some major outstanding questions. Currently, it does not meet the bar for acceptance.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper examines the approach of modeling aleatoric uncertainty by fitting a neural network, that estimates mean and variances of a heteorscedasitic Gaussian distribution, based on log likelihood maximization. The authors identify the problem that gradient based training on the netgative log likelihood (NLL) may result in suboptimal solutions where a high predicted variance compensates for the predicted mean being far from the true mean. To solve this problem, the authors suggest to adjust the log likelihood objective by weighting the log likelihood of each single data point by the corresponding beta-exponentiated variance estimate. This adjusted objective is referred to as beta-NLL.

All reviewers agreed that the identified problem and the proposed solution are interesting, that the paper is well written and organized, and that the contributions are significant and somewhat new. The main criticism was on the side of the empirical evaluation. It was criticized that the empirical analysis did not compare the proposed method to other approaches to solving the same problem, that the identified problem and the proposed method should be also analyzed on high-dimensional data, that the results on the synthetic experiments could be improved by investigating more than a single run and by incorporating the the MSE in corresponding Figure 1,  and that standard errors were not reported. 

Based on the reviews the authors added several new experiments and investigations in the revised version of their manuscript to improve their empirical analysis: 1) new experiments on high-dimensional data sets were conducted applying variational autoencoders on MNIST and Fashion MNIST and performing Depth-map prediction from images on the NYUv2 dataset. 2) For comparison several baseline methods were added to the experiments on the UCI and the dynamics datasets. 3) Three more UCI datasets (carbon, superconductivity, wine-white) were included in the empirical analysis. 4) An evaluation of calibration of predictive variance for the heteroscedastic sine dataset was added. 5) Several more repetitions of the experiment represented in Figure 1 were conducted. (6) An analysis of undersampling behavior on FetchPickAndPlace was added. Moreover, the authors reported two errors in their previous experiments that they discovered and corrected. 

All reviewers were satisfied with the changes in the revised version and the answers to their specific questions and increased their scores accordingly, now commonly voting for acceptance. The paper should therefore be accepted.
Decision: Accept (Poster)