Original: O N THE P ITFALLS OF H ETEROSCEDASTIC U NCERTAINTY E STIMATION WITH P ROBABILISTIC N EURAL N ETWORKS
Revised: O N THE P ITFALLS OF H ETEROSCEDASTIC U NCERTAINTY E STIMATION WITH P ROBABILISTIC N EURAL N ETWORKS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Capturing aleatoric uncertainty is a critical part of many machine learning systems.
Revised: Capturing aleatoric uncertainty is a critical part of many machine learning systems.

Original: 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.
Revised: 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.

Original: In this work, we examine this approach and identify potential hazards associated with the use of log-likelihood in conjunction with gradient-based optimizers.
Revised: In this work, we examine this approach and identify potential hazards associated with the use of log-likelihood in conjunction with gradient-based optimizers.

Original: First, we present a synthetic example illustrating how this approach can lead to very poor but stable parameter estimates.
Revised: First, we present a synthetic example illustrating how this approach can lead to very poor but stable parameter estimates.

Original: Second, we identify the culprit to be the log-likelihood loss, along with certain conditions that exacerbate the issue.
Revised: Second, we identify the culprit to be the log-likelihood loss, along with certain conditions that exacerbate the issue.

Original: Third, we present an alternative formulation in which each data point’s contribution to the loss is weighted by the β -exponentiated variance estimate.
Revised: Third, we present an alternative formulation in which each data point’s contribution to the loss is weighted by the β -exponentiated variance estimate.

Original: We show that using an appropriate β largely mitigates the issue in our illustrative example.
Revised: We show that using an appropriate β largely mitigates the issue in our illustrative example.

Original: 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.
Revised: 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.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Endowing models with the ability to capture uncertainty is of crucial importance in machine learning.
Revised: Endowing models with the ability to capture uncertainty is of crucial importance in machine learning.

Original: Uncertainty can be usefully categorized into two main types: epistemic uncertainty and aleatoric uncertainty (Kiureghian & Ditlevsen, 2009).
Revised: Uncertainty can be usefully categorized into two main types: epistemic uncertainty and aleatoric uncertainty (Kiureghian & Ditlevsen, 2009).

Original: Epistemic uncertainty accounts for subjective uncertainty in the model, one that is reducible given sufﬁcient data.
Revised: Epistemic uncertainty accounts for subjective uncertainty in the model, one that is reducible given sufﬁcient data.

Original: By contrast, aleatoric uncertainty captures the stochasticity inherent in the observations and can itself be subdivided into homoscedastic and heteroscedastic uncertainty.
Revised: By contrast, aleatoric uncertainty captures the stochasticity inherent in the observations and can itself be subdivided into homoscedastic and heteroscedastic uncertainty.

Original: Homoscedastic uncertainty is noise that stays constant across different inputs, whereas heteroscedastic uncertainty is one that varies depending on the inputs to the model.
Revised: Homoscedastic uncertainty is noise that stays constant across different inputs, whereas heteroscedastic uncertainty is one that varies depending on the inputs to the model.

Original: There are well-established beneﬁts for modeling each type of uncertainty.
Revised: There are well-established beneﬁts for modeling each type of uncertainty.

Original: For instance, capturing epistemic uncertainty enables effective budgeted data collection in active learning (Gal et al., 2017), allows for efﬁcient exploration in reinforcement learning (Osband et al., 2016), and is indispensable in cost-sensitive decision making (Amodei et al., 2016).
Revised: For instance, capturing epistemic uncertainty enables effective budgeted data collection in active learning (Gal et al., 2017), allows for efﬁcient exploration in reinforcement learning (Osband et al., 2016), and is indispensable in cost-sensitive decision making (Amodei et al., 2016).

Original: On the other hand, quantifying aleatoric uncertainty enables learning dynamics models of stochastic processes (e.g. for model-based or ofﬂine reinforcement learning)
Revised: On the other hand, quantifying aleatoric uncertainty enables learning dynamics models of stochastic processes (e.g. for model-based or ofﬂine reinforcement learning)

Original: (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).
Revised: (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).

Original: Many modern applications of machine learning rely on neural networks and deep learning tools to achieve state-of-the-art performance.
Revised: Many modern applications of machine learning rely on neural networks and deep learning tools to achieve state-of-the-art performance.

Original: However, most neural network models are not readily equipped with a capacity for uncertainty estimation.
Revised: However, most neural network models are not readily equipped with a capacity for uncertainty estimation.

Original: To address this shortcoming, a multitude of approaches have been proposed in the past few years.
Revised: To address this shortcoming, a multitude of approaches have been proposed in the past few years.

Original: Nevertheless, the majority of work in this domain has focused on epistemic uncertainty quantiﬁcation (Blundell et al., 2015; Gal & Ghahramani, 2016) or uncertainty estimation for classiﬁcation (Hendrycks & Gimpel, 2017; Guo et al., 2017; Ovadia et al., 2019; Mukhoti et al., 2021).
Revised: Nevertheless, the majority of work in this domain has focused on epistemic uncertainty quantiﬁcation (Blundell et al., 2015; Gal & Ghahramani, 2016) or uncertainty estimation for classiﬁcation (Hendrycks & Gimpel, 2017; Guo et al., 2017; Ovadia et al., 2019; Mukhoti et al., 2021).

Original: In this work, we examine a common approach for quantifying aleatoric uncertainty in neural network regression.
Revised: We examine a common approach for quantifying aleatoric uncertainty in neural network regression.

Original: 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.
Revised: 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.

Original: Then, the parameters of the network can be learned using maximum likelihood estimation (MLE), i.e. by minimizing the negative 
Revised: Then, the parameters of the network can be learned using maximum likelihood estimation (MLE), i.e. by minimizing the negative log-likelihood (NLL)

Original: [Figure]
Revised: [Figure]

Original: log-likelihood (NLL) criterion using stochastic gradient descent.
Revised: criterion using stochastic gradient descent.

Original: 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 overconﬁdent variance estimates.
Revised: 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 overconﬁdent variance estimates.

Original: Whereas strategies have been proposed to alleviate this speciﬁc 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 ﬁts.
Revised: Whereas strategies have been proposed to alleviate this speciﬁc 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 ﬁts.

Original: In this work, we analyse and propose a simple modiﬁcation to mitigate this issue.
Revised: In this work, we analyse and propose a simple modiﬁcation to mitigate this issue.

Original: 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.
Revised: 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.

Original: 1 for an illustrative example).
Revised: 1 for an illustrative example).

Original: Our analysis identiﬁes the high dependence of the gradients on the predictive variance as the primary culprit.
Revised: The primary culprit is the high dependence of the gradients on the predictive variance .

Revised: 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.

Original: More precisely, 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.
Revised: 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.

Original: 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 .
Revised: 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 .

Original: This formulation subsumes the standard
Revised: This formulation subsumes the standard

Original: NLL loss for β = 0 and allows to lessen the dependency of gradients on the variance estimates for 0 < β ≤ 1 .
Revised: NLL loss for β = 0 and allows to lessen the dependency of gradients on the variance estimates for 0 < β ≤ 1 .

Original: 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.
Revised: 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.

Original: Finally, we empirically show that our modiﬁed loss formulation largely mitigates the issue of poor ﬁts, achieving considerable improvements on a set of standard benchmarks while exhibiting more robustness to network architecture and learning rate conﬁgurations.
Revised: Finally, we empirically show that our modiﬁed loss formulation largely mitigates the issue of poor ﬁts, achieving considerable improvements on a set of standard benchmarks while exhibiting more robustness to network architecture and learning rate conﬁgurations.

Original:  P RELIMINARIES
Revised:  P RELIMINARIES

Original: Let X, Y be two random variables describing the input and target, following the joint distribution
Revised: Let X, Y be two random variables describing the input and target, following the joint distribution

Original: P ( X, Y ) .
Revised: P ( X, Y ) .

Original: We assume that Y is conditionally independent given X and that it follows some probability distribution P ( Y | X ) .
Revised: We assume that Y is conditionally independent given X and that it follows some probability distribution P ( Y | X ) .

Original: In the following, we use the common assumption that Y is normally distributed given X ; i.e. P ( Y | X ) =
Revised: In the following, we use the common assumption that Y is normally distributed given X ; i.e. P ( Y | X ) =

Original: N ( µ ( X ) , σ 2 ( X )) , where µ : R M (cid:55)→ R and σ 2 : R M
Revised: N ( µ ( X ) , σ 2 ( X )) , where µ : R M (cid:55)→ R and σ 2 : R M

Original: (cid:55)→ R + are respectively the true input-dependent mean and variance functions.
Revised: (cid:55)→ R + are respectively the true input-dependent mean and variance functions.

Original: 1 Equivalently, we can write
Revised: 1 Equivalently, we can write

Original: Y = µ ( X )+ (cid:15) ( X ) , with (cid:15) ( X ) ∼ N (0 , σ 2 ( X )) ; i.e. Y is generated from X by µ ( X ) plus a zero-mean
Revised: Y = µ ( X )+ (cid:15) ( X ) , with (cid:15) ( X ) ∼ N (0 , σ 2 ( X )) ; i.e. Y is generated from X by µ ( X ) plus a zero-mean

Original: Gaussian noise with variance σ 2 ( X ) .

Original: This input-dependent variance quantiﬁes the heteroscedastic uncertainty, or input-dependent aleatoric uncertainty.

Original: [Figure]
Revised: [Figure]

Original: An initial inhomogeneous feature space granularity (explained inthe text)results early on in different ﬁtting quality.
Revised: An initial inhomogeneous feature space granularity (explained in text) early on in different ﬁtting quality.

Original: The implicit weighting of the squared error in NLL can be seen as a biased datasampling with p ( x ) ∼ 1 σ 2 ( x ) (see Eq. 6).
Revised: The implicit weighting of the squared error in NLL can be seen as a biased data Gaussian noise with variance σ 2 ( X

Original: Badly ﬁt parts are increasingly ignored during training. On the right, the effect of our solution (Eq. 9) on the relative importance of data points is shown.
Revised: -dependent quantiﬁes heteroscedastic

Original: To learn estimates ˆ µ ( X ) , ˆ σ 2 ( X ) of the true mean and variance functions, it is common to use a neural network f θ parameterized by θ .
Revised: To learn estimates ˆ µ ( X ) , ˆ σ 2 ( X ) of the true mean and variance functions, it is common to use a neural network f θ parameterized by θ .

Original: Here, ˆ µ ( X ) and ˆ σ 2 ( X ) can be outputs of the ﬁnal layer (Nix & Weigend, 1994) or use two completely separate networks (Detlefsen et al.,
Revised: Here, ˆ µ ( X ) and ˆ σ 2 ( X ) can be outputs of the ﬁnal layer (Nix & Weigend, 1994) or use two completely separate networks (Detlefsen et al.,

Original: 2019).The variance output is hereby constrained to the positive region using a suitable activation function, e.g. softplus.
Revised: 2019).The variance output is hereby constrained to the positive region using a suitable activation function, e.g. softplus.

Original: The optimal parameters θ ∗ NLL can then be found using maximum likelihood estimation (MLE) by minimizing the negative log-likelihood (NLL) criterion L NLL under the distribution P ( X, Y ) :
Revised: The optimal parameters θ ∗ NLL can then be found using maximum likelihood estimation (MLE) by minimizing the negative log-likelihood (NLL) criterion L NLL under the distribution P ( X, Y ) :

Original: [Equation]
Revised: [Equation]

Original: In contrast, standard regression minimizes the mean squared error (MSE)
Revised: In contrast, standard regression minimizes the mean squared error (MSE)

Original: L MSE :
Revised: L MSE :

Original: [Equation]
Revised: [Equation]

Original: In practice, Eq. 1 and Eq. 2 are optimized using stochastic gradient descent with batches of samples drawn from P ( X, Y ) .
Revised: In practice, Eq. 1 and Eq. 2 are optimized using stochastic gradient descent with batches of samples drawn from P ( X, Y ) .

Original: The gradients of L NLL with respect to ˆ µ ( X ) , ˆ σ 2 ( X ) are given by
Revised: The gradients of L NLL with respect to ˆ µ ( X ) , ˆ σ 2 ( X ) are given by

Original: [Equation]
Revised: [Equation]

Original:  A NALYSIS
Revised:  A NALYSIS

Original: We now return to the example of trying to ﬁt a sinusoidal function from Sec. 1.
Revised: We now return to the example of trying to ﬁt a sinusoidal function from Sec. 1.

Original: Recall from Fig.that using the Gaussian NLL as the objective resulted in a sub-optimal ﬁt.
Revised: Recall from Fig.that using the Gaussian NLL as the objective resulted in a suboptimal ﬁt.

Original: In contrast, using MSE as the objective, the model converged without problems in reasonable time.
Revised: In contrast, using MSE as the objective, the model converged without problems in reasonable time.

Original: We now analyze the reasons behind this surprising result.
Revised: We now analyze the reasons behind this surprising result.

Original: From Eq.
Revised: From Eq.

Original: 3, we see that the true mean µ ( X ) is a minimizer of the NLL loss.
Revised: 3, we see that the true mean µ ( X ) is a minimizer of the NLL loss.

Original: It thus becomes clear that a) the solution found in Fig.
Revised: It thus becomes clear that a) the solution found in Fig.

Original: 1 is not the optimal one, and b) the NLL objective should, in principle, drive ˆ µ ( X ) to the optimal solution µ ( X ) .
Revised: 1 is not the optimal one, and b) the NLL objective should, in principle, drive ˆ µ ( X ) to the optimal solution µ ( X ) .

Original: So why is the model not converging?
Revised: So why is the model not converging?

Original: We identify two main culprits for the non-convergence of the Gaussian NLL objective:
Revised: We identify two main culprits for the non-convergence of the Gaussian NLL objective:

Original: Initial ﬂatness of the feature space can create an undercomplex but locally stable mean ﬁt.
Revised: Initial ﬂatness of the feature space can create an undercomplex but locally stable mean ﬁt.

Original: This ﬁt results from local symmetries and requires a form of symmetry breaking to escape. 2.
Revised: This ﬁt results from local symmetries and requires a form of symmetry breaking to escape. 2.

Original: The NLL loss scales the gradient of badly-predicted points down relative to well-predicted points, effectively undersampling those points.
Revised: The NLL loss scales the gradient of badly-predicted points down relative to well-predicted points, effectively undersampling those points.

Original: This effect worsens as training progresses.
Revised: This effect worsens as training progresses.

Original: These

Original:  culprits and their effect on training are illustrated in Fig.
Revised: These culprits and their effect on training are illustrated in Fig.

Original: 2 (left).
Revised: 2 (left).

Original: If the network cannot ﬁt 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.
Revised: If the network cannot ﬁt 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.

Original: This leads to down-weighting the data from these regions, fueling a vicious circle of self-amplifying the increasingly imbalanced weighting.
Revised: This leads to down-weighting the data from these regions, fueling a vicious circle of self-amplifying the increasingly imbalanced weighting.

Original: In the following, we analyse these effects and their reasons in more detail.
Revised: In the following, we analyse these effects and their reasons in more detail.

Original: [Figure]
Revised: [Figure]

Original: 3. S YMMETRY AND F EATURE N ON -L INEARITY
Revised: 3. S YMMETRY AND F EATURE N ON -L INEARITY

Original: It is instructive to see how the model evolves over training.
Revised: It is instructive to see how the model evolves over training.

Original: Figure 3 shows snapshots of the model at different stages of training.
Revised: Figure 3 shows snapshots of the model at different stages of training.

Original: It can be seen that the network ﬁrst learns essentially the best linear ﬁt while adapting the variance to match the spread of points around the predicted mean.
Revised: It can be seen that the network ﬁrst learns essentially the best linear ﬁt while adapting the variance to match the spread of points around the predicted mean.

Original: After the linear ﬁt has been established, the situation is locally stable.
Revised: After the linear ﬁt has been established, the situation is locally stable.

Original: That is, due to the symmetries of errors below and above the mean ﬁt, there is no incentive to change the situation.
Revised: That is, due to the symmetries of errors below and above the mean ﬁt, there is no incentive to change the situation.

Original: Escaping this situation thus requires some form of symmetry breaking.
Revised: Escaping this situation thus requires some form of symmetry breaking.

Original: 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.
Revised: 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.

Original: In addition, we hypothesize that an important form of symmetry breaking is how non-linear the feature space of the network is locally, in order to create the necessary non-linear ﬁt.
Revised: In addition, we hypothesize that an important form of symmetry breaking is how non-linear the feature space of the network is locally, in order to create the necessary non-linear ﬁt.

Original: Let us consider the non-linearity of the feature space.
Revised: Let us consider the non-linearity of the feature space.

Original: This quantity is not easy to capture.
Revised: This quantity is not easy to capture.

Original: To approximate it, we compute how much the Jacobian J f 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 vriance:
Revised: To approximate it, we compute how much the Jacobian J f 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: V ( x [Equation] =

Original: [Equation]
Revised: [Equation]

Original:  [Figure]
Revised: | [Equation] [Figure]

Original: [Equation]
Revised: [Equation]

Original: Figure 4 visualizes the Jacobian variance over the input space as a function of the training progress.
Revised: Figure 4 visualizes the Jacobian variance over the input space as a function of the training progress.

Original: Although initially relatively ﬂat [Figure] it gets more ﬁne-granular in parts of the input space – the parts which are later well ﬁt.
Revised: Although initially relatively ﬂat, it gets more ﬁne-granular in parts of the input space – the parts which are later well ﬁt.

Original: The region that has a low Jacobian variance gets stuck in this conﬁguration, see Fig.
Revised: The region that has a low Jacobian variance gets stuck in this conﬁguration, see Fig.

Original: This provides evidence that the non-linearity of the feature space is important for the success or failure of learning.
Revised: This provides evidence that the non-linearity of the feature space is important for the success or failure of learning.

Original: However, it does not answer the question of why gradient descent on the NLL loss would not break out of this situation
Revised: However, it does not answer the question of why gradient descent on the NLL loss would not break out of this situation

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: 3. W EIGHTING BY V ARIANCE E FFECTIVELY U NDERSAMPLES
Revised: 3. W EIGHTING BY V ARIANCE E FFECTIVELY U NDERSAMPLES

Original: The answer lies in an imbalanced weighting of data points across the input space.
Revised: The answer lies in an imbalanced weighting of data points across the input space.

Original: Recall that the gradient ∇ µ L NLL of the NLL with respect to the mean scales the error ˆ µ ( X ) − Y by[Equation] X
Revised: Recall that the gradient ∇ µ L NLL of the NLL with respect to the mean scales the error ˆ µ ( X ) − Y by

Original: [Equation]
Revised: [Equation]

Original: (Eq. 3).
Revised: (Eq. 3).

Original: As symmetry is broken and the true function starts to be ﬁt locally, the variance quickly shrinks in these areas to match the reduced MSE.
Revised: As symmetry is broken and the true function starts to be ﬁt locally, the variance quickly shrinks in these areas to match the reduced MSE.

Original: If the variance is well-calibrated, the gradient becomes ˆ µ ( X ) − Y ˆ σ 2 ( X ) ≈ ˆ µ ( X ) − Y
Revised: If the variance is well-calibrated, the gradient becomes ˆ µ ( X ) − Y ˆ σ 2 ( X ) ≈ ˆ µ ( X ) − Y

Original: (ˆ µ ( X ) − Y
Revised: (ˆ µ ( X [Equation] Y

Original: [Equation]
Revised: [Equation]

Original: 2 = 1 ˆ µ ( X ) − Y .
Revised: 2 = 1 ˆ µ ( X ) − Y .

Original: That is, 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-amplifying scenario.
Revised: That is, 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-amplifying scenario.

Original: Thus, the NLL exhibits the opposite of the MSE loss’s well-known behavior to focus on high-error samples.
Revised: Thus, the NLL exhibits the opposite of the MSE loss’s well-known behavior to focus on high-error samples.

Original: If the true variance σ 2 on the well-ﬁt regions is small compared to the error on the badly-ﬁt regions, or if the ratio of well-ﬁt to badly-ﬁt points is high, this can completely prevent learning progress on the badly-ﬁt regions.
Revised: If the true variance σ 2 on the well-ﬁt regions is small compared to the error on the badly-ﬁt regions, or if the ratio of well-ﬁt to badly-ﬁt points is high, this can completely prevent learning progress on the badly-ﬁt regions.

Original: Another way to view this is to interpret the different weighting of points as changing the training distribution P
Revised: Another way to view this is to interpret the different weighting of points as changing the training distribution P

Original: ( X, Y )
Revised: ( X, Y )

Original: to a modiﬁed distribution ˜ P ( X, Y ) in which points with high error have a lower probability of getting sampled.
Revised: to a modiﬁed distribution ˜ P ( X, Y ) in which points with high error have a lower probability of getting sampled.

Original: This can be shown by deﬁning ˜ P ( X, Y ) =
Revised: This can be shown by deﬁning ˜ P ( X, Y ) =

Original: Z − 1 P ( X,Y ) σ 2 ( X ) , where

Original: Z  = (cid:82) P [Equation] x,Y ) σ 2 ( x ) d x  is a normalizing constant, and recognizing that the gradient of the NLL equals the gradient of the MSE loss in Eq.
Revised: Z − 1 P ( X,Y ) σ 2 ( X ) , where Z = (cid:82) P ( x,y ) σ 2 ( x ) d x d y is a normalizing constant, and recognizing that the gradient of the NLL is proportional to the gradient of the MSE loss in Eq.

Original: 2 under the modiﬁed data distribution ˜ P ( X, Y ) :
Revised: 2 under the modiﬁed data distribution ˜ P ( X, Y ) :

Original: [Equation]
Revised: [Equation]

Original: In Fig.
Revised: In Fig.

Original: 5, we plot ˜ P ( X, Y ) over training time for our sinusoidal example.
Revised: 5, we plot ˜ P ( X, Y ) over training time for our sinusoidal example.

Original: 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) .
Revised: 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).

Revised: We show that this behavior also carries over to a real-world dataset in Sec.

Original:  As a side note, optimization with SGD typically assumes a ﬁxed training distribution; training under shifting distributions may thus make optimization with SGD more difﬁcult.
Revised: A.2. As a side note, optimization with SGD typically assumes a ﬁxed training distribution; training under shifting distributions may thus make optimization with SGD more difﬁcult.

Original: It also was recently shown that such a shifting training setup can hurt generalization (Igl et al., 2021).
Revised: It also was recently shown that such a shifting training setup can hurt generalization (Igl et al., 2021).

Original: These two points may further contribute to the sub-optimal performance of probabilistic neural networks experienced in practice.
Revised: These two points may further contribute to the suboptimal performance of probabilistic neural networks experienced in practice.

Original: Sometimes “weighting-by-variance” is seen as a feature of the Gaussian NLL (Kendall & Gal, 2017), namely thatit introduces a self-regularizing property by allowing the network to “ignore” outlier points with high error.
Revised: 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.

Original: 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-ﬁt regions, as shown above.
Revised: 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-ﬁt regions, as shown above.

Original: In our solution, explained below, we allow to control the amount ofthis self-regularization.
Revised: In our method (Sec. 4), we allow to control the extent of self-regularization.

Original:  M ETHOD
Revised:  M ETHOD

Original: To mitigate these problems with NLL training, we develop a solution in the following.
Revised: To mitigate these problems with NLL training, we develop a solution in the following.

Original: We start with a method that uses standard squared error losses to estimate the moments of the distribution to be predicted, which we term “Moment Matching”.
Revised: 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).

Original: It ﬁxes the problem with premature convergence when using L NLL and yields consistent training results, but  leads to underestimated variances and will act as one of the baselines in the empirical evaluation.
Revised: This method ﬁxes the problem with premature convergence when using L NLL and yields consistent training results, but it also leads to underestimated variances and will act as one of the baselines in our empirical evaluations.

Original: We then consider the distribution of residual errors, which are drastically different between NLL training and MSE training.
Revised: In Sec. 4.2, we then consider the distribution of residual errors which are drastically different between NLL training and MSE training.

Original: This suggests a problem-speciﬁc design choice that leads us to our proposed  β − NLL introduced in Sec. 4.3.
Revised: This suggests a problem-speciﬁc design choice that leads us to our proposed solution, which we term β − NLL , in Sec. 4.3.

Original: It allows choosing a meaningful loss-interpolation between NLL and MSE while keeping calibrated uncertainty estimates.
Revised: This solution method allows choosing a meaningful loss-interpolation between NLL and MSE while keeping calibrated uncertainty estimates.

Original: 4. M OMENT M ATCHING
Revised: 4. M OMENT M ATCHING

Original: A natural solution to the problem of estimating distributions would be to estimate the sufﬁcient statistics of the distribution.
Revised: A natural solution to the problem of estimating distributions would be to estimate the sufﬁcient statistics of the distribution.

Original: In the case of the Gaussian distribution, these would be the ﬁrst two moments µ, σ 2 fully describing the distribution.
Revised: In the case of the Gaussian distribution, these would be the ﬁrst two moments µ, σ 2 fully describing the distribution.

Original: So why not deﬁning losses based on the moment estimators?
Revised: So why not deﬁning losses based on the moment estimators?

Original: Starting from the conditional mean of the predictions y : E [ Y | X ] , we can deﬁne the squared deviation as a loss and see that the MSE is an upper bound:
Revised: Starting from the conditional mean of the predictions y : E [ Y | X ] , we can deﬁne the squared deviation as a loss and see that the MSE is an upper bound:

Original: [Equation]
Revised: [Equation]

Original: Notice that L µMM is the standard MSE loss.
Revised: Notice that L µ MM is the standard MSE loss.

Original: Thus, ﬁtting the mean ˆ µ ( x ) can follow standard (nondistributional) regression procedures.
Revised: Thus, ﬁtting the mean ˆ µ ( x ) can follow standard (nondistributional) regression procedures.

Original: Interestingly, due to the moment-matching viewpoint, we[Figure] can analogously deﬁne a loss for the variance (second central moment).
Revised: Interestingly, due to the moment-matching viewpoint, we can analogously deﬁne a loss for the variance (second central moment).

Original: The variance is deﬁned as
Revised: The variance is deﬁned as

Original: E (cid:2) ( Y − µ ( X ) [Equation] 2 | X (cid:3) ; analogously, we get the following loss: 
Revised: E (cid:2) ( Y − µ ( X )) 2 | X (cid:3) ; analogously, we get the following loss: E (cid:104)(cid:0) ( Y − ˆ µ ( X )) 2 − ˆ σ 2 ( X ) (cid:1) 2 | X (cid:105) .

Original: [Equation]
Revised: [Figure]

Original: To use the same physical unit as L µMM , we reformulate it in terms of the standard deviation as:
Revised: To use the same physical unit as L µ MM , we reformulate it in terms of the standard deviation as:

Original: [Equation]
Revised: [Equation]

Original: Thus, the moment matching loss is simply given by the sum of both losses: L MM =
Revised: Thus, the moment matching loss is simply given by the sum of both losses: L MM =

Original: L µMM + L σMM .
Revised: L µ MM + L σ MM .

Original: 4. R ESIDUAL D ISTRIBUTION AND A LLOCATION OF F UNCTION A PPROXIMATOR C APACITY
Revised: 4. R ESIDUAL D ISTRIBUTION AND A LLOCATION OF F UNCTION A PPROXIMATOR C APACITY

Original: The moment matching loss gives the same weighting to all data points.
Revised: The moment matching loss gives the same weighting to all data points.

Original: Since we are dealing with function approximators, the question is whether their capacity should be used equally.
Revised: Since we are dealing with function approximators, the question is whether their capacity should be used equally.

Original: 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.
Revised: 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.

Original: This is appropriate in case these variances are caused by true aleatoric uncertainty in the data.
Revised: This is appropriate in case these variances are caused by true aleatoric uncertainty in the data.

Original: 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.
Revised: 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.

Original: The handling of these difﬁcult points is different for the different loss functions.
Revised: The handling of these difﬁcult points is different for the different loss functions.

Original: Figure 6 shows the distribution of the residuals for a dynamics prediction dataset containing easy and hard to model areas.
Revised: Figure 6 shows the distribution of the residuals for a dynamics prediction dataset containing easy and hard to model areas.

Original: We observe a drastic difference between the NLL loss and the MSE/moment matching loss.
Revised: We observe a drastic difference between the NLL loss and the MSE loss (or equivalently solely optimizing L µ MM ).

Original: How important are the data points with high uncertainty?
Revised: How important are the data points with high uncertainty?

Original: Are these outliers or stem from truly noisy regions?
Revised: Are these outliers or stem from truly noisy regions?

Original: In this case, we would be willing to allocate less of the function approximator’s capacity to them.
Revised: In this case, we would be willing to allocate less of the function approximator’s capacity to them.

Original: Or are these simply difﬁcult samples?
Revised: Or are these simply difﬁcult samples?

Original: 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).
Revised: 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).

Original: Following this insight, this suggests that there is no one-loss-ﬁts-all, but rather that the modeler should have the choice to select which behavior is desired for the application at hand.
Revised: Following this insight, this suggests that there is no one-loss-ﬁts-all, but rather that the modeler should have the choice to select which behavior is desired for the application at hand.

Original: 4. V ARIANCE - WEIGHTING THE GRADIENTS OF THE NLL
Revised: 4. V ARIANCE - WEIGHTING THE GRADIENTS OF THE NLL

Original: In this section, we introduce our proposed solution to the problems occurring when training with
Revised: In this section, we introduce our proposed solution to the problems occurring when training with

Original: L NLL .
Revised: L NLL .

Original: The most important problem to solve is the premature convergence of NLL training to highly suboptimal mean ﬁts.
Revised: The most important problem to solve is the premature convergence of NLL training to highly suboptimal mean ﬁts.

Original: We identiﬁed in Sec. 3 the main reason being the relative down-weighting of badly ﬁt data points with its self-amplifying character.
Revised: In Sec. 3, we identiﬁed the relative down-weighting of badly ﬁt data points with its self-amplifying character as the main reason for this.

Original: 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 σ(Sec. 3.2).
Revised: 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).

Original: Consequently, we propose modifying this distribution by introducing a parameter β allowing to interpolate between NLL’s and a completely uniform data point importance.
Revised: Consequently, we propose modifying this distribution by introducing a parameter β allowing to interpolate between NLL’s and a completely uniform data point importance.

Original: The resulting sampling distribution is given by P ( x ) ∝ σ 2 β σ 2 and illustrated in Fig.
Revised: The resulting sampling distribution is given by P ( x ) ∝ σ 2 β σ 2 and illustrated in Fig.

Original: 2 (right).
Revised: 2 (right).

Original: How could this weighting be achieved?
Revised: How could this weighting be achieved?

Original: We simply introduce the variance-weighting term σ 2 β to the
Revised: We simply introduce the variance-weighting term σ 2 β to the

Original: L NLL loss such that it acts as a factor on the gradient.
Revised: L NLL loss such that it acts as a factor on the gradient.

Original: We denote the resulting loss as β − NLL given [Figure] by:
Revised: We denote the resulting loss as β − NLL :

Original: [Equation]
Revised: [Equation]

Original: where (cid:98)·(cid:99) denotes the stop gradient operation. The stop gradient makes the variance weighting term to be treated as a prefactor or learning rate. In this way, the gradients of L β − NLL are:
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: Naturally, for β = 0 , we recover the original NLL loss.
Revised: Naturally, for β = 0 , we recover the original NLL loss.

Original: For β = 1 the gradient w.r.t.
Revised: For β = 1 the gradient w.r.t.

Original: µ in Eq.
Revised: µ in Eq.

Original: 10 is equivalent to the one of MSE.
Revised: 10 is equivalent to the one of MSE.

Original: However, for the variance, the gradient in Eq.
Revised: However, for the variance, the gradient in Eq.

Original: 11 is a new quantity with 2 σ 2 in the denominator.
Revised: 11 is a new quantity with 2 σ 2 in the denominator.

Original: This is in contrast to the moment matching loss (Eq.
Revised: This is in contrast to the moment matching loss (Eq.

Original: For values 0 < β < 1 , we get different loss interpolations.
Revised: For values 0 < β < 1 , we get different loss interpolations.

Original: Particularly interesting is the case of β = 0 .
Revised: Particularly interesting is the case of β = 0 .

Original: 5 , where the data points are weighted with 1 σ (inverse standard deviation instead of inverse variance).
Revised: 5 , where the data points are weighted with 1 σ (inverse standard deviation instead of inverse variance).

Original: In Sec. 5, we show that β
Revised: In Sec. 5, we show that β

Original: = 0 .
Revised: = 0 .

Original: 5 generally achieves the best trade-off between accuracy and log-likelihood.
Revised: 5 generally achieves the best trade-off between accuracy and log-likelihood.

Original: As a remark, the new loss L β − NLL is not meant for performance evaluation – it is designed to result in meaningful gradients.
Revised: As a remark, the new loss L β − NLL is not meant for performance evaluation – it is designed to result in meaningful gradients.

Original: We found that its absolute value can be counter-intuitive due to the prefactor.
Revised: We found that its absolute value can be counter-intuitive due to the prefactor.

Original: The model performance during training should be monitored with the original (negative) log-likelihood and optionally the RMSE for the mean ﬁt.
Revised: The model performance during training should be monitored with the original (negative) log-likelihood and optionally the RMSE for the mean ﬁt.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We investigate the proposed changes in the loss function for training neural networks to predict a
Revised: We investigate the proposed changes in the loss function for training neural networks to predict a

Original: Gaussian distribution on toy examples and real-world datasets, the UCIdataset, and challenging dynamics model datasets.
Revised: Gaussian distribution on synthetic examples and real-world datasets featuring the UCI and challenging dynamics model datasets.

Original: As motivated above, the standard NLL loss often leads to suboptimal ﬁts, for instance, where the mean ﬁt is far from its desired quality.
Revised: As motivated above, the standard NLL loss often leads to suboptimal ﬁts, e.g. where the mean ﬁt is far from ideal.

Original: 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).
Revised: 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).

Original: 5. S YNTHETIC D ATASETS
Revised: 5. S YNTHETIC D ATASETS

Original: Sinusoidal without heteroscedastic noise We ﬁrst consider the dataset already used as an illustrative example in Fig.
Revised: Sinusoidal without heteroscedastic noise We ﬁrst consider our illustrative example of Fig.

Original: 1 – a simple sine curve with  small additive noise: y = sin( x ) + ξ , with ξ being Gaussian noise with standard deviation σ = 0 .
Revised: 1 – a simple sine curve with a small additive noise: y = sin( x ) + ξ , with ξ being Gaussian noise with standard deviation σ = 0 .

Original: One would expect that a network with sufﬁcient capacity can easily learn to ﬁt this function.
Revised: One would expect that a network with sufﬁcient capacity can easily learn to ﬁt this function.

Revised: Figure 7 inspects this over a range of architectures and learning rates.

Original: We ﬁnd that for the standard NLL ( β = 0 ) the network does not converge to a reasonable mean ﬁt.
Revised: We ﬁnd that for the standard NLL ( β = 0 ) the network does not converge to a reasonable mean ﬁt.

Original: There is a trend that larger networks and learning rates show better results, but when comparing this to β − NLL with β = 0 .
Revised: There is a trend that larger networks and learning rates show better results, but when comparing this to β − NLL with β = 0 .

Original: 5 we see that the networks are indeed able to ﬁt the function without problems.
Revised: 5 we see that the networks are indeed able to ﬁt the function without problems.

Original: As expected, the same holds for the mean squared error loss (MSE) which is recovered for β = 1 .
Revised: As expected, the same holds for the mean squared error loss (MSE) and β − NLL with β = 1 .

Original: The quality of the ﬁt w.r.t.
Revised: The quality of the ﬁt w.r.t.

Original: NLL is shown in Fig.
Revised: NLL is shown in Fig.

Original: S1 for completeness.
Revised: S3 for completeness.

Original: [Figure]

Original: [Table]

Original: Sinusoidal with heteroscedastic noise A commonly used illustrative example introduced in Detlefsen et al.
Revised: Sinusoidal with heteroscedastic noise A commonly used illustrative example introduced in Detlefsen et al.

Original: (2019) is a sine curve with increasing amplitude and noise:
Revised: (2019) is a sine curve with increasing amplitude and noise:

Original: y = x sin( x ) + xξ 1 + ξ 2 , with ξ  being Gaussian noise with standard deviation σ = 0 .
Revised: y = x sin( x ) + xξ 1 + ξ 2 , with ξ 1 and ξ 2 being Gaussian noise with standard deviation σ = 0 .

Original: The functional form is much easier, so ﬁtting the mean is achieved with all losses.
Revised: The functional form is much easier, so ﬁtting the mean is achieved with all losses.

Original: Here we sanity-check that the newly introduced loss isstill delivering good uncertainty estimates.
Revised: Here we sanity-check that the newly introduced loss is

Original: Figure 8 displays the predictions of the best models (w.r
Revised: [Figure] [Figure] 10  s t d gt std data std [Figure]

Original: [Table]
Revised: [Table]

Original: t.
Revised: still delivering good uncertainty estimates. Figure 8 (a-e) displays the predictions of the best models (w.r.t.

Original: NLL validation loss).
Revised: NLL validation loss) and (f) compares the predicted uncertainties over 10 independent trials.

Original: 5. UCI R EGRESSION D ATASETS
Revised: 5. UCI R EGRESSION D ATASETS

Original: As a standard real-world benchmark in predictive uncertainty estimation, we consider the UCI dataset (Hern´andez-Lobato & Adams, 2015).
Revised: As a standard real-world benchmark in predictive uncertainty estimation, we consider the UCI datasets (Hern´andez-Lobato & Adams, 2015).

Original: Table 1 reports the log-likelihoods of the test data comparingthe different loss variants.
Revised: Table 1 gives an overview comparing different loss variants.

Original: We report RMSE in Table S1.
Revised: We refer to Sec. A.3 for the full results on all 12 datasets.

Original: The results are encouraging: β − NLL achievesthe same predictive log-likelihood as the NLL lossor better, while clearly improving the predictive accuracy on most datasets.
Revised: 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.

Original: 5. D YNAMICS MODELS
Revised: 5. D YNAMICS MODELS

Original: 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 difﬁculty, ObjectSlide and Fetch-PickAndPlace.
Revised: 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 difﬁculty, ObjectSlide and Fetch-PickAndPlace.

Original: In both tasks, the goal is to predict how an object will move from the current state and the agent’s action.
Revised: In both tasks, the goal is to predict how an object will move from the current state and the agent’s action.

Original: Whereas ObjectSlide (Seitzer et al., 2021) is a simple 1D-environment,
Revised: Whereas ObjectSlide (Seitzer et al., 2021) is a simple 1D-environment,

Original: Fetch-PickAndPlace (Plappert et al., 2018) is a complex 3D robotic-manipulation environment.
Revised: Fetch-PickAndPlace (Plappert et al., 2018) is a complex 3D robotic-manipulation environment.

Original: The models are trained on state-action trajectories collected byrunning RL agents.
Revised: The models are trained on state-action trajectories collected by RL agents.

Original: See Sec.
Revised: See Sec.

Original: B for more details.
Revised: B for more details.

Original: For both datasets and all loss functions, we ﬁrst perform a grid search over different hyperparameter conﬁgurations (see  C.2).
Revised: For both datasets and all loss functions, we ﬁrst perform a grid search over different hyperparameter conﬁgurations (see Sec. C.2).

Original: The results also serve for a sensitivity analysis presented in Fig.
Revised: The results also serve for a sensitivity analysis presented in Fig.

Original: 9: how much does one have to tune hyperparameters for each loss function to work well in practice?
Revised: 9: how much does one have to tune hyperparameters for each loss function to work well in practice?

Original: We then take the best performing conﬁgurations and evaluate them on a hold-out test set (Table 2).
Revised: We then take the best performing conﬁgurations and evaluate them on a hold-out test set (Table 2).

Original: The sensitivity analysis reveals that both the NLL- and the MM-loss are vulnerable to hyperparameter settings, whereas the β − NLL loss achieves good results over a wide range of settings.
Revised: The sensitivity analysis reveals that both the NLL- and the MM-loss are vulnerable to hyperparameter settings, whereas the β − NLL loss achieves good results over a wide range of settings.

Original: When focusing on the best found model conﬁgurations for each loss in 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 ﬁts.
Revised: When focusing on the best found model conﬁgurations for each loss in 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 MMloss results in poor log-likelihood ﬁts.

Original: Our method performs well on both accuracy and log-likelihood ﬁt, for a range of β -values, where β
Revised: Our method performs well on both accuracy and log-likelihood ﬁt, for a range of β -values, where β

Original: = 0 .
Revised: = 0 .

Original: 5 generally achieves the best trade-off between both.
Revised: 5 generally achieves the best trade-off between both.

Original: [Figure]
Revised: [Figure]

Original: suffers due to high variance.

Original: [Table]
Revised: [Table]

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: We highlight a deep problem frequently occurring when optimizing probabilistic neural networks using the common NLL loss: training gets stuck in suboptimal function ﬁts.
Revised: We highlight a deep problem frequently occurring when optimizing probabilistic neural networks using the common NLL loss: training gets stuck in suboptimal function ﬁts.

Original: Due to the predictive variance-based weighting of mean prediction errors in the NLL loss, these bad ﬁts often stay unnoticed, unless the MSE is evaluated.
Revised: Due to the predictive variance-based weighting of mean prediction errors in the NLL loss, these bad ﬁts often stay unnoticed, unless the MSE is evaluated.

Original: 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 ﬁt regions (due to low feature non-linearity) get increasingly less weight and result in premature convergence.
Revised: 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 ﬁt regions (due to low feature non-linearity) get increasingly less weight and result in premature convergence.

Original: We propose a simple solution by introducing a family of loss functions called β − NLL .
Revised: We propose a simple solution by introducing a family of loss functions called β − NLL .

Original: Effectively, the gradient of the original NLL loss is scaled by the β -exponentiated per-sample variance.
Revised: Effectively, the gradient of the original NLL loss is scaled by the β -exponentiated per-sample variance.

Original: This allows for a meaningful interpolation between NLL loss and MSE loss while providing wellbehaved uncertainty estimates (even for the MSE case).
Revised: This allows for a meaningful interpolation between NLL loss and MSE loss while providing wellbehaved uncertainty estimates (even for the MSE case).

Original: The hyperparameter β gives practitioners the choice to control the self-regularization strength of NLL: how unimportant should high-noise regions or difﬁcult-to-predict data points be in the ﬁtting process.
Revised: The hyperparameter β gives practitioners the choice to control the self-regularization strength of NLL: how unimportant should high-noise regions or difﬁcult-to-predict data points be in the ﬁtting process.

Original: In most cases β = 0 .
Revised: In most cases, β = 0 .

Original: 5 will be a good starting point.
Revised: 5 will be a good starting point.

Original: We hope that our simple solution will improve the usability and performance of predictive uncertainty deep networks and increase their applicability.
Revised: We hope that our simple solution will improve the usability and performance of predictive uncertainty deep networks and increase their applicability.

Original: R EPRODUCIBILITY S TATEMENT
Revised: R EPRODUCIBILITY S TATEMENT

Original: All settings are described in detail in Sec.
Revised: All settings are described in detail in Sec.

Original: B and Sec.
Revised: B and Sec.

Original: C. We make full code and data available under https://sites.google.com/view/pitfalls-uncertainty .
Revised: C. We make full code and data available under https://sites.google.com/view/pitfalls-uncertainty .

Original: Both will be made publicly available after the review phase.
Revised: Both will be made publicly available after the review phase.

Original: A PPENDIX A A DDITIONAL R ESULTS 
Revised: A PPENDIX A A DDITIONAL R ESULTS A. S YNTHETIC D ATASET

Original: Synthetic Dataset – NLL ﬁt

Original: In Fig.
Revised: In Fig.

Original: S1, we provide the results for the NLL metric on the sinusoidal dataset.
Revised: S

Original: [Figure]
Revised: [Figure]

Original: UCI Dataset In addition to the results presented in the main text in Table 1 (where we report log-likelihoods of the predictions)
Revised: we replicate the experiment from Fig. 1 (training the NLL loss on a sinusoidal for 10updates) on several more random seeds. In Fig.

Revised: S

Original: [Figure]
Revised: [Figure]

Original: we include the achieved RMSE values in Table S
Revised: 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) [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: Also for β = 1 , which corresponds to MSE for ﬁtting the mean good uncertainty predictions are obtained with our β − NLL as testiﬁed by the low scores. A. A NALYSIS OF S AMPLING P ROBABILITIES ON F ETCH -P ICK A ND P LACE A. UCI D ATASETS In this section we include results for predictive log-likelihood (Table S1) and RMSE (Table S2) for all UCI datasets we evaluated on.

Original: [Table]
Revised: [Figure]

Original: B D ATASETS AND T RAINING S ETTINGS
Revised: B D ATASETS AND T RAINING S ETTINGS

Original: 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 .
Revised: 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 .

Original: 4 sin(2 πx )
Revised: 4 sin(2 πx )

Original: + ξ to them to create the targets y , where ξ is Gaussian noise with a standard deviation of 0 . 01 .
Revised: + ξ to them to create the targets y , where ξ is Gaussian noise with a standard deviation of 0 . 01 .

Original: Sinusoidal heteroscedastic noise We use the synthetic data as introduced in Detlefsen et al.
Revised: Sinusoidal heteroscedastic noise We use the synthetic data as introduced in Detlefsen et al.

Original: From the functional form y = x sin( x ) + xξ 1 + ξ 2 , with Gaussian noise with standard deviation σ = 0 .
Revised: From the functional form y = x sin( x ) + xξ 1 + ξ 2 , with Gaussian noise with standard deviation σ = 0 .

Original: 3 for ξ 1 and ξ 2 , we sample 500 points uniformly spaced in the interval [0 , 10] .
Revised: 3 for ξ 1 and ξ 2 , we sample 500 points uniformly spaced in the interval [0 , 10] .

Original: The model is a
Revised: The model is a

Original: MLP with one hidden layer with 50 units and tanh activation function (as used in Detlefsen et al.
Revised: MLP with one hidden layer with 50 units and tanh activation function (as used in Detlefsen et al.

Original: and Stirn & Knowles (2020)).
Revised: and Stirn & Knowles (2020)).

Original: UCI Datasets We use the UCI datasets suite commonly used to benchmark uncertainty estimation, stemming from the UCI Machine Learning Repository [-
-]
Revised: UCI Datasets We use the UCI datasets suite commonly used to benchmark uncertainty estimation, stemming from the UCI Machine Learning Repository [-
-]3 .

Original: In particular, we use the training-test protocol from (Hern´andez-Lobato & Adams, 2015; Gal & Ghahramani, 2016), and their data splits 4 .
Revised: In particular, we use the training-test protocol from (Hern´andez-Lobato & Adams, 2015; Gal & Ghahramani, 2016), and their data splits 4 .

Original: Inputs and targets are whitened on the training set.
Revised: Inputs and targets are whitened on the training set.

Original: Metrics are reported in the original scale of the data.
Revised: Metrics are reported in the original scale of the data.

Original: Each dataset is divided into 20 randomly sampled train-test splits (80%-20%).
Revised: Each dataset is divided into 20 randomly sampled train-test splits (80%-20%).

Original: 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.
Revised: 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.

Original: 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 100updates 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.
Revised: 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.

Original: The reported performance and standard deviations are taken as averages over all test splits.
Revised: The reported performance and standard deviations are taken as averages over all test splits.

Original: Note that the performance we report is not comparable with other publications, as performance is known to differ strongly over different data splits.
Revised: Note that the performance we report is not comparable with other publications, as performance is known to differ strongly over different data splits.

Original: Some other works also perform early-stopping on the test set, which distorts the results.
Revised: Some other works also perform early-stopping on the test set, which distorts the results.

Original: As model, we use a single-layer relu hidden network with 50 neurons, aside from “protein”, where we use 100 neurons.
Revised: We use a single-layer relu hidden network with 50 neurons, aside from “protein”, where we useneurons.

Original: The batch size is 256 .
Revised: The batch size is 256 .

Revised: [Table]

Revised: 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”).

Revised: For this purpose, we performed a two-sided Kolmogorov-Smirnov with a p ≤ 0 .

Revised: 05 signiﬁcance level.

Original: ObjectSlide This environment consists of an agent whose task it is to slide an object to a target location (Seitzer et al., 2021).
Revised: ObjectSlide This environment consists of an agent whose task it is to slide an object to a target location (Seitzer et al., 2021).

Original: 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.
Revised: 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.

Original: The forward prediction task consists of predicting the change in object position in the next state from current state and action.
Revised: The forward prediction task consists of predicting the change in object position in the next state from current state and action.

Original: 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.
Revised: 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.

Original: Inputs and targets are whitened on the training set.
Revised: Inputs and targets are whitened on the training set.

Original: Metrics are reported in the original scale of the data.
Revised: Metrics are reported in the original scale of the data.

Original: 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. 
Revised: 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. [Table]

Original: Fetch-PickAndPlace We use the Fetch-PickAndPlace environment (Plappert et al., 2018) from
Revised: Fetch-PickAndPlace We use the Fetch-PickAndPlace environment (Plappert et al., 2018) from

Original: OpenAI Gym (Brockman et al., 2016) as a challenging real-world scenario.
Revised: OpenAI Gym (Brockman et al., 2016) as a challenging real-world scenario.

Original: The task of the agent is to use position-controlled 7 DoF robotic arm to lift an object to a target location in space.
Revised: The task of the agent is to use position-controlled 7 DoF robotic arm to lift an object to a target location in space.

Original: The state space is 25-dimensional and the action space is 4-dimensional.
Revised: The state space is 25-dimensional and the action space is 4-dimensional.

Original: As in ObjectSlide, the prediction task is to predict the 3D-dimensional change in object position from the current state and action.
Revised: As in ObjectSlide, the prediction task is to predict the 3D-dimensional change in object position from the current state and action.

Original: 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.
Revised: 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.

Original: Inputs and targets are whitened on the training set.
Revised: Inputs and targets are whitened on the training set.

Original: Metrics are reported in the original scale of the data.
Revised: Metrics are reported in the original scale of the data.

Original: We train for a maximum ofepochs with a batch size of 256 , and evaluate the model with the best validation log-likelihood on the test set afterwards.
Revised: We train for a maximum ofepochs with a batch size of 256 , and evaluate the model with the best validation log-likelihood on the test set afterwards.

Original: C H YPERPARAMETER S ETTINGS AND I MPLEMENTATION D ETAILS
Revised: C H YPERPARAMETER S ETTINGS AND I MPLEMENTATION D ETAILS

Original: For all experiments, we used the Adam optimizer (Kingma & Ba, 2015) with standard settings β 1 = 0 .
Revised: For all experiments, we used the Adam optimizer (Kingma & Ba, 2015) with standard settings β 1 = 0 .

Original: 9 , β 2 = 0 .
Revised: 9 , β 2 = 0 .

Original: We parametrize the Gaussian distribution using two linear layers on top of shared features produced by a MLP.
Revised: We parametrize the Gaussian distribution using two linear layers on top of shared features produced by a MLP.

Original: The variance ˆ σ 2 ( x ) is constrained to the positive region using the softplus( x ) =
Revised: The variance ˆ σ 2 ( x ) is constrained to the positive region using the softplus( x ) =

Original: log(1 + exp( x )) activation function We additionally add a small constant of 10 −to prevent the variance from collapsing to zero, and clamp the maximum variance to 1 000.
Revised: 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.

Original: C. S INUSOIDAL R EGRESSION P ROBLEM
Revised: C. S INUSOIDAL R EGRESSION P ROBLEM

Original: The sinusoidal ﬁt in Fig.
Revised: The sinusoidal ﬁt in Fig.

Original: 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 .
Revised: 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 .

Original: For the experiment in
Revised: For the experiment in

Original: Sec. 5.1, we scan over learning rates and architectures with different hidden layers and units per layer, as detailed in Table S2.
Revised: Sec. 5.1, we scan over learning rates and architectures with different hidden layers and units per layer, as detailed in Table S3.

Original: [Table]
Revised: [Table]

Original: C. O BJECT S LIDE AND F ETCH -P ICK A ND P LACE

Original: For each tested loss function, we performed a grid search for the ObjectSlide and Fetch-PickAndPlace datasets, We report the parameters in Table S3.
Revised: For each tested loss function, we performed a grid search for the ObjectSlide and Fetch-PickAndPlace datasets, We report the parameters in Table S4.

Original: For the results in Table 2, we retrained the model conﬁgurations with the best validation log-likelihood on the grid search, and evaluated them on the hold-out test set.
Revised: For the results in Table 2, we retrained the model conﬁgurations with the best validation log-likelihood on the grid search, and evaluated them on the hold-out test set.

Original: The best found conﬁguration on ObjectSlide coincided to 3 hidden layers with 128 neurons, relu activation and learning rate 0 .
Revised: The best found conﬁguration on ObjectSlide coincided to 3 hidden layers with 128 neurons, relu activation and learning rate 0 .

Original: The best found conﬁguration on Fetch-PickAndPlace coincided to 4 hidden layers with 128 neurons, relu activation and learning rates 0 .
Revised: The best found conﬁguration on Fetch-PickAndPlace coincided to 4 hidden layers with 128 neurons, relu activation and learning rates 0 .

Original: 0003 for L NLL , L β − NLLwith β = 0 .
Revised: 0003 for L NLL , L β − NLL

Original: 25 and β = 0 .

Original: 5 , and learning rate 0 .

Original: 001 for L MSE , L MM , L β − NLL with β = 0 .

Original: 75 and β = 1 . 0 .

Original: We run 96 conﬁgurations per loss function.

Original: [Table]
Revised: [Table]