Year: 2021

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 in which each data point’s contribution to the loss is weighted by the β-exponentiated variance estimate. We show that using an appropriate β largely mitigates the issue in our illustrative example. Fourth, we evaluate this approach on a range of standard benchmarks from the literature and show that it achieves considerable improvements and performs more robustly with respect to hyperparameters, both in predictive RMSE and log-likelihood criteria.

heading: 1 INTRODUCTION
text: Endowing models with the ability to capture uncertainty is of crucial importance in machine learning. Uncertainty can be usefully 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 is noise that stays constant across different inputs, whereas heteroscedastic uncertainty is one that varies depending on the inputs to the model.
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 aleatoric uncertainty enables learning dynamics models of stochastic processes (e.g. for model-based or offline reinforcement learning) (Chua et al., 2018; Yu et al., 2020), improves performance in semantic segmentation, depth regression and object detection (Kendall & Gal, 2017; Harakeh & Waslander, 2021), and allows for risk-sensitive decision making (Dabney et al., 2018; Vlastelica et al., 2021).
Many modern applications of machine learning rely on neural networks and deep learning tools to achieve state-of-the-art performance. However, most neural network models are not readily equipped with a capacity for uncertainty estimation. To address this shortcoming, a multitude of approaches have been proposed in the past few years. Nevertheless, the majority of work in this domain has focused on epistemic uncertainty quantification (Blundell et al., 2015; Gal & Ghahramani, 2016) or uncertainty estimation for classification (Hendrycks & Gimpel, 2017; Guo et al., 2017; Ovadia et al., 2019; Mukhoti 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 a 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 of 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 is the de-facto standard (Nix & Weigend, 1994; Lakshminarayanan et al., 2017; Kendall & Gal, 2017; Chua et al., 2018; Kloss et al., 2021), is known to be subject to overconfident variance estimates. Whereas strategies have been proposed to alleviate this specific issue (Detlefsen et al., 2019; Hu et al., 2020; Stirn & Knowles, 2020), we argue that an equally important issue is that this procedure can additionally lead to subpar mean fits. In this work, we analyse and propose a simple modification to mitigate this issue.
Summary of contributions We demonstrate a pitfall of optimizing the NLL loss for neural network regression, one that hinders the training of accurate mean predictors (see Fig. 1 for an illustrative example). The primary culprit is the high dependence of the gradients on the predictive variance. While such dependence is generally known to be responsible for instabilities in joint optimization of mean and variance estimators (Takahashi et al., 2018; Stirn & Knowles, 2020), we identify a fresh perspective on how this dependence can further be problematic. Namely, we hypothesize that the issue arises due to the NLL loss scaling down the gradient of poorly-predicted data points relative to the well-predicted ones, leading to effectively undersampling the poorly-predicted data points.
We then introduce an alternative loss formulation that counteracts this by weighting the contribution of each data point to the overall loss by its β-exponentiated variance estimate, where β controls the extent of dependency of gradients on predictive variance. This formulation subsumes the standard NLL loss for β = 0 and allows to lessen the dependency of gradients on the variance estimates for 0 < β ≤ 1. Interestingly, using β = 1 completely removes such dependency for training the mean estimator, yielding the standard mean squared error (MSE) loss – but with the additional capacity of uncertainty estimation. Finally, we empirically show that our modified loss formulation largely mitigates the issue of poor fits, achieving considerable improvements on a set of standard benchmarks while exhibiting more robustness to network architecture and learning rate configurations.

heading: 2 PRELIMINARIES
text: Let X,Y be two random variables describing the input and target, following the joint distribution P (X,Y ). We assume that Y is conditionally independent givenX 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) = N (µ(X), σ2(X)), where µ : RM 7→ R and σ2 : RM 7→ R+ are respectively the true input-dependent mean and variance functions.1 Equivalently, we can write Y = µ(X)+ (X), with (X) ∼ N (0, σ2(X)); i.e. Y is generated fromX by µ(X) plus a zero-mean Gaussian noise with variance σ2(X). This input-dependent variance quantifies the heteroscedastic uncertainty, or input-dependent aleatoric uncertainty.
1For notational convenience, we focus on univariate regression but point out that the work extends to the multivariate case as well.
To learn estimates µ̂(X), σ̂2(X) of the true mean and variance functions, it is common to use a neural network fθ parameterized by θ. Here, µ̂(X) and σ̂2(X) can be outputs of the final layer (Nix & Weigend, 1994) or use two completely separate networks (Detlefsen et al., 2019).The variance output is hereby constrained to the positive region using a suitable activation function, e.g. softplus. The optimal parameters θ∗NLL can then be found using maximum likelihood estimation (MLE) by minimizing the negative log-likelihood (NLL) criterion LNLL under the distribution P (X,Y ):
θ∗NLL = argmin θ LNLL(θ;D) = argmin θ E X,Y
[ 1
2 log σ̂2(X) + (Y − µ̂(X))2 2σ̂2(X) + const
] . (1)
In contrast, standard regression minimizes the mean squared error (MSE) LMSE:
θ∗MSE = argmin θ LMSE(θ;D) = argmin θ E X,Y
[ (Y − µ̂(X))2
2
] . (2)
In practice, Eq. 1 and Eq. 2 are optimized using stochastic gradient descent with batches of samples drawn from P (X,Y ). The gradients of LNLL with respect to µ̂(X), σ̂2(X) are given by
∇µ̂LNLL(θ) = E X,Y [ µ̂(X)− Y σ̂2(X) ] , ∇σ̂2LNLL(θ) = E X,Y [ σ̂2(X)− (Y − µ̂(X))2 2 ( σ̂2(X) )2 ] . (3, 4)

heading: 3 ANALYSIS
text: We now return to the example of trying to fit a sinusoidal function from Sec. 1. Recall from Fig. 1 that using the Gaussian NLL as the objective resulted in a suboptimal fit. In contrast, using MSE as the objective, the model converged without problems in reasonable time. We now analyze the reasons behind this surprising result.
From Eq. 3, we see that the true mean µ(X) is a minimizer of the NLL loss. It thus becomes clear that a) the solution found in Fig. 1 is not the optimal one, and b) the NLL objective should, in principle, drive µ̂(X) to the optimal solution µ(X). So why is the model not converging? We identify two main culprits for the non-convergence of the Gaussian NLL objective:
1. Initial flatness of the feature space can create an 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 as training progresses.
These culprits and their effect on training are illustrated in Fig. 2 (left). If the network cannot fit a certain region yet, because its feature space (spanned by the last hidden layer) is too coarse, it results in a high effective data variance. This leads to down-weighting the data from these regions, fueling a vicious circle of self-amplifying the increasingly imbalanced weighting. In the following, we analyse these effects and their reasons in more detail.

heading: 3.1 SYMMETRY AND FEATURE NON-LINEARITY
text: It is instructive to see how the model evolves over 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 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 from the stochasticity of mini-batch sampling in SGD, or natural asymmetries contained in the dataset, e.g. outlier points. In addition, we hypothesize that the local non-linearity of the feature space plays an important role in creating the necessary non-linear fit.
Thus, let us consider the non-linearity of the feature space. This quantity is not easy to capture. To approximate it, we compute how much the Jacobian Jf of the features f(x) with respect to the input varies in an L2-ball with radius r around a point x, denoted as the Jacobian variance:2
V (x) = 1 |B| ∑ x′∈S
( Jf (x
′)− 1|S| ∑ x′′∈B Jf (x ′′)
)2 , B = {x′ : ‖x− x′‖2 ≤ r}. (5)
Figure 4 visualizes the Jacobian variance over the input space as a function of the training progress. Although initially relatively flat, it becomes more fine-granular in parts of the input space – the parts which are 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 the success or failure of learning. However, why does gradient descent not break out of this situation?
0 4 8 12 Input X
0
1
2
3
U pd
at e
St ep
s ×
10 6
10−4
100
Ja co
bi an
V ar
ia nc
e
Figure 4: Jacobian variance V (x) (see Eq. 5) over training time. 0 4 8 12 Input X
0
1
2
3
U pd
at e
St ep
s ×
10 6
10−4 10−2 Sa m pl in g Pr ob
ab ili
ty
Figure 5: Probability of sampling a data point at x over training time.

heading: 3.2 WEIGHTING BY VARIANCE EFFECTIVELY UNDERSAMPLES
text: The answer lies in an imbalanced weighting of data points across the input space. Recall that the gradient∇µLNLL of the NLL with respect to the mean scales the error µ̂(X)− Y by 1σ̂2(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 MSE. If the variance is well-calibrated, the gradient becomes µ̂(X)−Y σ̂2(X) ≈ µ̂(X)−Y (µ̂(X)−Y )2 = 1 µ̂(X)−Y . Data points with already low error will get their contribution in the batch gradient scaled up relatively to high error data points – a “rich get richer” self-amplification. 2This is a form of second-order derivative computed numerically, which also gives non-zero results for networks with ReLU activation (in contrast to, for example, the Hessian).
Thus, NLL acts opposite to MSE which focuses on high-error samples. If the true variance σ2 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 prevented on the badly-fit regions.
Another way to view this is to interpret the different weighting of points as changing the training distribution P (X,Y ) to a modified distribution P̃ (X,Y ) in which points with high error have a lower probability of getting sampled. This can be shown by defining P̃ (X,Y ) = Z−1 P (X,Y )σ2(X) ,
where Z = ∫ P (x,y)
σ2(x) dxdy is a normalizing constant, and recognizing that the gradient of the NLL is
proportional to the gradient of the MSE loss in Eq. 2 under the modified data distribution P̃ (X,Y ): ∇µLNLL(θ) = Z · E(X,Y )∼P̃ (X,Y )[µ(X)− Y ] ∝ ∇µE(X,Y )∼P̃ (X,Y ) [ (Y − µ(X))2
2
] . (6)
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 it is highly unlikely to sample points from this region (10−5 as opposed to 10−3 for uniform sampling). We show that this behavior also carries over to a real-world dataset in Sec. A.2. As a side note, optimization with SGD typically assumes a fixed training distribution; training under shifting distributions may thus make optimization with SGD more difficult. It also was recently shown that such a shifting training setup can hurt generalization (Igl et al., 2021). These two points may further contribute to the suboptimal performance of probabilistic neural networks experienced in practice.
Sometimes “weighting-by-variance” is seen as a feature of the Gaussian NLL (Kendall & Gal, 2017) that introduces a self-regularizing property by allowing the network to “ignore” outlier points with high error. This can be desirable if the predicted variance corresponds to the inherent unpredictability (noise), but it is undesirable if it causes premature convergence and ignorance of hard-to-fit regions, as shown above. In our method (Sec. 4), we allow to control the extent of self-regularization.

heading: 4 METHOD
text: To mitigate these problems with NLL training, we develop a solution in the following. We start with a method that uses standard squared error losses to estimate the moments of the target distribution, which we term “moment matching” (MM; Sec. 4.1). This method fixes the problem with premature convergence when using LNLL and yields consistent training results, but it also leads to underestimated variances and will act as one of the baselines in our empirical evaluations. In Sec. 4.2, we then consider the distribution of residual errors which are drastically different between NLL training and MSE training. This suggests a problem-specific design choice that leads us to our proposed solution, which we term β−NLL, in Sec. 4.3. This solution method allows choosing a meaningful loss-interpolation between NLL and MSE while keeping calibrated uncertainty estimates.

heading: 4.1 MOMENT MATCHING
text: A natural solution to the problem of estimating distributions would be to estimate the sufficient statistics of the distribution. In the case of the Gaussian distribution, these would be the first two moments µ, σ2 fully describing the distribution. So why not defining losses based on the moment estimators? Starting from the conditional mean of the predictions y: E[Y | X], we can define the squared deviation as a loss and see that the MSE is an upper bound:
(E[Y | X]− µ̂(X))2 = E[(Y − µ̂(X)) | X]2 ≤ E [ (Y − µ̂(X))2 | X ] := LµMM . (7)
Notice that LµMM is the standard MSE loss. Thus, fitting the mean µ̂(x) can follow standard (nondistributional) regression procedures. Interestingly, due to the moment-matching viewpoint, we can analogously define a loss for the variance (second central moment). The variance is defined as E [ (Y − µ(X))2 | X ] ; analogously, we get the following loss: E [( (Y − µ̂(X))2 − σ̂2(X)
)2 | X]. To use the same physical unit as LµMM, we reformulate it in terms of the standard deviation as:
LσMM := E [(√ (Y − µ̂(X))2 − σ̂(X) )2 ∣∣∣∣ X ] . (8)
Thus, the moment matching loss is simply given by the sum of both losses: LMM = LµMM + LσMM.

heading: 4.2 RESIDUAL DISTRIBUTION AND ALLOCATION OF FUNCTION APPROXIMATOR CAPACITY
text: The moment matching loss gives the same weighting to all data points. Since we are dealing with function approximators, the question is whether their capacity should be used equally. As discussed in Sec. 3.2, the NLL loss gives high weight to data points with low (predicted) variance and low weight to those with high variance. This is appropriate in case these variances are caused by true aleatoric uncertainty in the data. However, due to the use of function approximators, there is also the case where data points cannot be well predicted, resulting in high predicted variance, although the ground truth is corrupted by little noise. The handling of these difficult points is different for the different loss functions. Figure 6 shows the distribution of the residuals for a dynamics prediction dataset containing easy and hard to model areas. We observe a drastic difference between the NLL loss and the MSE loss (or equivalently solely optimizing LµMM). How important are the data points with high uncertainty? Are these outliers or stem from truly noisy regions? In this case, we would be willing to allocate less of the function approximator’s capacity to them. Or are these simply difficult samples? In our dynamics prediction example, we analysed the data points with high uncertainty and found that they were actually the most important ones to get correct (because they captured non-trivial interactions in the physical world). Following this insight, this suggests that there is no one-loss-fits-all, but rather that the modeler should have the choice to select which behavior is desired for the application at hand.

heading: 4.3 VARIANCE-WEIGHTING THE GRADIENTS OF THE NLL
text: In this section, we introduce our proposed solution to the problems occurring when training with LNLL. The most important problem to solve 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 with its self-amplifying character as the main reason for this. Effectively, NLL weights the mean squared error per data point with 1σ2 , which can be interpreted as sampling data points with P (x) ∝ 1σ2 (Sec. 3.2). 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) ∝ σ2βσ2 and illustrated in Fig. 2 (right). How could this weighting be achieved? We simply introduce the variance-weighting term σ2β to the LNLL loss such that it acts as a factor on the gradient. We denote the resulting loss as β−NLL:
Lβ−NLL := E X,Y
[ bσ̂2β(X)c ( 1
2 log σ̂2(X) + (Y − µ̂(X))2 2σ̂2(X) + const
)] , (9)
where b·c denotes the stop gradient operation. By stopping the gradient, the variance weighting term acts as an adaptive, input-dependent learning rate. In this way, the gradients of Lβ−NLL are:
∇µ̂Lβ−NLL(θ) = E X,Y [ µ̂(X)− Y σ̂2−2β(X) ] ,∇σ̂2Lβ−NLL(θ) = E X,Y [ σ̂2(X)− (Y − µ̂(X))2 2σ̂4−2β(X) ] . (10, 11)
Naturally, for β = 0, we recover the original NLL loss. 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. This is in contrast to the moment matching loss (Eq. 8). For values 0 < β < 1, we get different loss interpolations. Particularly interesting is the case of β = 0.5, where the data points are weighted with 1σ (inverse standard deviation instead of inverse variance). In Sec. 5, we show that β = 0.5 generally achieves the best trade-off between accuracy and log-likelihood. A Pytorch implementation of the loss function is provided in Sec. C.5.
As a remark, the new loss Lβ−NLL is not meant for performance evaluation – it is designed to result in meaningful gradients. We found that its absolute value can be counter-intuitive due to the prefactor. The model performance during training should be monitored with the original (negative) log-likelihood and optionally the RMSE for the mean fit.

heading: 5 EXPERIMENTS
text: We investigate the proposed changes in the loss function for training neural networks to predict a Gaussian distribution on synthetic examples and real-world datasets featuring the UCI, challenging dynamics model datasets, and regression tasks on MNIST, Fashion-MNIST, and depth-map prediction.
As motivated above, the standard NLL loss often leads to suboptimal fits, e.g. where the mean fit is far from ideal. To analyze this quantitatively, we report the performance for many hyperparameter settings, such as learning rates and network architectures, and consider both the quality in predictions with respect to the root mean squared error (RMSE) and the negative log-likelihood (NLL).

heading: 5.1 SYNTHETIC DATASETS
text: Sinusoidal without heteroscedastic noise We first consider our illustrative example of Fig. 1 – a simple sine curve with a small additive noise: y = sin(x) + ξ, with ξ being Gaussian noise with standard deviation σ = 0.01. One would expect that a network with sufficient capacity can easily learn to fit this function. Figure 7 inspects this over a range of architectures and learning rates.
We find that for the standard NLL (β = 0) the network does not converge to a reasonable mean fit. There is a trend that larger networks and learning rates show better results, but when comparing this to β−NLL with β = 0.5 we see that the networks are indeed able to fit the function without problems. As expected, the same holds for the mean squared error loss (MSE) and β−NLL with β = 1. The quality of the fit w.r.t. NLL is shown in Fig. S3 for completeness.
Sinusoidal with heteroscedastic noise A commonly used illustrative example introduced in Detlefsen et al. (2019) is a sine curve with increasing amplitude and noise: y = x sin(x) + xξ1 + ξ2, with ξ1 and ξ2 being Gaussian noise with standard deviation σ = 0.3. The functional form is much easier, so fitting the mean is achieved with all losses. Here we sanity-check that the newly introduced loss is still delivering good uncertainty estimates. 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.
Lβ−NLL 0 2 -3.25 ± 0.31 -2.44 ± 0.36 12.74 ± 0.47 -2.86 ± 5.18 4 6.08 ± 0.65 2.21 ± 0.21 0.0022 ± 0.0006 1.22 ± 0.47 Lβ−NLL 0.25 3 -3.31 ± 0.51 -2.13 ± 0.22 13.63 ± 0.36 -1.97 ± 1.14 5 5.79 ± 0.74 1.82 ± 0.18 0.0012 ± 0.0004 1.73 ± 1.00 Lβ−NLL 0.5 5 -3.29 ± 0.36 -1.89 ± 0.20 13.77 ± 0.58 -2.47 ± 1.68 6 5.61 ± 0.65 1.10 ± 0.10 0.0006 ± 0.0002 2.35 ± 1.44 Lβ−NLL 0.75 5 -3.27 ± 0.34 -2.05 ± 0.27 13.74 ± 0.44 -1.87 ± 0.55 7 5.67 ± 0.73 1.04 ± 0.12 0.0004 ± 0.0001 1.97 ± 1.03 Lβ−NLL 1 2 -3.23 ± 0.33 -2.89 ± 0.61 13.59 ± 0.38 -2.27 ± 1.07 8 5.55 ± 0.77 1.44 ± 0.38 0.0004 ± 0.0001 2.08 ± 1.13 LMM 0 -3.49 ± 0.38 -4.05 ± 0.25 -2.46 ± 66.93 -11.2 ± 31.0 4 6.28 ± 0.82 2.38 ± 0.23 0.0005 ± 0.0001 3.02 ± 1.38 LMSE — — — — — 11 4.96 ± 0.64 0.95 ± 0.08 0.0003 ± 0.0001 0.78 ± 0.25 Student-t 10 -3.07 ± 0.14 -2.22 ± 0.25 12.49 ± 0.50 -1.23 ± 0.55 5 5.82 ± 0.59 2.22 ± 0.18 0.0023 ± 0.0007 1.34 ± 0.63 xVAMP 7 -3.06 ± 0.15 -2.34 ± 0.15 12.66 ± 0.55 -0.99 ± 0.33 7 5.44 ± 0.64 2.06 ± 0.19 0.0019 ± 0.0008 0.99 ± 0.43 xVAMP* 7 -3.03 ± 0.13 -2.30 ± 0.22 12.85 ± 0.59 -1.04 ± 0.47 8 5.35 ± 0.73 2.03 ± 0.13 0.0017 ± 0.0006 1.13 ± 0.66 VBEM 2 -3.14 ± 0.07 -4.33 ± 0.18 8.14 ± 0.12 -2.65 ± 0.10 8 5.21 ± 0.58 1.36 ± 0.36 0.0010 ± 0.0005 1.66 ± 0.84 VBEM* 8 -2.99 ± 0.13 -1.92 ± 0.24 12.81 ± 0.98 -0.98 ± 0.24 8 5.17 ± 0.59 1.26 ± 0.23 0.0018 ± 0.0006 0.65 ± 0.20

heading: 5.2 REAL-WORLD DATASETS
text: UCI Regression Datasets As a standard real-world benchmark in predictive uncertainty estimation, we consider the UCI datasets (Hernández-Lobato & Adams, 2015). Table 1 gives an overview comparing different loss variants. We refer to Sec. A.3 for the full results on all 12 datasets. The results are encouraging: β−NLL achieves predictive log-likelihood on par with or better than the NLL loss while clearly improving the predictive accuracy on most datasets.
Dynamics models As a major application of uncertainty estimation lies in model-based reinforcement learning (RL), we test the different loss functions on two dynamics predictions tasks of varying difficulty, ObjectSlide and Fetch-PickAndPlace. In both tasks, the goal is to predict how an object will move from the current state and the agent’s action. Whereas ObjectSlide (Seitzer et al., 2021) is a simple 1D-environment, Fetch-PickAndPlace (Plappert et al., 2018) is a complex 3D roboticmanipulation environment. The models are trained on state-action trajectories collected by RL agents. See Sec. B for more details.
For both datasets and all loss functions, we perform a grid search over different hyperparameter configurations (see Sec. C.2) for a sensitivity analysis to hyperparameters settings, presented in Fig. 9. It reveals that both the NLL- and the MM-loss are vulnerable to wrong hyperparameter settings, whereas the β−NLL loss achieves good results over a wide range of settings. The best performing configurations for each loss are evaluated on a hold-out test set (Table 2). One can see that the NLL loss results in poor predictive performance and also exhibits quite a high variance across random seeds; the MM-loss results in poor log-likelihood fits. Our method yields high accuracy and log-likelihood fits, for a range of β-values, where β = 0.5 generally achieves the best trade-off.
Generative models and depths-map prediction In Fig. S5 we present results on variational autoencoders with probabilistic decoders, showing qualitatively similar trends as on UCI. More details are provided in Sec. A.4. On the task of depth regression we modify a state-of-the-art method (Bhat et al., 2021) and test it on the NYUv2 dataset (Silberman et al., 2012) with our loss (Fig. S6). β-NLL yields an improved mean fit in comparison to plain NLL while also producing sharper depth maps, as detailed in Sec. A.5.

heading: 6 DISCUSSION
text: We highlight a deep problem frequently occurring when optimizing probabilistic neural networks using the common NLL loss: training gets stuck in suboptimal function fits. Due to the predictive variance-based weighting of mean prediction errors in the NLL loss, these bad fits often stay unnoticed, unless the MSE is evaluated. With our analysis on feature space non-linearity and induced data-set weighting (interpreted as a biased dataset sampling), we reveal the underlying reason: initially badly fit regions (due to low feature non-linearity) get increasingly less weight and result in premature convergence. We propose a simple solution by introducing a family of loss functions called β−NLL. Effectively, the gradient of the original NLL loss is scaled by the β-exponentiated per-sample variance. This allows for a meaningful interpolation between NLL loss and MSE loss while providing wellbehaved uncertainty estimates (even for the MSE case). The hyperparameter β gives practitioners the choice to control the self-regularization strength of NLL: how unimportant should high-noise regions or difficult-to-predict data points be in the fitting process. In most cases, β = 0.5 will be a good starting point. We hope that our simple solution will improve the usability and performance of predictive uncertainty deep networks and increase their applicability.

heading: REPRODUCIBILITY STATEMENT
text: All settings are described in detail in Sec. B and Sec. C. We make full code and data available under https://sites.google.com/view/pitfalls-uncertainty. Both will be made publicly available after the review phase.

heading: APPENDIX

heading: A ADDITIONAL RESULTS

heading: A.1 SYNTHETIC DATASET
text: In Fig. S1, we replicate the experiment from Fig. 1 (training the NLL loss on a sinusoidal for 107 updates) on several more random seeds. In order to test the dependence of our reported issue on the optimizer, we repeat the experiment from Fig. 1, but with different optimizers instead of Adam with β1 = 0.9, β2 = 0.999. The results are shown in 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 is occurring stably across optimization settings. In Fig. S3, we provide the results for the NLL metric on the sinusoidal dataset, complimentary to the results in Fig. 7 (see discussion in Sec. 5.1).
0 3 6 9 12 Input X
-0.4
0.0
0.4
Ta rg
et Y
0 3 6 9 12 Input X
-0.4
0.0
0.4
Ta rg
et Y
0 3 6 9 12 Input X
-0.4
0.0
0.4
Ta rg
et Y
Figure S1: Repeating the experiment from Fig. 1, i.e. training with NLL loss for 107 update steps. The observed behavior is stable across different independent trials.
SG D,
No Mo
m.
SG D,
Mo m.=
0.9
SG D,
Mo m.=
0.9 9
0.005
0.01
0.05
L ea
rn in
g R
at e
0.10
0.15
0.20
0.25
0.30
R M
SE
RM SPr
op, α =
0.9
RM SPr
op, α =
0.9 9
Ad am
, β 1=
0.9
β 2= 0.9
Ad am
, β 1=
0.5
β 2= 0.9
0.0001
0.0005
0.001
L ea
rn in
g R
at e
0.10
0.15
0.20
0.25
0.30
R M
SE
Figure S2: Using different optimizers to train on the sinusoidal from Fig. 1 with NLL loss. The color code indicates the mean RMSE over 3 independent trials per optimizer setting. Black indicates that all trials diverged. Training was done for 2 · 106 update steps and used architecture 2 from Table S5. The observed behavior is stable across optimization settings.
(a) β = 0 (NLL) (b) β = 0.5 (c) β = 1
0 1 2 3 4 Architecture
3e-05
1e-04
3e-04
1e-03
L ea
rn in
g R
at e
−3
−2
−1
N L
L
0 1 2 3 4 Architecture
3e-05
1e-04
3e-04
1e-03
L ea
rn in
g R
at e
−3
−2
−1
N L
L
0 1 2 3 4 Architecture
3e-05
1e-04
3e-04
1e-03
L ea
rn in
g R
at e
−3
−2
−1
N L
L
Figure S3: Convergence properties analyzed on the sinusoidal toy regression problem. Same as Fig. 7, but for the negative log-likelihood (NLL). Due to the bad mean fit, the original NLL loss (β = 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 β−NLL as testified by the low NLL scores.

heading: A.2 ANALYSIS OF SAMPLING PROBABILITIES ON FETCH-PICKANDPLACE
text: In Fig. S4, we show how the analysis from Sec. 3.2 transfers to a real world dataset, namely FetchPickAndPlace. 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 a model trained with the MSE loss. T he 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 our analysis in Sec. 3.2.
10−8 10−5 10−2 Sampling Probability
0
200
400
600
C ou
nt
Initial LNLL LMSE (trained)
10−8 10−5 10−2 Sampling Probability
0
200
400
600
C ou
nt
Epoch 1
10−8 10−5 10−2 Sampling Probability
0
200
400
600
C ou
nt
Epoch 10
10−8 10−5 10−2 Sampling Probability
0
200
400
600
C ou
nt
Epoch 50
10−8 10−5 10−2 Sampling Probability
0
200
400
600
C ou
nt
Epoch 100
10−8 10−5 10−2 Sampling Probability
0
200
400
600
C ou
nt
Epoch 200
Figure S4: Undersampling behavior of LNLL on Fetch-PickAndPlace. The plot shows how the distribution of effective sampling probability evolves over training time, taken over 2 000 fixed training points sampled at the initial epoch. The dashed blue histogram shows the distribution of effective sampling probabilities when using the squared residuals from a model trained with MSE loss LMSE. This gives a reference distribution that LNLL should roughly match, taking into account the relative hardness of prediction on different samples. The LNLL drastically undersamples compared to the reference (note the log-scale), effectively never sampling some points.

heading: A.3 UCI DATASETS
text: 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-t distribution (xVAMP, VBEM (Stirn & Knowles, 2020)) tend to have better predictive log-likelihood than β−NLL, although there are also datasets where β−NLL performs better (“energy”, “naval”), or there is no statistically significant improvement (“housing”, “kin8m”, “wine-red”). We conjecture this is because the Student-t distribution might be better suited to model real-world data than the Gaussian, especially in low-data regimes. For RMSE, LMSE unsurprisingly performs best. Our β−NLL is often on-par with LMSE and on-par or better than the baselines, with the exception of “yacht”. At the same time, our method is very simple to implement (see Sec. C.5) and computationally lightweight compared to xVAMP and VBEM, which require the costly evaluation of a prior and Monte-Carlo sampling at each training step.
A.4 VARIATIONAL AUTOENCODERS
We test different loss functions on the task of generative modelling using Variational Autoencoders (VAEs) (Kingma & Welling, 2014). To this end, we parametrize the decoder distribution p(x | z) with N ( µ(z), σ2(z) ) , where µ(z) is the mean and σ2(z) the variance output from a neural network. We train the VAE by maximizing the ELBO Eq(z|x)[log p(x | z)]−DKL(q(z | x) || p(z)), plugging in different loss functions for log p(x | z). Following (Stirn & Knowles, 2020), we evaluate the log-posterior predictive likelihood logEq(z|x)[p(x | z)]. We approximate the expectation using a finite mixture of 20 Monte-Carlo samples from q(z | x). To compute the RMSE, we take the mean of
Table S1: Results for UCI Regression Datasets. Predictive log-likelihood (higher 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 †.
carbon concrete energy housing Loss (10721, 5, 3) (1030, 8, 1) (768, 8, 2) (506, 13, 1)
Lβ−NLL(β = 0) 13.13 ± 0.25 -3.25 ± 0.31 -2.44 ± 0.36 -2.86 ± 0.50 Lβ−NLL(β = 0.25) 12.99 ± 0.29 -3.31 ± 0.51 -2.13 ± 0.22 -2.75 ± 0.42 Lβ−NLL(β = 0.5) 12.36 ± 0.84 -3.29 ± 0.36 -1.89 ± 0.20 -2.64 ± 0.36† Lβ−NLL(β = 0.75) 10.16 ± 4.97 -3.27 ± 0.34 -2.05 ± 0.27† -2.72 ± 0.42† Lβ−NLL(β = 1.0) -0.47 ± 30.22 -3.23 ± 0.33 -2.89 ± 0.61 -2.85 ± 0.87 LMM 12.02 ± 0.83 -3.49 ± 0.38 -4.05 ± 0.25 -3.42 ± 1.01 LMSE — — — — Student-t 15.68 ± 0.30 -3.07 ± 0.14† -2.22 ± 0.25 -2.47 ± 0.24† xVAMP 13.20 ± 0.20 -3.06 ± 0.15† -2.34 ± 0.15 -2.43 ± 0.21† xVAMP* 13.14 ± 0.25 -3.03 ± 0.13† -2.30 ± 0.22 -2.45 ± 0.22† VBEM 5.81 ± 0.65 -3.14 ± 0.07 -4.33 ± 0.18 -2.56 ± 0.15 VBEM* 13.12 ± 0.37 -2.99 ± 0.13 -1.92 ± 0.24† -2.42 ± 0.22
kin8m naval power protein Loss (8192, 8, 1) (11934, 16, 2) (9568, 4, 1) (45730, 9, 1)
Lβ−NLL(β = 0) 1.140 ± 0.039† 12.74 ± 0.47 -2.807 ± 0.057 -2.80 ± 0.05 Lβ−NLL(β = 0.25) 1.142 ± 0.026† 13.63 ± 0.36† -2.801 ± 0.053 -2.79 ± 0.05 Lβ−NLL(β = 0.5) 1.141 ± 0.046† 13.77 ± 0.58 -2.805 ± 0.052 -2.78 ± 0.02 Lβ−NLL(β = 0.75) 1.137 ± 0.041† 13.74 ± 0.44† -2.806 ± 0.053 -2.79 ± 0.02 Lβ−NLL(β = 1.0) 1.126 ± 0.041 13.59 ± 0.38† -2.810 ± 0.051 -2.80 ± 0.03 LMM 0.999 ± 0.063 -2.46 ± 66.93 -2.918 ± 0.092 -2.98 ± 0.06 LMSE — — — — Student-t 1.155 ± 0.037 12.49 ± 0.50 -2.738 ± 0.026 -2.56 ± 0.02 xVAMP 1.147 ± 0.037† 12.66 ± 0.55 -2.788 ± 0.032 -2.74 ± 0.03 xVAMP* 1.147 ± 0.036† 12.85 ± 0.59 -2.785 ± 0.036 -2.71 ± 0.02 VBEM 1.019 ± 0.054 8.14 ± 0.12 -2.819 ± 0.018 -2.85 ± 0.01 VBEM* 1.138 ± 0.036† 12.81 ± 0.98 -2.783 ± 0.031 -2.73 ± 0.01
superconductivity wine-red wine-white yacht Loss (21263, 81, 1) (1599, 11, 1) (4898, 11, 1) (308, 6, 1)
Lβ−NLL(β = 0) -3.75 ± 0.17 -1.03 ± 0.24† -1.011 ± 0.016 -2.86 ± 5.18 Lβ−NLL(β = 0.25) -3.71 ± 0.15 -0.98 ± 0.12† -1.010 ± 0.016 -1.97 ± 1.14 Lβ−NLL(β = 0.5) -3.65 ± 0.07 -0.99 ± 0.16† -1.009 ± 0.014 -2.47 ± 1.68 Lβ−NLL(β = 0.75) -3.71 ± 0.09 -1.02 ± 0.22† -1.013 ± 0.013 -1.87 ± 0.55 Lβ−NLL(β = 1.0) -3.84 ± 0.07 -0.97 ± 0.10† -1.014 ± 0.021 -2.27 ± 1.07 LMM -4.30 ± 0.28 -1.22 ± 0.43 -1.116 ± 0.055 -11.24 ± 31.03 LMSE — — — — Student-t -3.40 ± 0.04 -0.94 ± 0.10† -1.001 ± 0.010† -1.23 ± 0.55† xVAMP -3.43 ± 0.03† -0.95 ± 0.06† -0.995 ± 0.016† -0.99 ± 0.33† xVAMP* -3.42 ± 0.02† -0.94 ± 0.06 -0.992 ± 0.015 -1.04 ± 0.47† VBEM -3.64 ± 0.15 -0.95 ± 0.07† -1.004 ± 0.013† -2.65 ± 0.10 VBEM* -3.42 ± 0.03† -0.94 ± 0.07† -0.999 ± 0.017† -0.98 ± 0.24
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 †.
carbon concrete energy housing Loss (10721, 5, 3) (1030, 8, 1) (768, 8, 2) (506, 13, 1)
Lβ−NLL(β = 0) 0.0027 ± 0.0001 6.08 ± 0.65 2.21 ± 0.21 3.56 ± 1.07† Lβ−NLL(β = 0.25) 0.0027 ± 0.0001 5.79 ± 0.74 1.82 ± 0.18 3.48 ± 1.15† Lβ−NLL(β = 0.5) 0.0027 ± 0.0002 5.61 ± 0.65 1.10 ± 0.10 3.42 ± 1.04† Lβ−NLL(β = 0.75) 0.0029 ± 0.0002 5.67 ± 0.73 1.04 ± 0.12 3.43 ± 1.07† Lβ−NLL(β = 1.0) 0.0034 ± 0.0004 5.55 ± 0.77† 1.44 ± 0.38 3.50 ± 0.95† LMM 0.0035 ± 0.0005 6.28 ± 0.82 2.38 ± 0.23 4.02 ± 1.18 LMSE 0.0029 ± 0.0001 4.96 ± 0.64 0.95 ± 0.08 3.24 ± 1.08† Student-t 0.0026 ± 0.0000 5.82 ± 0.59 2.22 ± 0.18 3.48 ± 1.17† xVAMP 0.0026 ± 0.0001 5.44 ± 0.64† 2.06 ± 0.19 3.23 ± 1.00† xVAMP* 0.0027 ± 0.0002† 5.35 ± 0.73† 2.03 ± 0.13 3.38 ± 1.15† VBEM 0.0037 ± 0.0008 5.21 ± 0.58† 1.36 ± 0.36 3.32 ± 1.06† VBEM* 0.0026 ± 0.0002 5.17 ± 0.59† 1.26 ± 0.23 3.19 ± 1.02
kin8m naval power protein Loss (8192, 8, 1) (11934, 16, 2) (9568, 4, 1) (45730, 9, 1)
Lβ−NLL(β = 0) 0.087 ± 0.004 0.0022 ± 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.05† Lβ−NLL(β = 0.5) 0.082 ± 0.003† 0.0006 ± 0.0002 4.04 ± 0.17† 4.31 ± 0.02† Lβ−NLL(β = 0.75) 0.081 ± 0.004† 0.0004 ± 0.0001† 4.04 ± 0.15† 4.28 ± 0.02† Lβ−NLL(β = 1.0) 0.081 ± 0.003† 0.0004 ± 0.0001† 4.06 ± 0.18† 4.31 ± 0.05† LMM 0.082 ± 0.003† 0.0005 ± 0.0001 4.07 ± 0.16† 4.32 ± 0.07† LMSE 0.081 ± 0.003 0.0003 ± 0.0001 4.01 ± 0.19 4.28 ± 0.07 Student-t 0.085 ± 0.005 0.0023 ± 0.0007 4.02 ± 0.16† 4.76 ± 0.24 xVAMP 0.081 ± 0.003† 0.0019 ± 0.0008 4.03 ± 0.17† 4.38 ± 0.05† xVAMP* 0.082 ± 0.003† 0.0017 ± 0.0006 4.03 ± 0.18† 4.31 ± 0.02† VBEM 0.082 ± 0.003† 0.0010 ± 0.0005 4.09 ± 0.15† 4.31 ± 0.01† VBEM* 0.082 ± 0.004† 0.0018 ± 0.0006 4.02 ± 0.18† 4.35 ± 0.09†
superconductivity wine-red wine-white yacht Loss (21263, 81, 1) (1599, 11, 1) (4898, 11, 1) (308, 6, 1)
Lβ−NLL(β = 0) 13.85 ± 0.38 0.636 ± 0.038† 0.662 ± 0.007† 1.22 ± 0.47 Lβ−NLL(β = 0.25) 13.37 ± 0.24 0.638 ± 0.036† 0.661 ± 0.008† 1.73 ± 1.00 Lβ−NLL(β = 0.5) 13.19 ± 0.47 0.635 ± 0.037† 0.662 ± 0.008† 2.35 ± 1.44 Lβ−NLL(β = 0.75) 13.25 ± 0.41 0.638 ± 0.035† 0.662 ± 0.007† 1.97 ± 1.03 Lβ−NLL(β = 1.0) 13.41 ± 0.35 0.639 ± 0.035† 0.659 ± 0.010 2.08 ± 1.13 LMM 13.66 ± 0.58 0.652 ± 0.044† 0.665 ± 0.006 3.02 ± 1.38 LMSE 12.49 ± 0.24 0.633 ± 0.036† 0.661 ± 0.008† 0.78 ± 0.25† Student-t 13.81 ± 0.50 0.636 ± 0.038† 0.662 ± 0.007† 1.34 ± 0.63 xVAMP 13.52 ± 0.51 0.635 ± 0.035† 0.662 ± 0.012† 0.99 ± 0.43 xVAMP* 13.44 ± 0.37 0.633 ± 0.035† 0.661 ± 0.010† 1.13 ± 0.66 VBEM 12.67 ± 0.48† 0.639 ± 0.041† 0.663 ± 0.012† 1.66 ± 0.84 VBEM* 13.30 ± 0.32 0.633 ± 0.040 0.663 ± 0.013† 0.65 ± 0.20
that mixture. We compare β−NLL with LMM, LNLL with a fixed variance of 1, Student-t (Takahashi et 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 at the expense of poor log-likelihood. Vice-versa, Student-t, xVAMP, xVAMP*, VBEM* achieves strong log-likelihoods but worse reconstruction errors. β−NLL with β > 0 provides a good compromise between log-likelihood and reconstruction error. Figure S5 shows qualitative examples. Our β-NLL loss with β > 0 allows to learn good reconstructions and meaningful uncertainties. Moreover, samples produced from sampling latents from the prior produce semantically meaningful images, indicating that adding our loss function does not break the disentangling properties of VAEs. Compare that to Student-t, xVAMP(*) and VBEM(*), which do not produce similarly clear images when sampling from the prior.
Table S3: Results for generative modelling with Variational Autoencoders on MNIST and FashionMNIST. We report RMSE and posterior predictive log-likelihood (LL) with standard deviation over 5 independent trials.
MNIST Fashion-MNIST Loss RMSE ↓ LL ↑ RMSE ↓ LL ↑
LNLL (σ2 = 1) 0.153 ± 0.002 -730 ± 0 0.143 ± 0.002 -729 ± 0 Lβ−NLL(β = 0) 0.237 ± 0.002 2116 ± 55 0.170 ± 0.001 1940 ± 104 Lβ−NLL(β = 0.25) 0.181 ± 0.004 2511 ± 88 0.140 ± 0.002 2010 ± 39 Lβ−NLL(β = 0.5) 0.151 ± 0.003 2220 ± 25 0.125 ± 0.003 1639 ± 52 Lβ−NLL(β = 0.75) 0.142 ± 0.001 1954 ± 50 0.131 ± 0.001 1331 ± 32 Lβ−NLL(β = 1.0) 0.152 ± 0.001 1706 ± 30 0.138 ± 0.002 1142 ± 26 LMM 0.260 ± 0.000 385 ± 11 0.295 ± 0.000 -151 ± 4 Student-t 0.273 ± 0.002 4291 ± 103 0.182 ± 0.002 2857 ± 9 xVAMP 0.225 ± 0.002 2989 ± 268 0.161 ± 0.001 2158 ± 100 xVAMP* 0.225 ± 0.001 3062 ± 215 0.160 ± 0.002 2150 ± 131 VBEM 0.114 ± 0.001 719 ± 9 0.108 ± 0.000 660 ± 2 VBEM* 0.176 ± 0.008 3213 ± 238 0.150 ± 0.003 2244 ± 78

heading: A.5 DEPTH REGRESSION
text: We evaluate β−NLL on the task of 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 function. Note that we remove the Mini-ViT Transformer module from the model, and thus our results are not directly comparable with those reported by (Bhat et al., 2021).
Table S4 presents quantitative results. β−NLL with β > 0 achieves better RMSE than the NLL loss. β−NLL with β = 0.5 again provides a good trade-off, achieving similar RMSE as LMSE and performing better on some of the other metrics. Figure S6 shows qualitative examples. Compared to LNLL, the depth maps predicted by β−NLL with β > 0 are noticeably sharper.

heading: B DATASETS AND TRAINING SETTINGS
text: 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(x) = 0.4 sin(2πx) + ξ to them to create the targets y, where ξ is Gaussian noise with a standard deviation of 0.01.
Sinusoidal heteroscedastic noise We use the synthetic data as introduced in Detlefsen et al. (2019). From the functional form y = x sin(x) + xξ1 + ξ2, with Gaussian noise with standard deviation σ = 0.3 for ξ1 and ξ2, we sample 500 points uniformly spaced in the interval [0, 10]. The model is a MLP with one hidden layer with 50 units and tanh activation function (as used in Detlefsen et al. (2019) and Stirn & Knowles (2020)).
MNIST Fashion-MNIST dataset
Lβ−NLL(β = 0) mean
std ∼posterior ∼ N (0, 1)
Lβ−NLL(β = 0.25) mean
std ∼posterior ∼ N (0, 1)
Lβ−NLL(β = 0.5) mean
std ∼posterior ∼ N (0, 1)
Lβ−NLL(β = 0.75) mean
std ∼posterior ∼ N (0, 1)
Lβ−NLL(β = 1.0) mean
std ∼posterior ∼ N (0, 1)
LNLL (σ2 = 1) mean
std ∼posterior ∼ N (0, 1)
LMM mean
std ∼posterior ∼ N (0, 1)
Student-t mean
std ∼posterior ∼ N (0, 1)
xVAMP mean
std ∼posterior ∼ N (0, 1)
xVAMP* mean
std ∼posterior ∼ N (0, 1)
VBEM mean
std ∼posterior ∼ N (0, 1)
VBEM* mean
std ∼posterior ∼ N (0, 1)
Figure S5: Generative modelling with Variational Autoencoders on MNIST and Fashion-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 priorN (0, 1). LNLL(σ2 = 1) refers to Gaussian log-likelihood with a fixed variance of 1. Values are clipped to the interval [0, 1]. Examples are not cherry-picked.
Input
Ground Truth
Lβ−NLL (β = 0.0)
pred.
std.
Lβ−NLL (β = 0.5)
pred.
std.
Lβ−NLL (β = 0.75)
pred.
std.
Lβ−NLL (β = 1.0)
pred.
std.
LMSE pred.
Figure S6: Example results for depth regression on the NYUv2 dataset (Silberman et al., 2012). First two 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 in form of the predicted standard deviation (std.). Results for Lβ−NLL with β > 0 are noticeably sharper. The last column shows a magnified view on the previous picture.
Table S4: Results for depth regression on the NYUv2 dataset (Silberman et al., 2012). We adapt a stateof-the-art network for depth regression from AdaBins (Bhat et 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 to AdaBins’. 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, even though it is a loss function specifically designed for the task of depth regression. We refer to Bhat et al. (2021) for a description of the metrics.
Method δ1 ↑ δ2 ↑ δ3 ↑ REL ↓ RMSE ↓ log10 ↓ LL ↑ Lβ−NLL(β = 0) 0.8855 0.9796 0.9959 0.1094 0.3854 0.0462 -4.52 Lβ−NLL(β = 0.25) 0.8887 0.9812 0.9956 0.1081 0.3818 0.0458 -8.15 Lβ−NLL(β = 0.5) 0.8885 0.9813 0.9956 0.1093 0.3789 0.0458 -7.50 Lβ−NLL(β = 0.75) 0.8902 0.9804 0.9952 0.1095 0.3800 0.0462 -7.35 Lβ−NLL(β = 1.0) 0.8872 0.9813 0.9958 0.1088 0.3845 0.0467 -5.10 LMSE 0.8890 0.9806 0.9960 0.1086 0.3776 0.0461 — L1 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 — DAV (Chen et al., 2019) 0.882 0.980 0.996 0.108 0.412 — —
UCI Datasets We use the UCI datasets suite commonly used to benchmark uncertainty estimation, stemming from the UCI Machine Learning Repository3. In particular, we use the training-test protocol from (Hernández-Lobato & Adams, 2015; Gal & Ghahramani, 2016), and their data splits 4.
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 divide the training set into 80% training data and 20% validation data and search for an optimal learning rate from the set {10−4, 3 · 10−4, 7 · 10−4, 10−3, 3 · 10−3, 7 · 10−3} by monitoring log-likelihood on the validation set. We train for a maximum of 20 000 updates, besides on the larger “kin8m”, “power plant”, “protein” and “naval” datasets, where we train for a maximum of 100 000 updates. We perform early stopping with a patience of 50 epochs, retrain the model with the best found learning rate on 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, as 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, aside from “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 statistically indistinguishable from the respective best method in Table 1 (”Ties”). For this purpose, we performed a two-sided Kolmogorov-Smirnov with a p ≤ 0.05 significance level.
ObjectSlide This environment consists of an agent whose task it is to slide an object to a target location (Seitzer et al., 2021). The continuous state space consists of 4 dimensions: agent and object position and velocity, and the continuous action space is a one-dimensional movement command. The forward prediction task consists of predicting the change in object position in the next state from current state and action. The dataset we use consists of 180 000 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 a maximum of 5 000 epochs with a batch size of 256, and evaluate the model with the best validation log-likelihood on the test set afterwards.
3https://archive.ics.uci.edu 4available under https://github.com/yaringal/DropoutUncertaintyExps
Fetch-PickAndPlace We use the Fetch-PickAndPlace environment (Plappert et al., 2018) from OpenAI Gym (Brockman et al., 2016) as a challenging real-world scenario. The task of the agent is to use position-controlled 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 3D-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 epochs with a batch size of 256, and evaluate the model with the best validation log-likelihood on the test set afterwards.
NYUv2 Depth Regression We use the dataset in the variant provided by Lee et al. (2019) 5.

heading: C HYPERPARAMETER SETTINGS AND IMPLEMENTATION DETAILS
text: For all experiments, we used the Adam optimizer (Kingma & Ba, 2015) with standard settings β1 = 0.9, β2 = 0.999. We parametrize the Gaussian distribution using two linear layers on top of shared features produced by a MLP. The variance σ̂2(x) is constrained to the positive region using the softplus(x) = log(1 + exp(x)) activation function. We additionally add a small constant of 10−8 to prevent the variance from collapsing to zero, and clamp the maximum variance to 1 000.

heading: C.1 SINUSOIDAL REGRESSION PROBLEM
text: The sinusoidal fit in Fig. 1 results from a network with two hidden layers and 128 neurons each, tanh activation, and optimized with learning rate 5 · 10−4 and a batch size of 100. For the experiment in Sec. 5.1, we scan over learning rates and architectures with different hidden layers and units per layer, as detailed in Table S5.
Table S5: Architectures used for the sinusoidal regression task. Fully connected feed-forward neural networks with tanh activation function.
Architecture # 0 1 2 3 4
# Hidden Layers 2 2 2 3 3 # Units per Layer 32 64 128 128 256

heading: C.2 OBJECTSLIDE AND FETCH-PICKANDPLACE
text: For each tested loss function, we performed a grid search for the ObjectSlide and Fetch-PickAndPlace datasets, We report the parameters in Table S6. For the results in Table 2, we retrained the model configurations with the best validation log-likelihood on the grid search, and evaluated them on the hold-out test set.
The best found configuration on ObjectSlide coincided to 3 hidden layers with 128 neurons, relu activation and learning rate 0.001. The best found configuration on Fetch-PickAndPlace coincided to 4 hidden layers with 128 neurons, relu activation and learning rates 0.0003 for LNLL, Lβ−NLL with β = 0.25 and β = 0.5, and learning rate 0.001 for LMSE, LMM, Lβ−NLL with β = 0.75 and β = 1.0.
C.3 VARIATIONAL AUTOENCODERS
We 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, a decoder with three layers of 128, 256, 512 neurons and ReLU activations. The latent space is 10-dimensional for MNIST and 25-dimensional for FashionMNIST. We train the VAEs for a maximum of 1 000 epochs, using Adam with a learning rate of 0.0003 and a batch size of 256. Early stopping with a patience of 50 epochs
5available under https://github.com/cogaplex-bts/bts
Table S6: Hyperparameter settings for grid search on ObjectSlide and Fetch-PickAndPlace datasets. We run 96 configurations per loss function.
Hyperparameter Set of Values
Learning Rate {3 · 10−5, 10−4, 3 · 10−4, ·10−3} # Hidden Layers {2, 3, 4}
# Units per Layer {128, 256, 386, 512} Activation {tanh, relu}
is performed on the log-likelihood of the validation set. The validation set consists of 20% of the MNIST/FashionMNIST training set.

heading: C.4 DEPTH REGRESSION
text: We use the official implementation of AdaBins (Bhat et al., 2021)6, thereby reproducing their exact training settings and evaluation protocol. We remove the AdaBins/mini-ViT Transformer from the model. Instead, the feature map output by the U-Net is reduced to two channels using a 1 × 1 convolution, where we use the first channel as the mean prediction and the second channels as the variance prediction. In this case, both mean and variance are constrained to positive numbers by a softplus activation. On top of that, we add a 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.
C.5 IMPLEMENTATION OF BETA-NLL IN PYTORCH
1 def beta_nll_loss(mean, variance, target, beta): 2 """Compute beta-NLL loss 3 4 :param mean: Predicted mean of shape B x D 5 :param variance: Predicted variance of shape B x D 6 :param target: Target of shape B x D 7 :returns: Loss per batch element of shape B 8 """ 9 loss = 0.5 * ((target - mean) ** 2 / variance + variance.log())
10 11 if beta > 0: 12 loss = loss * variance.detach() ** beta 13 14 return loss.sum(axis=-1)
6https://github.com/shariqfarooq123/AdaBins

Review summaries:
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.

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**  For 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.