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: 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, termed β − NLL , in which each data point’s contribution to the loss is weighted by the β -exponentiated variance estimate.
Revised: Third, we present an alternative formulation, termed β − NLL , 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 domains and tasks and show that it achieves considerable improvements and performs more robustly concerning hyperparameters, both in predictive RMSE and log-likelihood criteria.
Revised: Fourth, we evaluate this approach on a range of domains and tasks and show that it achieves considerable improvements and performs more robustly concerning hyperparameters, both in predictive RMSE and log-likelihood criteria.

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 categorized into two main types: epistemic uncertainty and aleatoric uncertainty (Kiureghian & Ditlevsen, 2009).
Revised: Uncertainty can be 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 corresponds to noise that is constant across the input space, whereas heteroscedastic uncertainty corresponds to noise that varies with the input.
Revised: Homoscedastic uncertainty corresponds to noise that is constant across the input space, whereas heteroscedastic uncertainty corresponds to noise that varies with the input.

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 of 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 of 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: 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 log-likelihood (NLL) criterion using stochastic gradient descent.
Revised: 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.

Original: This simple procedure, which is the de-facto standard (Nix & Weigend, 1994; Lakshminarayanan et al., 2017; Kendall & Gal, 2017; Chua et al., 2018), 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), 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; Stirn & Knowles, 2020), we argue that an equally important
Revised: Whereas strategies have been proposed to alleviate this speciﬁc issue (Detlefsen et al., 2019; Stirn & Knowles, 2020), we argue that an equally important

Original: [Figure]
Revised: [Figure]

Original: T a r
Revised: T a r

Original: [Figure]
Revised: [Figure]

Original: issue is that this procedure can additionally lead to subpar mean ﬁts .
Revised: issue is that this procedure can additionally lead to subpar mean ﬁts .

Original: In this work, we analyze and propose a simple modiﬁcation to mitigate this issue.
Revised: In this work, we analyze 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: The primary culprit is the high dependence of the gradients on the predictive variance .
Revised: The primary culprit is the high dependence of the gradients on the predictive variance .

Original: 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.
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: 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.
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 , termed β − NLL , 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 , termed β − NLL , 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 NLL loss for β = 0 and allows to lessen the dependency of gradients on the variance estimates for 0 < β ≤ 1 .
Revised: This formulation subsumes the standard 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 range of domains and tasks while exhibiting more robustness to hyperparameter 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 range of domains and tasks while exhibiting more robustness to hyperparameter 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 ) .
Revised: Gaussian noise with variance σ 2 ( X ) .

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

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., 2019).
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., 2019).

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

Original: [Figure]
Revised: [Figure]

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
Revised: In contrast

Original: [Equation]
Revised: [Equation]

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

Original: L MSE :
Revised: L MSE :

Original: [Equation]
Revised: [Equation]

Original: In practice
Revised: In practice

Original: [Equation]
Revised: [Equation]

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

Original: The gradients of L NLL w.r.t.
Revised: The gradients of L NLL w.r.t.

Original: (with respect to) ˆ µ ( X ) , ˆ σ 2 ( X ) are given by
Revised: (with respect to) ˆ µ ( X ) , ˆ σ 2 ( X ) are given by

Original: [Equation]
Revised: [Equation]

Original: (cid:35) .
Revised: (cid:35) .

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 suboptimal ﬁ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 to the optimal mean ﬁt in a reasonable time.
Revised: In contrast, using MSE as the objective the model converged to the optimal mean ﬁt in a 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 the minimizer of the NLL loss.
Revised: 3, we see that the true mean µ ( X ) is the 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 does the model not converge to the optimal solution?
Revised: So, why does the model not converge to the optimal solution?

Original: We identify two main culprits for this behavior of the Gaussian NLL objective:
Revised: We identify two main culprits for this behavior 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 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 would then perceive 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 would then perceive a high effective data variance.

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

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

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 during training as shown in Fig.
Revised: It is instructive to see how the model evolves during training as shown in Fig.

Original: The network ﬁrst learns essentially the best linear ﬁt while adapting the variance to match the residuals.
Revised: The network ﬁrst learns essentially the best linear ﬁt while adapting the variance to match the residuals.

Original: The situation
Revised: The situation

Original: [Figure]
Revised: [Figure]

Original: is locally stable.
Revised: 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: Symmetry breaking is required for further progress.
Revised: Symmetry breaking is required for further progress.

Original: One form of symmetry breaking comes with the inherent stochasticity of mini-batch sampling in SGD, or the natural asymmetries contained in the dataset due to, e.g., outliers.
Revised: One form of symmetry breaking comes with the inherent stochasticity of mini-batch sampling in SGD, or the natural asymmetries contained in the dataset due to, e.g., outliers.

Original: Moreover, we hypothesize that the local non-linearity of the feature space plays an important role in creating the necessary non-linear ﬁt.
Revised: Moreover, we hypothesize that the local non-linearity of the feature space plays an important role in creating 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 for a dataset D , we compute how much the Jacobian J f of the features f ( x ) w.r.t.
Revised: To approximate it for a dataset D , we compute how much the Jacobian J f of the features f ( x ) w.r.t.

Original: the input varies in an L 2 -ball with radius r around a point x , denoted as the Jacobian variance:
Revised: the input varies in an L 2 -ball with radius r around a point x , denoted as the Jacobian variance:

Original: V
Revised: V

Original: [Equation]
Revised: [Equation]

Original: ) =
Revised: ) =

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, it becomes more granular in parts of the input space, the parts which are later well ﬁt.
Revised: Although initially relatively ﬂat, it becomes more granular in parts of the input space, the parts which are later well ﬁt.

Original: The region with low Jacobian variance remains stuck in this conﬁguration (see
Revised: The region with low Jacobian variance remains stuck in this conﬁguration (see

Original: Fig. 1).
Revised: Fig. 1).

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

Original: However, why does gradient descent not break out of this situation?
Revised: However, why does gradient descent not break out of this situation?

Original: [Figure]
Revised: [Figure]

Original: the mean of matrix V ( x ) (see Eq. 5).
Revised: the mean of matrix V ( x ) (see Eq. 5).

Original: [Figure]
Revised: [Figure]

Original: 3. I NVERSE -V ARIANCE W EIGHTING E FFECTIVELY U NDERSAMPLES
Revised: 3. I NVERSE -V ARIANCE W EIGHTING 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 w.r.t.
Revised: Recall that the gradient ∇ ˆ µ L NLL of the NLL w.r.t.

Original: the mean scales the error ˆ µ ( X ) − Y by 1 ˆ σ 2 ( X )
Revised: the mean scales the error ˆ µ ( X ) − Y by 1 ˆ σ 2 ( X )

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 ) 2 = 1 ˆ µ ( X ) − Y .
Revised: (ˆ µ ( X ) − Y ) 2 = 1 ˆ µ ( X ) − Y .

Original: Data points with already low error will get their contribution in the batch gradient scaled up relatively to high error data points – “rich get richer” self-ampliﬁcation.
Revised: Data points with already low error will get their contribution in the batch gradient scaled up relatively to high error data points – “rich get richer” self-ampliﬁcation.

Original: Thus, NLL acts contrary to MSE which focuses on high-error samples.
Revised: Thus, NLL acts contrary to MSE which focuses on high-error samples.

Original: If the true variance σ 2 on the well-ﬁt regions is much smaller than the errors on the badly-ﬁt regions, or there are much more well-ﬁt than badly-ﬁt points, then learning progress is completely hindered on the badly-ﬁt regions.
Revised: If the true variance σ 2 on the well-ﬁt regions is much smaller than the errors on the badly-ﬁt regions, or there are much more well-ﬁt than badly-ﬁt points, then learning progress is completely hindered 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 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.
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).

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

Original: B.3.
Revised: B.3.

Original: Sometimes, “inverse-variance weighting” is seen as a feature of the Gaussian NLL (Kendall & Gal, 2017) which introduces a self-regularizing property by allowing the network to “ignore” outlier points with high error.
Revised: Sometimes, “inverse-variance weighting” is seen as a feature of the Gaussian NLL (Kendall & Gal, 2017) which 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 data-inherent unpredictability (noise), but it is undesirable if it causes premature convergence and ignorance of hardto-ﬁt regions, as shown above.
Revised: This can be desirable if the predicted variance corresponds to data-inherent unpredictability (noise), but it is undesirable if it causes premature convergence and ignorance of hardto-ﬁt regions, as shown above.

Original: In our method, we enable control over the extent of self-regularization.
Revised: In our method, we enable control over the extent of self-regularization.

Original:  M ETHOD
Revised:  M ETHOD

Original: In this section, we develop a solution method to mitigate these issues with NLL training.
Revised: In this section, we develop a solution method to mitigate these issues with NLL training.

Original: Our approach, which we term β − NLL , allows choosing an arbitrary loss-interpolation between NLL and
Revised: Our approach, which we term β − NLL , allows choosing an arbitrary loss-interpolation between NLL and

Original: MSE while keeping calibrated uncertainty estimates.
Revised: MSE while keeping calibrated uncertainty estimates.

Original: 4. V ARIANCE -W EIGHTING THE G RADIENTS OF THE NLL
Revised: 4. V ARIANCE -W EIGHTING THE G RADIENTS OF THE NLL

Original: The problem we want to address is the premature convergence of NLL training to highly suboptimal mean ﬁts.
Revised: The problem we want to address is the premature convergence of NLL training to highly suboptimal mean ﬁts.

Original: In Sec. 3, we identiﬁed the relative down-weighting of badly-ﬁt data points in the NLL loss together with its self-amplifying characteristic as the main culprit.
Revised: In Sec. 3, we identiﬁed the relative down-weighting of badly-ﬁt data points in the NLL loss together with its self-amplifying characteristic as the main culprit.

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 σ 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 .

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 :
Revised: We denote the resulting loss as β − NLL :

Original: [Equation]
Revised: [Equation]

Original: where (cid:98)·(cid:99) denotes the stop gradient operation.
Revised: where (cid:98)·(cid:99) denotes the stop gradient operation.

Original: By stopping the gradient, the variance-weighting term acts as an adaptive, input-dependent learning rate.
Revised: By stopping the gradient, the variance-weighting term acts as an adaptive, input-dependent learning rate.

Original: In this way, the gradients of L β − NLL are:
Revised: In this way, the gradients of L β − NLL are:

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: 8 is equivalent to the one of MSE.
Revised: 8 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: 9 is a new quantity with 2 σ 2 in the denominator.
Revised: 9 is a new quantity with 2 σ 2 in the denominator.

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 our experiments (Sec. 5), we ﬁnd that β
Revised: In our experiments (Sec. 5), we ﬁnd 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: A Pytorch implementation of the loss function is provided in Sec.
Revised: A Pytorch implementation of the loss function is provided in Sec.

Original: D.5.
Revised: D.5.

Original: Note that the new loss L β − NLL is not meant for performance evaluation, rather it is designed to result in meaningful gradients.
Revised: Note that the new loss L β − NLL is not meant for performance evaluation, rather it is designed to result in meaningful gradients.

Original: Due to the weighting term, the loss value does not reﬂect the model’s quality.
Revised: Due to the weighting term, the loss value does not reﬂect the model’s quality.

Original: The model performance during training should be monitored with the original negative log-likelihood objective and, optionally, with RMSE for testing the quality of the mean ﬁt.
Revised: The model performance during training should be monitored with the original negative log-likelihood objective and, optionally, with RMSE for testing the quality of the mean ﬁt.

Original: [Figure]
Revised: [Figure]

Original: 4. A LLOCATION OF F UNCTION A PPROXIMATOR C APACITY
Revised: 4. A LLOCATION OF F UNCTION A PPROXIMATOR C APACITY

Original: Even though L MSE , L NLL , and L β − NLL all have the same optima w.r.t.
Revised: Even though L MSE , L NLL , and L β − NLL all have the same optima w.r.t.

Original: the mean (and also the variance in the case of L NLL and L β − NLL ), optimizing them leads to very different solutions.
Revised: the mean (and also the variance in the case of L NLL and L β − NLL ), optimizing them leads to very different solutions.

Original: In particular, because these losses weight data points differently, they assign the capacity of the function approximator differently .
Revised: In particular, because these losses weight data points differently, they assign the capacity of the function approximator differently .

Original: Whereas the MSE loss gives the same weighting to all data points, the
Revised: Whereas the MSE loss gives the same weighting to all data points, the

Original: NLL loss gives high weight to data points with low predicted variance and low weight to those with high variance.
Revised: NLL loss gives high weight to data points with low predicted variance and low weight to those with high variance.

Original: β − NLL interpolates between the two.
Revised: β − NLL interpolates between the two.

Original: The behavior of the NLL loss is appropriate if these variances are caused by true aleatoric uncertainty in the data.
Revised: The behavior of the NLL loss is appropriate if these variances are caused by true aleatoric uncertainty in the data.

Original: However, due to the use of function approximation, there is also the case where data points cannot be well predicted (maybe only transiently).
Revised: However, due to the use of function approximation, there is also the case where data points cannot be well predicted (maybe only transiently).

Original: This would result in high predicted variance, although the ground truth is corrupted by little noise.
Revised: This would result in high predicted variance, although the ground truth is corrupted by little noise.

Original: The different loss functions thus vary in how they handle these difﬁcult data points .
Revised: The different loss functions thus vary in how they handle these difﬁcult data points .

Original: An example of how the differences between the losses manifest in practice is illustrated in Fig.
Revised: An example of how the differences between the losses manifest in practice is illustrated in Fig.

Original: Here we show the distribution of the residuals for a dynamics prediction dataset containing easy and hard to model areas.
Revised: Here we show the distribution of the residuals for a dynamics prediction dataset containing easy and hard to model areas.

Original: The NLL loss essentially ignores a fraction of the data by predicting high uncertainty.
Revised: The NLL loss essentially ignores a fraction of the data by predicting high uncertainty.

Original: By analyzing those data points, we found that they were actually the most important data points to model correctly (because they captured non-trivial interactions in the physical world).
Revised: By analyzing those data points, we found that they were actually the most important data points to model correctly (because they captured non-trivial interactions in the physical world).

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

Original: Are they outliers (i.e. do they stem from truly noisy regions), to which we would be willing to allocate less of the function approximator’s capacity?
Revised: Are they outliers (i.e. do they stem from truly noisy regions), to which we would be willing to allocate less of the function approximator’s capacity?

Original: Or are they just difﬁcult samples that are important to ﬁt correctly?
Revised: Or are they just difﬁcult samples that are important to ﬁt correctly?

Original: The answer is taskdependent and, as such, there is no one-loss-ﬁts-all solution.
Revised: The answer is taskdependent and, as such, there is no one-loss-ﬁts-all solution.

Original: Rather, the modeler should choose which behavior is desired.
Revised: Rather, the modeler should choose which behavior is desired.

Original: Our β − NLL loss makes this choice available through the β parameter.
Revised: Our β − NLL loss makes this choice available through the β parameter.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In our experiments, we ask the following questions and draw the following conclusions:
Revised: In our experiments, we ask the following questions and draw the following conclusions:

Original: Sec. 5.1: Does β − NLL ﬁx the pitfall with NLL’s convergence?
Revised: Sec. 5.1: Does β − NLL ﬁx the pitfall with NLL’s convergence?

Original: Yes , β − NLL converges to good mean and uncertainty estimates across a range of β values.
Revised: Yes , β − NLL converges to good mean and uncertainty estimates across a range of β values.

Original: Sec.
Revised: Sec.

Original: 5.2: Does β − NLL improve over NLL in practical settings?
Revised: 5.2: Does β − NLL improve over NLL in practical settings?

Original: How sensitive to hyperparameters is β − NLL ?
Revised: How sensitive to hyperparameters is β − NLL ?

Original: We investigate a diverse set of real-world domains: regression on the UCI datasets, dynamics model learning, generative modeling on MNIST and Fashion-MNIST, and depth-map prediction from natural images.
Revised: We investigate a diverse set of real-world domains: regression on the UCI datasets, dynamics model learning, generative modeling on MNIST and Fashion-MNIST, and depth-map prediction from natural images.

Original: Yes , β − NLL generally performs better than NLL and is considerably easier to tune.
Revised: Yes , β − NLL generally performs better than NLL and is considerably easier to tune.

Original: Sec.
Revised: Sec.

Original: 5.3: How does β − NLL compare to other loss functions for distributional regression?
Revised: 5.3: How does β − NLL compare to other loss functions for distributional regression?

Original: We compare with a range of approaches: learning to match the moments of a Gaussian (termed “moment matching” (MM); see Sec.
Revised: We compare with a range of approaches: learning to match the moments of a Gaussian (termed “moment matching” (MM); see Sec.

Original: A), using a Student’s t-distribution instead of a Gaussian (Detlefsen et al., 2019), or putting different priors on the variance and using variational inference (xVAMP, xVAMP*, VBEM, VBEM*
Revised: A), using a Student’s t-distribution instead of a Gaussian (Detlefsen et al., 2019), or putting different priors on the variance and using variational inference (xVAMP, xVAMP*, VBEM, VBEM*

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

Original: It depends.
Revised: It depends.

Original: Different losses make different trade-offs, which we discuss in Sec.
Revised: Different losses make different trade-offs, which we discuss in Sec.

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

Original: estimates.
Revised: estimates.

Original: Moment matching (e) signiﬁcantly underestimates the variance everywhere.
Revised: Moment matching (e) signiﬁcantly underestimates the variance everywhere.

Original: We refer the reader to Sec.
Revised: We refer the reader to Sec.

Original: C and Sec.
Revised: C and Sec.

Original: D for a description of datasets and training settings.
Revised: D for a description of datasets and training settings.

Original: We evaluate the quality of predictions quantitatively in terms of the root mean squared error (RMSE) and the negative log-likelihood (NLL).
Revised: We evaluate the quality of predictions quantitatively in terms of 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 perform an extended investigation of our illustrative example from Fig.
Revised: Sinusoidal without heteroscedastic noise We ﬁrst perform an extended investigation of our illustrative example from Fig.

Original: 1 – a sine curve with a small additive noise: y = 0 .
Revised: 1 – a sine curve with a small additive noise: y = 0 .

Original: 4 sin(2 πx ) + ξ , with ξ being Gaussian noise with standard deviation σ = 0 .
Revised: 4 sin(2 π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.

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

Original: We ﬁnd that for the standard NLL loss ( β = 0 ), the networks do not converge to a reasonable mean ﬁt.
Revised: We ﬁnd that for the standard NLL loss ( β = 0 ), the networks do 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 any issues.
Revised: 5 we see that the networks are indeed able to ﬁt the function without any issues.

Original: As expected, the same holds for the mean squared error loss (MSE) and β − NLL with β = 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: S3.
Revised: S3.

Original: Sinusoidal with heteroscedastic noise We sanity-check that β − NLL is still delivering good uncertainty estimates on the illustrative example from Detlefsen et al.
Revised: Sinusoidal with heteroscedastic noise We sanity-check that β − NLL is still delivering good uncertainty estimates on the illustrative example from Detlefsen et al.

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

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

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

Original: Fitting the mean is achieved with all losses.
Revised: Fitting the mean is achieved with all losses.

Original: On the training range, β − NLL with β > 0 learns virtually the same uncertainties as the NLL loss.
Revised: On the training range, β − NLL with β > 0 learns virtually the same uncertainties as the NLL loss.

Original: 5. R EAL -W ORLD D ATASETS
Revised: 5. R EAL -W ORLD D ATASETS

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

Original: Table 1 gives an overview comparing different loss variants.
Revised: Table 1 gives an overview comparing different loss variants.

Original: We refer to Sec.
Revised: We refer to Sec.

Original: B.4 for the full results on all 12 datasets.
Revised: B.4 for the full results on all 12 datasets.

Original: The results are encouraging: β − NLL achieves predictive log-likelihoods on par with or better than the
Revised: The results are encouraging: β − NLL achieves predictive log-likelihoods on par with or better than the

Original: NLL loss while clearly improving the predictive accuracy on most datasets.
Revised: NLL loss while clearly improving the predictive accuracy on most datasets.

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: 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 difﬁculty, ObjectSlide, and Fetch-PickAndPlace.
Revised: 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 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, Fetch-PickAndPlace
Revised: Whereas ObjectSlide (Seitzer et al., 2021) is a simple 1D-environment, Fetch-PickAndPlace

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

Original: The models are trained on trajectories collected by RL agents.
Revised: The models are trained on trajectories collected by RL agents.

Original: For both datasets, we perform a grid search over different hyperparameter conﬁgurations (see Sec. D.2) for a sensitivity analysis to hyperparameters settings, presented in Fig.
Revised: For both datasets, we perform a grid search over different hyperparameter conﬁgurations (see Sec. D.2) for a sensitivity analysis to hyperparameters settings, presented in Fig.

Original: It reveals that NLL is vulnerable to the choice of hyperparameters, whereas β − NLL achieves good results over a wide range of conﬁgurations.
Revised: It reveals that NLL is vulnerable to the choice of hyperparameters, whereas β − NLL achieves good results over a wide range of conﬁgurations.

Original: The best performing conﬁgurations for each loss are then evaluated on a hold-out test set (Table 2).
Revised: The best performing conﬁgurations for each loss are then evaluated on a hold-out test set (Table 2).

Original: One can see that the NLL loss results in poor predictive performance and also exhibits quite a high variance across random seeds.
Revised: One can see that the NLL loss results in poor predictive performance and also exhibits quite a high variance across random seeds.

Original: Our method yields high accuracy and log-likelihood ﬁts for a range of β values, with β = 0 .
Revised: Our method yields high accuracy and log-likelihood ﬁts for a range of β values, with β = 0 .

Original: 5 generally achieving the best trade-off.
Revised: 5 generally achieving the best trade-off.

Original: Generative modeling and depth-map prediction For generative modeling, we train variational autoencoders (Kingma & Welling, 2014) with probabilistic decoders on MNIST and Fashion-MNIST.
Revised: Generative modeling and depth-map prediction For generative modeling, we train variational autoencoders (Kingma & Welling, 2014) with probabilistic decoders on MNIST and Fashion-MNIST.

Original: For the task of depth regression, we modify a state-of-the-art method (AdaBins; Bhat et al.
Revised: For the task of depth regression, we modify a state-of-the-art method (AdaBins; Bhat et al.

Original: and test it on the NYUv2 dataset (Silberman et al., 2012) with our loss (Fig. S6).
Revised: and test it on the NYUv2 dataset (Silberman et al., 2012) with our loss (Fig. S6).

Original: Table 3 presents selected results for both tasks, yielding similar trends as before.
Revised: Table 3 presents selected results for both tasks, yielding similar trends as before.

Original: We refer to Sec.
Revised: We refer to Sec.

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

Original: B.6 for more details, including qualitative results.
Revised: B.6 for more details, including qualitative results.

Original: 5. C OMPARISON TO O THER L OSS F UNCTIONS
Revised: 5. C OMPARISON TO O THER L OSS F UNCTIONS

Original: The previous sections have demonstrated that our β − NLL loss has clear advantages over the NLL loss.
Revised: The previous sections have demonstrated that our β − NLL loss has clear advantages over the NLL loss.

Original: However, the comparison to other loss functions requires a more nuanced discussion.
Revised: However, the comparison to other loss functions requires a more nuanced discussion.

Original: First,
Revised: First,

Original: [Table]
Revised: [Table]

Original: can be found in Table S3 and Table S
Revised: can be found in Table S3 and Table S

Original: [Table]
Revised: [Table]

Original: we also test an alternative loss function based on matching the moments of the Gaussian ( L MM ; see
Revised: we also test an alternative loss function based on matching the moments of the Gaussian ( L MM ; see

Original: Sec.
Revised: Sec.

Original: A).
Revised: A).

Original: While this loss results in high accuracy, it is unstable to train and exhibits poor likelihoods.
Revised: While this loss results in high accuracy, it is unstable to train and exhibits poor likelihoods.

Original: Second, we test several loss functions based on a Student’s t-distribution, including xVAMP and
Revised: Second, we test several loss functions based on a Student’s t-distribution, including xVAMP and

Original: VBEM (Stirn & Knowles, 2020).
Revised: VBEM (Stirn & Knowles, 2020).

Original: These approaches generally achieve better likelihoods than the β − NLL ; we conjecture this is because their ability to maintain uncertainty about the variance results in a better ﬁt (in terms of KL divergence) when the variance is wrongly estimated.
Revised: These approaches generally achieve better likelihoods than the β − NLL ; we conjecture this is because their ability to maintain uncertainty about the variance results in a better ﬁt (in terms of KL divergence) when the variance is wrongly estimated.

Original: In terms of predictive accuracy, β − NLL outperforms the Student’s t-based approaches.
Revised: In terms of predictive accuracy, β − NLL outperforms the Student’s t-based approaches.

Original: Exceptions are some of the UCI datasets with limited data (e.g. “concrete” and “yacht”) where xVAMP and VBEM are on par or better than β − NLL .
Revised: Exceptions are some of the UCI datasets with limited data (e.g. “concrete” and “yacht”) where xVAMP and VBEM are on par or better than β − NLL .

Original: This is likely because these methods can mitigate overﬁtting by placing a prior on the variance.
Revised: This is likely because these methods can mitigate overﬁtting by placing a prior on the variance.

Original: However, xVAMP and VBEM are also non-trivial to implement and computationally heavy: both need MC samples to evaluate the prior; for xVAMP, training time roughly doubles as evaluating the prior also requires a second forward pass through the network.
Revised: However, xVAMP and VBEM are also non-trivial to implement and computationally heavy: both need MC samples to evaluate the prior; for xVAMP, training time roughly doubles as evaluating the prior also requires a second forward pass through the network.

Original: In contrast, β − NLL is simple to implement (see Sec.
Revised: In contrast, β − NLL is simple to implement (see Sec.

Original: D.5)
Revised: D.5)

Original: and introduces no additional computational costs.
Revised: and introduces no additional computational costs.

Original:  C ONCLUSION
Revised:  C ONCLUSION

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

Original: With our analysis, we reveal the underlying reason: initially badly-ﬁt regions receive increasingly less weight in the loss which results in premature convergence.
Revised: With our analysis, we reveal the underlying reason: initially badly-ﬁt regions receive increasingly less weight in the loss which results 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 the NLL and MSE loss functions while providing well-behaved uncertainty estimates.
Revised: This allows for a meaningful interpolation between the NLL and MSE loss functions while providing well-behaved uncertainty estimates.

Original: The hyperparameter β gives practitioners the choice to control the self-regularization strength of NLL: how important 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 important 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 think the problem discussed in this paper is primarily why practitioners using the Gaussian distribution in regression or generative modeling tasks often opt for a constant or homoscedastic (global) variance, as opposed to the more general heteroscedastic (data-dependent) variance.
Revised: We think the problem discussed in this paper is primarily why practitioners using the Gaussian distribution in regression or generative modeling tasks often opt for a constant or homoscedastic (global) variance, as opposed to the more general heteroscedastic (data-dependent) variance.

Original: We hope that our simple solution contributes to changing this situation by improving the usability and performance of modeling data uncertainty with deep neural networks.
Revised: We hope that our simple solution contributes to changing this situation by improving the usability and performance of modeling data uncertainty with deep neural networks.

Original: A CKNOWLEDGMENTS
Revised: A CKNOWLEDGMENTS

Original: The authors thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Maximilian Seitzer.
Revised: The authors thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Maximilian Seitzer.

Original: Georg Martius is a member of the Machine Learning Cluster of Excellence, EXC number 2064/1 – Project number 390727645.
Revised: Georg Martius is a member of the Machine Learning Cluster of Excellence, EXC number 2064/1 – Project number 390727645.

Original: We acknowledge the ﬁnancial support from the German Federal Ministry of Education and Research (BMBF) through the T¨ubingen
Revised: We acknowledge the ﬁnancial support from the German Federal Ministry of Education and Research (BMBF) through the T¨ubingen

Original: AI Center (FKZ: 01IS18039B)
Revised: AI Center (FKZ: 01IS18039B)

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: C and Sec.
Revised: C and Sec.

Original: D. We make full code and data available under https://github.com/martius-lab/pitfalls-of-uncertainty .
Revised: D. We make full code and data available under https://github.com/martius-lab/beta-nll .

Original: A PPENDIX A T HE M OMENT M ATCHING L OSS
Revised: A PPENDIX A T HE M OMENT M ATCHING L OSS

Original: We also investigated an alternative loss function designed to counter the imbalanced weighting of data points by the NLL loss, which we call “moment matching” (MM).
Revised: We also investigated an alternative loss function designed to counter the imbalanced weighting of data points by the NLL loss, which we call “moment matching” (MM).

Original: It uses standard squared error losses to estimate the moments of the target distribution , i.e. directly estimating the sufﬁcient statistics of the target distribution.
Revised: It uses standard squared error losses to estimate the moments of the target distribution , i.e. directly estimating the sufﬁcient statistics of the target distribution.

Original: In our experiments, we found that this loss function generally ﬁxes the problems with premature convergence when using L NLL , but it also leads to underestimation of variances and exhibits considerable training instabilities.
Revised: In our experiments, we found that this loss function generally ﬁxes the problems with premature convergence when using L NLL , but it also leads to underestimation of variances and exhibits considerable training instabilities.

Original: To fully describe a Gaussian distribution, only the ﬁrst two moments µ, σ 2 need to be estimated.
Revised: To fully describe a Gaussian distribution, only the ﬁrst two moments µ, σ 2 need to be estimated.

Original: So, why not directly deﬁne losses based on the moment estimators?
Revised: So, why not directly deﬁne losses based on the moment estimators?

Original: Starting from the conditional mean of the targets, 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 targets, 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, learning a mean ﬁt ˆ µ ( x ) can follow the standard (non-distributional) regression procedures.
Revised: Thus, learning a mean ﬁt ˆ µ ( x ) can follow the standard (non-distributional) regression procedures.

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

Original: Variance is deﬁned as E (cid:2) ( Y − µ ( X )) 2 | X (cid:3) .
Revised: Variance is deﬁned as E (cid:2) ( Y − µ ( X )) 2 | X (cid:3) .

Original: As such, by analogy, we can deﬁne the following loss: L ˆ σ 2 MM :=
Revised: As such, by analogy, we can deﬁne the following loss: L ˆ σ 2 MM :=

Original: E (cid:104)(cid:0) ( Y − ˆ µ ( X )) 2 − ˆ σ 2 ( X ) (cid:1) 2 | X (cid:105) .
Revised: E (cid:104)(cid:0) ( Y − ˆ µ ( X )) 2 − ˆ σ 2 ( X ) (cid:1) 2 | X (cid:105) .

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 can be expressed simply by the sum of the two losses: L MM =
Revised: Thus, the moment matching loss can be expressed simply by the sum of the two losses: L MM =

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

Original: We found that using L ˆ σ MM instead of L ˆ σ 2 MM makes it easier to balance the losses.
Revised: We found that using L ˆ σ MM instead of L ˆ σ 2 MM makes it easier to balance the losses.

Original: Interestingly, our β − NLL subsumes the moment matching loss if we allow for different values of β to be used for mean and variance estimation.
Revised: Interestingly, our β − NLL subsumes the moment matching loss if we allow for different values of β to be used for mean and variance estimation.

Original: In particular, using β − NLL with β = 1 for the mean and β = 2 for the variance results in the same gradients as using 12 L ˆ µ MM + 14 L ˆ σ 2 MM as the loss.
Revised: In particular, using β − NLL with β = 1 for the mean and β = 2 for the variance results in the same gradients as using 12 L ˆ µ MM + 14 L ˆ σ 2 MM as the loss.

Original: Note that we did not investigate this connection further, and also did not perform any experiments with different values of β for mean and variance.
Revised: Note that we did not investigate this connection further, and also did not perform any experiments with different values of β for mean and variance.

Original: B A DDITIONAL R ESULTS B. F URTHER C ONSEQUENCES OF I NVERSE -V ARIANCE W EIGHTING
Revised: B A DDITIONAL R ESULTS B. F URTHER C ONSEQUENCES OF I NVERSE -V ARIANCE W EIGHTING

Original: In Sec. 3.2, we interpreted the inverse-variance weighting of the NLL as training under a different data distribution ˜ P ( X, Y ) in which points with high error have a low probability of being sampled.
Revised: In Sec. 3.2, we interpreted the inverse-variance weighting of the NLL as training under a different data distribution ˜ P ( X, Y ) in which points with high error have a low probability of being sampled.

Original: A consequence of this is that the distribution ˜ P ( X, Y ) continuously shifts while the model is improving its ﬁt.
Revised: A consequence of this is that the distribution ˜ P ( X, Y ) continuously shifts while the model is improving its ﬁt.

Original: For datasets where the noise level is low on parts of the input space and the underlying function can be modeled to high accuracy, an interesting phenomenon can be observed: even though the training curve for RMSE looks like it indicates convergence (i.e. it ﬂattens out), the model can actually still be adapting with the changing training distribution.
Revised: For datasets where the noise level is low on parts of the input space and the underlying function can be modeled to high accuracy, an interesting phenomenon can be observed: even though the training curve for RMSE looks like it indicates convergence (i.e. it ﬂattens out), the model can actually still be adapting with the changing training distribution.

Original: It is just that the changes happen on data points that already have such low prediction error relative to the average error that further improvements are virtually invisible in the average error.
Revised: It is just that the changes happen on data points that already have such low prediction error relative to the average error that further improvements are virtually invisible in the average error.

Original: The actual training progress can be revealed using a histogram of the prediction errors on a log-scale, in the manner of Fig.
Revised: The actual training progress can be revealed using a histogram of the prediction errors on a log-scale, in the manner of Fig.

Original: Eventually, the training distribution stabilizes when for all data points, either the prediction error reaches the noise level at that point, or the variance can not decrease further at that point because it reaches a manually set lower bound.
Revised: Eventually, the training distribution stabilizes when for all data points, either the prediction error reaches the noise level at that point, or the variance can not decrease further at that point because it reaches a manually set lower bound.

Original: If no lower bound on the variance is set, the distribution may never stabilize, leading to training instabilities when the variance reaches close to zero.
Revised: If no lower bound on the variance is set, the distribution may never stabilize, leading to training instabilities when the variance reaches close to zero.

Original: B. S YNTHETIC D ATASET
Revised: B. S YNTHETIC D ATASET

Original: In Fig.
Revised: In Fig.

Original: S
Revised: S

Original: [Figure]
Revised: [Figure]

Original: we replicate the experiment from Fig.
Revised: we replicate the experiment from Fig.

Original: 1 (training the NLL loss on a sinusoidal for 10updates) on several more random seeds.
Revised: 1 (training the NLL loss on a sinusoidal for 10updates) on several more random seeds.

Original: In order to test the dependence of our reported issue on the optimizer
Revised: In order to test the dependence of our reported issue on the optimizer

Original: [Figure]
Revised: [Figure]

Original: we repeat the experiment from Fig.
Revised: we repeat the experiment from Fig.

Original: but use different optimizers than Adam with β 1 = 0 .
Revised: but use different optimizers than Adam with β 1 = 0 .

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

Original: The results are shown in Fig. S2.
Revised: The results are shown in Fig. S2.

Original: We ﬁnd that none of the conﬁgurations reaches below an RMSE of 0 .
Revised: We ﬁnd that none of the conﬁgurations reaches below an RMSE of 0 .

Original: 1 (the optimal value corresponds to an RMSE of 0 . 01 ), indicating that the issue is occurring stably across optimization settings.
Revised: 1 (the optimal value corresponds to an RMSE of 0 . 01 ), indicating that the issue is occurring stably across optimization settings.

Original: In Fig.
Revised: In Fig.

Original: S
Revised: S

Original: [Figure]
Revised: [Figure]

Original: we provide the results for the
Revised: we provide the results for the

Original: NLL metric on the sinusoidal dataset, complementing the results of Fig.
Revised: NLL metric on the sinusoidal dataset, complementing the results of Fig.

Original: 7 (see discussion in Sec. 5.1).
Revised: 7 (see discussion in Sec. 5.1).

Original: [Figure]
Revised: [Figure]

Original: The observed behavior is stable across different independent trials.
Revised: The observed behavior is stable across different independent trials.

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

Original: B. A NALYSIS OF S AMPLING P ROBABILITIES ON F ETCH -P ICK A ND P LACE
Revised: B. A NALYSIS OF S AMPLING P ROBABILITIES ON F ETCH -P ICK A ND P LACE

Original: In Fig.
Revised: In Fig.

Original: S
Revised: S

Original: [Figure]
Revised: [Figure]

Original: we show how the analysis from Sec.
Revised: we show how the analysis from Sec.

Original: 3.2 transfers to a real world dataset, namely FetchPickAndPlace.
Revised: 3.2 transfers to a real world dataset, namely FetchPickAndPlace.

Original: The ﬁgure shows how the distribution of effective sampling probability evolves during training (over a ﬁxed 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.
Revised: The ﬁgure shows how the distribution of effective sampling probability evolves during training (over a ﬁxed 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.

Original: The mismatch between the two distributions demonstrates that optimizing the
Revised: The mismatch between the two distributions demonstrates that optimizing the

Original: NLL loss drastically undersamples in comparison to the reference, effectively never sampling some data points.
Revised: NLL loss drastically undersamples in comparison to the reference, effectively never sampling some data points.

Original: This further corroborates our analysis in Sec. 3.
Revised: This further corroborates our analysis in Sec. 3.

Original: [Figure]
Revised: [Figure]

Original: B. UCI D ATASETS
Revised: B. UCI D ATASETS

Original: In this section
Revised: In this section

Original: [Figure]
Revised: [Figure]

Original: we include results for predictive log-likelihood (Table S1) and RMSE (Table S2) for all UCI datasets we evaluated on.
Revised: we include results for predictive log-likelihood (Table S1) and RMSE (Table S2) for all UCI datasets we evaluated on.

Original: We ﬁnd that baselines based on the Student’s t-distribution (xVAMP and VBEM (Stirn & Knowles, 2020)) tend to have better predictive log-likelihood than β − NLL , although there are also datasets where β − NLL performs better (“energy” and “naval”) or where there is no statistically signiﬁcant improvement (“housing”, “kin8m”, and “wine-red”).
Revised: We ﬁnd that baselines based on the Student’s t-distribution (xVAMP and VBEM (Stirn & Knowles, 2020)) tend to have better predictive log-likelihood than β − NLL , although there is also a dataset where β − NLL performs better (“naval”) or where there is no statistically signiﬁcant improvement (“housing”, “kin8m”, “wine-red”, “wine-white”).

Original: For RMSE, L MSE unsurprisingly performs best.
Revised: For RMSE, L MSE unsurprisingly performs best.

Original: Our β − NLL is often on par with L MSE and on par or better than the other baselines, except for “yacht”.
Revised: Our β − NLL is often on par with L MSE and on par or better than the other baselines, except for “yacht”.

Original: At the same time, our method is very simple to implement (see Sec.
Revised: At the same time, our method is very simple to implement (see Sec.

Original: D.5)
Revised: D.5)

Original: and computationally lightweight compared to xVAMP and VBEM, which require the costly evaluation of a prior and
Revised: and computationally lightweight compared to xVAMP and VBEM, which require the costly evaluation of a prior and

Original: Monte-Carlo sampling at each training step.
Revised: Monte-Carlo sampling at each training step.

Original: B. G ENERATIVE M ODELING WITH V ARIATIONAL A UTOENCODERS
Revised: B. G ENERATIVE M ODELING WITH V ARIATIONAL A UTOENCODERS

Original: We test different loss functions on the task of generative modeling using variational autoencoders (VAEs) (Kingma & Welling, 2014).
Revised: We test different loss functions on the task of generative modeling using variational autoencoders (VAEs) (Kingma & Welling, 2014).

Original: To this end, we parameterize the decoder distribution p ( x | z ) with N (cid:0) µ ( z ) , σ 2 ( z ) (cid:1) , where the mean µ ( z ) and variance σ 2 ( z ) are outputs of a neural network.
Revised: To this end, we parameterize the decoder distribution p ( x | z ) with N (cid:0) µ ( z ) , σ 2 ( z ) (cid:1) , where the mean µ ( z ) and variance σ 2 ( z ) are outputs of a neural network.

Original: We train the VAE by maximizing the ELBO E q ( z | x ) [log p ( x | z )] − D KL ( q ( z | x ) || p ( z )) , plugging in different loss functions for log p ( x | z ) .
Revised: We train the VAE by maximizing the ELBO E q ( z | x ) [log p ( x | z )] − D KL ( q ( z | x ) || p ( z )) , plugging in different loss functions for log p ( x | z ) .

Original: Following Stirn & Knowles (2020), we evaluate the logposterior predictive likelihood log E q ( z | x ) [ p ( x | z )] .
Revised: Following Stirn & Knowles (2020), we evaluate the logposterior predictive likelihood log E q ( z | x ) [ p ( x | z )] .

Original: We approximate the expectation using a ﬁnite mixture of 20 Monte-Carlo samples from q ( z | x ) .
Revised: We approximate the expectation using a ﬁnite mixture of 20 Monte-Carlo samples from q ( z | x ) .

Original: To compute the RMSE, we take the mean of that mixture.
Revised: To compute the RMSE, we take the mean of that mixture.

Original: We compare β − NLL against L MM , L NLL with a ﬁxed variance of 1 , Student-t (Takahashi et al., 2018; Detlefsen et al., 2019), xVAMP and VBEM (Stirn & Knowles, 2020).
Revised: We compare β − NLL against L MM , L NLL with a ﬁxed variance of 1 , Student-t (Takahashi et al., 2018; Detlefsen et al., 2019), xVAMP and VBEM (Stirn & Knowles, 2020).

Original: We evaluate on MNIST and FashionMNIST.
Revised: We evaluate on MNIST and FashionMNIST.

Original: Table S3 presents quantitative results.
Revised: Table S3 presents quantitative results.

Original: VBEM achieves the best reconstruction error at the expense of poor log-likelihood.
Revised: VBEM achieves the best reconstruction error at the expense of poor log-likelihood.

Original: Vice versa, Student-t, xVAMP, xVAMP*, and VBEM* achieve strong log-likelihoods but worse reconstruction errors.
Revised: Vice versa, Student-t, xVAMP, xVAMP*, and VBEM* achieve strong log-likelihoods but worse reconstruction errors.

Original: β − NLL with β > 0 provides a good compromise between log-likelihood and reconstruction error.
Revised: β − NLL with β > 0 provides a good compromise between log-likelihood and reconstruction error.

Original: Figure S
Revised: Figure S

Original: [Table] [Table]
Revised: [Table] [Table]

Original: shows qualitative examples.
Revised: shows qualitative examples.

Original: Our β − NLL loss with β > 0 allows to learn good reconstructions and meaningful uncertainties.
Revised: Our β − NLL loss with β > 0 allows to learn good reconstructions and meaningful uncertainties.

Original: Moreover, images produced by sampling latents from the prior are semantically meaningful, indicating that adding our loss function does not break the disentangling properties of VAEs.
Revised: Moreover, images produced by sampling latents from the prior are semantically meaningful, indicating that adding our loss function does not break the disentangling properties of VAEs.

Original: Compare that to Student-t, xVAMP(*), and VBEM(*), which do not produce similarly clear images when sampling from the prior.
Revised: Compare that to Student-t, xVAMP(*), and VBEM(*), which do not produce similarly clear images when sampling from the prior.

Original: [Table]
Revised: [Table]

Original: B. D EPTH R EGRESSION
Revised: B. D EPTH R EGRESSION

Original: We evaluate β − NLL on the task of depth regression on the NYUv2 dataset (Silberman et al., 2012).
Revised: We evaluate β − NLL on the task of depth regression on the NYUv2 dataset (Silberman et al., 2012).

Original: For this purpose, we use a state-of-the-art method for depth regression, AdaBins (Bhat et al., 2021), and train it with different loss functions.
Revised: For this purpose, we use a state-of-the-art method for depth regression, AdaBins (Bhat et al., 2021), and train it with different loss functions.

Original: Note that we remove the Mini-ViT Transformer module from the model, thus our results are not directly comparable with those reported by Bhat et al.
Revised: Note that we remove the Mini-ViT Transformer module from the model, thus our results are not directly comparable with those reported by Bhat et al.

Original: Table S4 presents quantitative results.
Revised: Table S4 presents quantitative results.

Original: β − NLL with β > 0 achieves better RMSE than the NLL loss.
Revised: β − NLL with β > 0 achieves better RMSE than the NLL loss.

Original: β − NLL with β = 0 .
Revised: β − NLL with β = 0 .

Original: 5 again provides a good trade-off, achieving similar RMSE as L MSE and performing better on some of the other metrics.
Revised: 5 again provides a good trade-off, achieving similar RMSE as L MSE and performing better on some of the other metrics.

Original: Figure S6 shows qualitative examples.
Revised: Figure S6 shows qualitative examples.

Original: Compared to
Revised: Compared to

Original: L NLL , the depth maps predicted by β − NLL with β > 0 are noticeably sharper.
Revised: L NLL , the depth maps predicted by β − NLL with β > 0 are noticeably sharper.

Original: refer
Revised: refer

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: MNIST
Revised: MNIST

Original: Fashion-MNIST dataset
Revised: Fashion-MNIST dataset

Original: L β − NLL ( β = 0 ) meanstd ∼ posterior ∼ N (0 , 1)
Revised: L β − NLL ( β = 0 ) meanstd ∼ posterior ∼ N (0 , 1)

Original: L β − NLL ( β = 0 . 25 ) meanstd ∼ posterior ∼ N (0 , 1) L β − NLL ( β = 0 . 5 ) meanstd ∼ posterior ∼ N (0 , 1)
Revised: L β − NLL ( β = 0 . 25 ) meanstd ∼ posterior ∼ N (0 , 1) L β − NLL ( β = 0 . 5 ) meanstd ∼ posterior ∼ N (0 , 1)

Original: L β − NLL ( β = 0 . 75 ) meanstd ∼ posterior ∼ N (0 , 1) L β − NLL ( β = 1 . 0 ) meanstd ∼ posterior ∼ N (0 , 1)
Revised: L β − NLL ( β = 0 . 75 ) meanstd ∼ posterior ∼ N (0 , 1) L β − NLL ( β = 1 . 0 ) meanstd ∼ posterior ∼ N (0 , 1)

Original: L NLL ( σ 2 = 1 ) meanstd ∼ posterior ∼ N (0 , 1)
Revised: L NLL ( σ 2 = 1 ) meanstd ∼ posterior ∼ N (0 , 1)

Original: L MM meanstd ∼ posterior ∼ N (0 , 1)
Revised: L MM meanstd ∼ posterior ∼ N (0 , 1)

Original: Student-t meanstd ∼ posterior ∼ N (0 , 1) xVAMP meanstd ∼ posterior ∼ N (0 , 1) xVAMP* meanstd ∼ posterior ∼ N (0 , 1)
Revised: Student-t meanstd ∼ posterior ∼ N (0 , 1) xVAMP meanstd ∼ posterior ∼ N (0 , 1) xVAMP* meanstd ∼ posterior ∼ N (0 , 1)

Original: VBEM meanstd ∼ posterior ∼ N (0 , 1)
Revised: VBEM meanstd ∼ posterior ∼ N (0 , 1)

Original: VBEM* meanstd ∼ posterior ∼ N (0 , 1)
Revised: VBEM* meanstd ∼ posterior ∼ N (0 , 1)

Original: [Figure]
Revised: [Figure]

Original: C D ATASETS AND T RAINING S ETTINGS
Revised: C 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 with heteroscedastic noise We use the synthetic data as introduced in Detlefsen et al.
Revised: Sinusoidal with 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 an MLP with one hidden layer of 50 units and tanh activations (as used in Detlefsen et al.
Revised: The model is an MLP with one hidden layer of 50 units and tanh activations (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.

Original: 3 In particular, we use the training-test protocol from (Hern´andez-Lobato & Adams, 2015; Gal & Ghahramani, 2016), and their data  splits.
Revised: 3 In particular, we use the training-test protocol from (Hern´andez-Lobato & Adams, 2015; Gal & Ghahramani, 2016), and their data splits 4 , except for “carbon”, “energy”, “naval”, “superconductivity”, and “wine-white”, where we generate our own random splits.

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 randomly 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, except for the larger “kin8m”, “power plant”, “protein”, and “naval” datasets where we train for a maximum of 100 000 updates.
Revised: We train for a maximum of 20 000 updates, except for the larger “kin8m”, “power plant”, “protein”, and “naval” datasets where we train for a maximum of 100 000 updates.

Original: 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 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: We use a single-layer relu hidden network with 50 neurons, except for “protein” where we useneurons.
Revised: We use a single-layer relu hidden network with 50 neurons, except for “protein” where we useneurons.

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

Original: Following Stirn & Knowles (2020), for each method, we report the number of datasets for which the method is statistically indistinguishable from the respective best method in Table 1 ( Ties ).
Revised: Following Stirn & Knowles (2020), for each method, we report the number of datasets for which the method is statistically indistinguishable from the respective best method in Table 1 ( Ties ).

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

Original: 05 signiﬁcance level.
Revised: 05 signiﬁcance level.

Original: ObjectSlide This environment consists of an agent whose task is to slide an object to a target location (Seitzer et al., 2021).
Revised: ObjectSlide This environment consists of an agent whose task 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 positions and velocities, and the continuous action space is a one-dimensional movement command.
Revised: The continuous state space consists of 4 dimensions: agent and object positions and velocities, 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 the current state and action.
Revised: The forward prediction task consists of predicting the change in object position in the next state from the 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.

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 a 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 a 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 3-dimensional change in object position from the current state and action.
Revised: As in ObjectSlide, the prediction task is to predict the 3-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 of epochs with a batch size of 256 , and evaluate the model with the best validation log-likelihood on the test set afterwards.

Original: NYUv2 Depth Regression We use the dataset in the variant provided by Lee et al.
Revised: NYUv2 Depth Regression We use the dataset in the variant provided by Lee et al.

Original: Training settings and evaluation protocol were taken from Bhat et al.
Revised: Training settings and evaluation protocol were taken from Bhat et al.

Original: We train for 25 epochs using a batch size of 16 and validate the model every 100 updates.
Revised: We train for 25 epochs using a batch size of 16 and validate the model every 100 updates.

Original: We use the model with the best “REL” metric on the validation set for testing.
Revised: We use the model with the best “REL” metric on the validation set for testing.

Original: D H YPERPARAMETER S ETTINGS AND I MPLEMENTATION D ETAILS
Revised: D 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 parameterize the Gaussian distribution using two linear layers on top of shared features produced by an MLP.
Revised: We parameterize the Gaussian distribution using two linear layers on top of shared features produced by an 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.
Revised: log(1 + exp( x )) activation function.

Original: 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.
Revised: 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: Some baselines use a Student’s t-distribution (Student’s t, xVAMP(*), VBEM(*)) as their predictive distribution.
Revised: Some baselines use a Student’s t-distribution (Student’s t, xVAMP(*), VBEM(*)) as their predictive distribution.

Original: This distribution results from integrating out the unknown variance of a Gaussian with a learned Gamma prior on the inverse variance (Detlefsen et al., 2019; Stirn & Knowles, 2020).
Revised: This distribution results from integrating out the unknown variance of a Gaussian with a learned Gamma prior on the inverse variance (Detlefsen et al., 2019; Stirn & Knowles, 2020).

Original: We parametrize the Gamma distribution in terms of data-dependent alpha and beta parameters, i.e. α ( x ) and β ( x ) , which are computed using linear layers on top of the shared features.
Revised: We parametrize the Gamma distribution in terms of data-dependent alpha and beta parameters, i.e. α ( x ) and β ( x ) , which are computed using linear layers on top of the shared features.

Original: In this case, the
Revised: In this case, the

Original: MLP has three outputs: mean, alpha, and beta.
Revised: MLP has three outputs: mean, alpha, and beta.

Original: Both alpha and beta are constrained to the positive region using the softplus activation.
Revised: Both alpha and beta are constrained to the positive region using the softplus activation.

Original: We add a positive constant of 1 . 001 for alpha and 10 − 8 · 0 .for beta.
Revised: We add a positive constant of 1 . 001 for alpha and 10 − 8 · 0 .for beta.

Original: Alpha is clamped to a maximum value of 1 000 and beta to 10 − 8 · 999 .
Revised: Alpha is clamped to a maximum value of 1 000 and beta to 10 − 8 · 999 .

Original: These values were chosen such that the resulting variance matches the range (10 − 8 , 1000] while ensuring that the degrees-of-freedom parameter ν of the Student’s t is always greater than 2 .
Revised: These values were chosen such that the resulting variance matches the range (10 − 8 , 1000] while ensuring that the degrees-of-freedom parameter ν of the Student’s t is always greater than 2 .

Original: For xVAMP and VBEM, the MLP outputs a fourth term, π ( x ) , representing the logits of a categorical distribution that speciﬁes the mixture weights of the prior.
Revised: For xVAMP and VBEM, the MLP outputs a fourth term, π ( x ) , representing the logits of a categorical distribution that speciﬁes the mixture weights of the prior.

Original: We initialize the prior parameters exactly the same as Stirn & Knowles (2020).
Revised: We initialize the prior parameters exactly the same as Stirn & Knowles (2020).

Original: For these methods, the objective function additionally contains a KL divergence between a Gamma distribution and a mixture-of-Gamma distributions.
Revised: For these methods, the objective function additionally contains a KL divergence between a Gamma distribution and a mixture-of-Gamma distributions.

Original: Following
Revised: Following

Original: Stirn & Knowles (2020), we approximate this KL divergence using 20 Monte-Carlo samples.
Revised: Stirn & Knowles (2020), we approximate this KL divergence using 20 Monte-Carlo samples.

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

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

Original: 1 results from a network of two hidden layers with 128 neurons per layer and tanh activations, optimized with a learning rate of 5 · 10 − 4 and a batch size of 100 .
Revised: 1 results from a network of two hidden layers with 128 neurons per layer and tanh activations, optimized with a learning rate of 5 · 10 − 4 and a batch size of 100 .

Original: For the experiment in Sec.
Revised: For the experiment in Sec.

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

Original: [Table]
Revised: [Table]

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

Original: For each tested loss function, we performed a grid search on the ObjectSlide and Fetch-PickAndPlace datasets.
Revised: For each tested loss function, we performed a grid search on the ObjectSlide and Fetch-PickAndPlace datasets.

Original: We report the parameters we scanned over in Table S6.
Revised: We report the parameters we scanned over in Table S6.

Original: Table S7 reports the model conﬁgurations with the best validation log-likelihood on the grid search.
Revised: Table S7 reports the model conﬁgurations with the best validation log-likelihood on the grid search.

Original: For the results in Table 2, we retrained the best model conﬁguration with ﬁve different random seeds and evaluated them on the hold-out test set. 
Revised: For the results in Table 2, we retrained the best model conﬁguration with ﬁve different random seeds and evaluated them on the hold-out test set. [Table] [Table]

Original: D. V ARIATIONAL A UTOENCODERS
Revised: D. V ARIATIONAL A UTOENCODERS

Original: We largely follow Stirn & Knowles (2020) for their training settings for the VAE experiment.
Revised: We largely follow Stirn & Knowles (2020) for their training settings for the VAE experiment.

Original: In particular, we use an encoder with three layers of 512 , 256 , 128 neurons and a decoder with three[Table] [Table] layers of 128 , 256 , 512 neurons, all with relu activations.
Revised: In particular, we use an encoder with three layers of 512 , 256 , 128 neurons and a decoder with three layers of 128 , 256 , 512 neurons, all with relu activations.

Original: The latent space is 10 -dimensional for
Revised: The latent space is 10 -dimensional for

Original: MNIST and 25 -dimensional for FashionMNIST.
Revised: MNIST and 25 -dimensional for FashionMNIST.

Original: We train the VAEs for a maximum of 1 000 epochs, using Adam with a learning rate of 0 .
Revised: We train the VAEs for a maximum of 1 000 epochs, using Adam with a learning rate of 0 .

Original: 0003 and a batch size of 256 .
Revised: 0003 and a batch size of 256 .

Original: Early-stopping with a patience of 50 epochs is performed on the log-likelihood of the validation set.
Revised: Early-stopping with a patience of 50 epochs is performed on the log-likelihood of the validation set.

Original: The validation set consists of 20% of the MNIST/FashionMNIST training set.
Revised: The validation set consists of 20% of the MNIST/FashionMNIST training set.

Original: D. D EPTH R EGRESSION
Revised: D. D EPTH R EGRESSION

Original: We use the ofﬁcial implementation of AdaBins (Bhat et al., 2021), 6 thereby reproducing their exact training settings and evaluation protocol.
Revised: We use the ofﬁcial implementation of AdaBins (Bhat et al., 2021), 6 thereby reproducing their exact training settings and evaluation protocol.

Original: We remove the AdaBins/mini-ViT Transformer from the model.
Revised: We remove the AdaBins/mini-ViT Transformer from the model.

Original: Instead, the feature map output by the U-Net is reduced to two channels using a 1 ×convolution, where we use the ﬁrst channel as the mean predictor and the second channel as the variance predictor.
Revised: Instead, the feature map output by the U-Net is reduced to two channels using a 1 ×convolution, where we use the ﬁrst channel as the mean predictor and the second channel as the variance predictor.

Original: In this setting, both mean and variance are constrained to positive numbers by a softplus activation.
Revised: In this setting, both mean and variance are constrained to positive numbers by a softplus activation.

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

Original: D. I MPLEMENTATION OF BETA -NLL IN P YTORCH
Revised: D. I MPLEMENTATION OF BETA -NLL IN P YTORCH