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

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

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

Original: In deep learning, a common approach to this end is to train a neural network to estimate the parameters of a heteroscedastic Gaussian distribution by maximizing the logarithm of the likelihood function under the observed data.
Revised: In deep learning, a common approach to this end is to train a neural network to estimate the parameters of a heteroscedastic Gaussian distribution by maximizing the logarithm of the likelihood function under the observed data.

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

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

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

Original: Third, we present an alternative formulation in which each data point’s contribution to the loss is weighted by the β -exponentiated variance estimate.
Revised: Third, we present an alternative formulation, 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 standard benchmarks from the literature and show that it achieves considerable improvements and performs more robustly with respect to hyperparameters, both in predictive RMSE and log-likelihood criteria.
Revised: Fourth, we evaluate this approach on a range of 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 beusefully 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 is noise that stays constant across different inputs, whereas heteroscedastic uncertainty is one that varies depending on the inputs to the model.
Revised: Homoscedastic uncertainty 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  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: Many modern applications of machine learning rely on neural networks and deep learning tools to achieve state-of-the-art performance.

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

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

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

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

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.

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: criterion using stochastic gradient descent. This simple procedure, which is the de-facto standard (Nix & Weigend, 1994; Lakshminarayanan et al., 2017; Kendall & Gal, 2017; Chua et al., 2018; Kloss et al., 2021), is known to be subject to overconﬁdent variance estimates.
Revised: T a r

Original: Whereas strategies have been proposed to alleviate this speciﬁc issue (Detlefsen et al., 2019; Hu et al., 2020; Stirn & Knowles, 2020), we argue that an equally important issue is that this procedure can additionally lead to subpar mean ﬁts.
Revised: [Figure] issue is that this procedure can additionally lead to subpar mean ﬁts .

Original: In this work, we analyse and propose a simple modiﬁcation to mitigate this issue.
Revised: In this work, we 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  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 set of standard benchmarks while exhibiting more robustness to network architecture and learning rate conﬁgurations.
Revised: Finally, we empirically show that our modiﬁed loss formulation largely mitigates the issue of poor ﬁts, achieving considerable improvements on a 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[Figure]
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: An initial inhomogeneous feature space granularity (explained in text) early on in different ﬁtting quality.

Original: The implicit weighting of the squared error in NLL can be seen as a biased data

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

Original: -dependent quantiﬁes  heteroscedastic 
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., 
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: 2019).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:  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: [Figure] The optimal parameters θ ∗ NLL can then be found using maximum likelihood estimation (MLE) by minimizing the negative log-likelihood (NLL) criterion L NLL under the distribution P ( X, Y ) :

Original: [Equation]
Revised: [Equation]

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

Original: L MSE :
Revised: L MSE :

Original: [Equation]
Revised: [Equation]

Original: In practice, Eq. 1 and Eq. 2 are optimized using stochastic gradient descent  with batches of samples drawn from P ( X, Y ) .
Revised: In practice [Equation] 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 with respect to ˆ µ ( X ) , ˆ σ 2 ( X ) are given by
Revised: The gradients of L NLL w.r.t. (with respect to) ˆ µ ( X ) , ˆ σ 2 ( X ) are given by

Original: [Equation]
Revised: [Equation]

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 without problems in  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 a 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 is the model not converging?
Revised: So, why does the model not converge to the optimal solution?

Original: We identify two main culprits for the non-convergence 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 results in a high effective data variance.
Revised: If the network cannot ﬁt a certain region yet because its feature space (spanned by the last hidden layer) is too coarse, it would then perceive a high effective data variance.

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

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

Original: [Figure]

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

Original: It is instructive to see how the model evolves over training.
Revised: It is instructive to see how the model evolves during training as shown in Fig.

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

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

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

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

Original: Escaping this situation thus requires some form of symmetry breaking.
Revised: Symmetry breaking is required for further progress.

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

Original: In addition, we hypothesize that an important form ofsymmetry breaking is how non-linear the feature space of the network is locally, in order to create 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, we compute how much the Jacobian J f of the features f ( x ) with respect to the input varies in an L2-ball with radius r around a point x , denoted as the Jacobian variance:V ( x [Equation]
Revised: To approximate it for a dataset D , we compute how much the Jacobian J f of the features f ( x ) w.r.t. the input varies in an L 2 -ball with radius r around a point x , denoted as the Jacobian variance:

Original: =
Revised: V

Original: [Equation]
Revised: [Equation]

Original: | [Equation] [Figure]
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 gets more ﬁne-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 that has a low Jacobian variance gets stuck in this conﬁguration, see Fig. 
Revised: The region with low Jacobian variance remains stuck in this conﬁguration (see Fig. 1).

Original: This provides evidence that the non-linearity of the feature space is important forthe success or failure of learning.
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,it does not answer the question of why  gradient descenton the NLL loss would not break out of this situation
Revised: However, why does gradient descent not break out of this situation?

Original: [Figure]
Revised: [Figure]

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

Original: [Figure]
Revised: [Figure]

Original: 3.  W EIGHTINGBY V ARIANCE 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 with respect to the mean scales the error ˆ µ ( X ) − Y by [Equation]
Revised: Recall that the gradient ∇ ˆ µ L NLL of the NLL w.r.t. 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(ˆ µ ( X [Equation] Y
Revised: If the variance is well-calibrated, the gradient becomes ˆ µ ( X ) − Y ˆ σ 2 ( X ) ≈ ˆ µ ( X ) − Y

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

Original: That is, data points with already low error will get their contribution in the batch gradient scaled up relatively to high error data points –a “rich get richer” self-amplifying scenario.
Revised: 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,the NLL exhibits the opposite of the MSE loss’s well-known behavior to focus 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 small compared to the error on the badly-ﬁt regions, or if the ratio of well-ﬁt to badly-ﬁt points is high, this can completely prevent learning progress  on the badly-ﬁt regions.
Revised: If the true variance σ 2 on the well-ﬁt regions is 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: A.2.
Revised: B.3.

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

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

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

Original: Sometimes “weighting-by-variance” is seen as a feature of the Gaussian NLL (Kendall & Gal, 2017) that 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 the inherent unpredictability (noise), but it is undesirable if it causes premature convergence and ignorance of hard-to-ﬁt regions, as shown above.
Revised: This can be desirable if the predicted variance corresponds to 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 (Sec. 4), we allow to control  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: To mitigate these problems with NLL training, we develop a solution in the following.
Revised: In this section, we develop a solution method to mitigate these issues with NLL training.

Original: We start with a method that uses standard squared error losses to estimate the moments of the target distribution, which we term “moment matching” (MM; Sec. 4.1).

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

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

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

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

Original: 4. M OMENT M ATCHING

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

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

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

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

Original: [Equation]

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

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

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

Original: The variance is deﬁned as

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

Original: [Figure]

Original: To use the same physical unit as L µ MM , we reformulate it in terms of the standard deviation as: [Equation] Thus, the moment matching loss is simply given by the sum of both losses: L MM = L µ MM + L σ MM .
Revised: Our approach, which we term β − NLL , allows choosing an arbitrary loss-interpolation between NLL and

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

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

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

Original: As discussed in Sec. 3.2, the NLL loss gives high weight to data points with low (predicted) variance and low weight to those with high variance.

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

Original: However, due to the use of function approximators, there is also the case where data points cannot be well predicted, resulting in high predicted variance, although the ground truth is corrupted by little noise.

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

Original: Figure 6 shows the distribution of the residuals for a dynamics prediction dataset containing easy and hard to model areas. We observe a drastic difference between the NLL loss and the MSE loss (or equivalently solely optimizing L µ MM ). How important are the data points with high uncertainty? Are these outliers or stem from truly noisy regions?
Revised: MSE while keeping calibrated uncertainty estimates.

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

Original: Or are these simply difﬁcult samples?

Original: In our dynamics prediction example, we analysed the data points with high uncertainty and found that they were actually the most important ones to get correct (because they captured non-trivial interactions in the physical world).

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

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

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

Original: L NLL . Themost important problem  to solve 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  with its self-amplifying character as the main reason for this.
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 (Sec. 3.2).
Revised: Effectively, NLL weights the mean-squared-error per data point with 1 σ 2 , which can be interpreted as sampling data points with P ( x ) ∝ 1 σ 2 .

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: [Figure]
Revised: where (cid:98)·(cid:99) denotes the stop gradient operation. By stopping the gradient, the variance-weighting term acts as an adaptive, input-dependent learning rate. In this way, the gradients of L β − NLL are:

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

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

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

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

Original: In Sec. 5, we show that β
Revised: In 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: As a remark, the new loss L β − NLL is not meant for performance evaluation – it is designed to result in meaningful gradients.
Revised: A Pytorch implementation of the loss function is provided in Sec. D.5. Note that the new loss L β − NLL is not meant for performance evaluation, rather it is designed to result in meaningful gradients.

Original: We found that its absolute value can be counter-intuitive due to the prefactor.
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 and optionally the RMSE for  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. [Figure]

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

Revised: Even though L MSE , L NLL , and L β − NLL all have the same optima w.r.t. the mean (and also the variance in the case of L NLL and L β − NLL ), optimizing them leads to very different solutions. In particular, because these losses weight data points differently, they assign the capacity of the function approximator differently .

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

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

Revised: β − NLL interpolates between the two.

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

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

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

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

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

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

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

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

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

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?

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

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

Original: We investigate the proposed changes in the  loss function for training neural networks to predict a
Revised: Rather, the modeler should choose which behavior is desired. Our β − NLL loss makes this choice available through the β parameter.  E XPERIMENTS In our experiments, we ask the following questions and draw the following conclusions:

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

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

Original: Gaussian distribution on synthetic examples and real-world datasets featuring the UCI and challenging dynamics model datasets.
Revised: Sec. 5.2: Does β − NLL improve over NLL in practical settings? How sensitive to hyperparameters is β − NLL ? 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: As motivated above, the standard NLL loss often leads to suboptimal ﬁts, e.g. where the mean ﬁt is far from ideal.
Revised: Yes , β − NLL generally performs better than NLL and is considerably easier to tune. Sec. 5.3: How does β − NLL compare to other loss functions for distributional regression?

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

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*

Revised: ) (Stirn & Knowles, 2020).

Revised: It depends.

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

Revised: [Figure]

Revised: [Figure]

Revised: estimates.

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

Original: To analyze this quantitatively, we report the performance for many hyperparameter settings, such as learning rates and network architectures, and consider both the quality in predictions with respect to the root mean squared error (RMSE) and the negative log-likelihood (NLL).
Revised: We refer the reader to Sec. C and Sec. D for a description of datasets and training settings. We evaluate the quality of predictions quantitatively in terms of the root mean squared error (RMSE) and the negative log-likelihood (NLL).

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

Original: Sinusoidal without heteroscedastic noise We ﬁrst consider our illustrative example of Fig.
Revised: Sinusoidal without heteroscedastic noise We ﬁrst perform an extended investigation of our illustrative example from Fig.

Original: 1 – asimple sine curve with a small additive noise: y = sin( x ) + ξ , with ξ being Gaussian noise with standard deviation σ = 0 .
Revised: 1 – a sine curve with a small additive noise: y = 0 . 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  ( β = 0 ) the network does 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 problems.
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 for completeness.
Revised: S3.

Original: Sinusoidal with heteroscedastic noise A commonly used illustrative example introduced in 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: (2019) is a sine curve with increasing amplitude and noise: y = x sin( x ) + xξ 1 + ξ 2 , with ξ 1 and ξ 2 being Gaussian noise with standard deviation σ = 0 .
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: The functional form is much easier, so ﬁtting the mean is achieved with all losses.
Revised: Figure 8 (a-e) displays the predictions of the best models (w.r.t.

Original: Here we sanity-check that the newly introduced loss is [Figure] [Figure] 10 s t d gt std data std
Revised: NLL validation loss) and (f) compares the predicted uncertainties over 10 independent trials.

Original: [Figure] [Table] still delivering good uncertainty estimates. Figure 8 (a-e) displays the predictions of the best models (w.r.t.
Revised: Fitting the mean is achieved with all losses.

Original:  NLL validation loss) and (f) compares the predicted uncertainties over 10 independent trials.
Revised: On the training range, β − NLL with β > 0 learns virtually the same uncertainties as the NLL loss.

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

Original:  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: A.3 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-likelihood on par with or better than the NLL loss while clearly improving the predictive accuracy on most datasets.
Revised: The results are encouraging: β − NLL achieves predictive log-likelihoods on par with or better than the NLL loss while clearly improving the predictive accuracy on most datasets.

Original: 5. D YNAMICS MODELS
Revised: [Table] [Figure]

Original:  As a major application of uncertainty estimation lies in model-based reinforcement learning (RL), we test the different loss functions on two dynamics predictions tasks of varying difﬁculty, ObjectSlide and Fetch-PickAndPlace.
Revised: 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, 
Revised: Whereas ObjectSlide (Seitzer et al., 2021) is a simple 1D-environment, Fetch-PickAndPlace

Original: Fetch-PickAndPlace (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 onstate-action trajectories collected by RL agents.
Revised: The models are trained on trajectories collected by RL agents.

Original: See Sec. B for more details. For both datasets and all loss functions, weﬁrst perform a grid search over different hyperparameter conﬁgurations (see Sec. C.2).
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: The results also serve for a sensitivity analysis presented in Fig. 9: how much does one have to tune hyperparameters for each loss function to work well in practice?
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: We then take the best performing conﬁgurations and evaluate them 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: The sensitivity analysis reveals thatboth the NLL- and theMM-loss are vulnerable to hyperparameter settings, whereas the β − NLL lossachieves good results over awide range of settings.
Revised: One can see that the NLL loss results in poor predictive performance and also exhibits quite a high variance across random seeds. Our method yields high accuracy and log-likelihood ﬁts for a range of β values, with β = 0 .

Original: When focusing on the best found model conﬁgurations for each loss in Table 2, one can see that the NLL loss results in poor predictive performance and also exhibits quite a high variance across random seeds; the MMloss results in poor log-likelihood ﬁts.
Revised: 5 generally achieving the best trade-off. Generative modeling and depth-map prediction For generative modeling, we train variational autoencoders (Kingma & Welling, 2014) with probabilistic decoders on MNIST and Fashion-MNIST.

Original: Our method performs well on both accuracy and log-likelihood ﬁt, for a range of β -values, where β
Revised: For the task of depth regression, we modify a state-of-the-art method (AdaBins; Bhat et al.

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

Original: 5 generally achieves the best trade-off between both.
Revised: Table 3 presents selected results for both tasks, yielding similar trends as before.

Original: [Figure]
Revised: We refer to Sec. B.5 and Sec. B.6 for more details, including qualitative results. 5. C OMPARISON TO O THER L OSS F UNCTIONS

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

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

Revised: First,

Revised: [Table]

Revised: can be found in Table S3 and Table S

Original: [Table]
Revised: [Table]

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

Original: D ISCUSSION
Revised: Sec. A). While this loss results in high accuracy, it is unstable to train and exhibits poor likelihoods. Second, we test several loss functions based on a Student’s t-distribution, including xVAMP and

Revised: VBEM (Stirn & Knowles, 2020).

Original: We highlight a deep problem frequently occurring whenoptimizing probabilistic neural networks using the common NLL loss: training gets stuck in suboptimal function ﬁts.
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. In terms of predictive accuracy, β − NLL outperforms the Student’s t-based approaches. 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 . This is likely because these methods can mitigate overﬁtting by placing a prior on the variance.

Original: Due to  the predictive variance-based weighting of mean prediction errors in the NLL loss, these bad ﬁts often stay unnoticed, unless the MSE is evaluated.
Revised: 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. In contrast, β − NLL is simple to implement (see Sec. D.5) and introduces no additional computational costs.

Original:  With our analysis on feature space non-linearity and induced data-set weighting (interpreted as a biased dataset sampling), we reveal the underlying reason: initially badly ﬁt regions (due to low feature non-linearity) get increasingly less weight and result in premature convergence.
Revised: We highlight a problem frequently occurring when optimizing probabilistic neural networks using the common NLL loss: training gets stuck in suboptimal function ﬁts. 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  NLLloss and MSE loss  while providing wellbehaved uncertainty estimates (even for the MSE case).
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 unimportant should high-noise regions or difﬁcult-to-predict data points be in the ﬁtting process.
Revised: The hyperparameter β gives practitioners the choice to control the self-regularization strength of NLL: how 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.

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 will improve the usability and performance of predictive uncertainty  deep networks and increase their applicability.
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. A CKNOWLEDGMENTS

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

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

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

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

Original: C. We make full code and data available under https://sites.google.com/view/pitfalls-uncertainty . 
Revised: D. We make full code and data available under https://github.com/martius-lab/pitfalls-of-uncertainty . A PPENDIX A T HE M OMENT M ATCHING L OSS

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

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: Both will be made publicly available after the review phase.
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. To fully describe a Gaussian distribution, only the ﬁrst two moments µ, σ 2 need to be estimated. So, why not directly deﬁne losses based on the moment estimators? 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: [Equation]

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

Revised: Thus, learning a mean ﬁt ˆ µ ( x ) can follow the standard (non-distributional) regression procedures.

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

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

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

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

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

Revised: [Equation]

Revised: Thus, the moment matching loss can be expressed simply by the sum of the two losses: L MM =

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

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

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

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.

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: A PPENDIX A A DDITIONAL R ESULTS A. S YNTHETIC D ATASET
Revised: B A DDITIONAL R ESULTS B. F URTHER C ONSEQUENCES OF I NVERSE -V ARIANCE W EIGHTING

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.

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

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.

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.

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

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.

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.

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. In order to test the dependence of our reported issue on the optimizer

Original: In Fig. 
Revised: [Figure] we repeat the experiment from Fig. but use different optimizers than Adam with β 1 = 0 . 9 , β 2 = 0 .

Revised: The results are shown in Fig. S2.

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

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

Revised: In Fig.

Original: S
Revised: S

Original: [Figure]
Revised: [Figure]

Original: we provide the results for the NLL metric on the sinusoidal dataset, complimentary to the results in Fig.
Revised: we provide the results for the 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]

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

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

Original: Also for β = 1 , which corresponds to MSE for ﬁtting the mean good uncertainty predictions are obtained with our β − NLL as testiﬁed by the low scores.

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

Revised: In Fig. S

Original: In this section we include results for predictive log-likelihood (Table S1) and RMSE (Table S2) for all UCI datasets we evaluated on.
Revised: [Figure] we show how the analysis from Sec. 3.2 transfers to a real world dataset, namely FetchPickAndPlace. 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 mismatch between the two distributions demonstrates that optimizing the

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

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

Original: [Figure]
Revised: [Figure]

Revised: B. UCI D ATASETS

Revised: In this section

Revised: [Figure]

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

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 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: For RMSE, L MSE unsurprisingly performs best.

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

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

Revised: D.5)

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

Revised: Monte-Carlo sampling at each training step.

Original: B D ATASETS AND T RAINING S ETTINGS
Revised: B. G ENERATIVE M ODELING WITH V ARIATIONAL A UTOENCODERS

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

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.

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

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

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

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

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

Revised: We evaluate on MNIST and FashionMNIST.

Revised: Table S3 presents quantitative results.

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

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

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

Revised: Figure S

Revised: [Table]

Revised: [Table]

Revised: shows qualitative examples.

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

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.

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

Revised: [Table]

Revised: B. D EPTH R EGRESSION

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

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.

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.

Revised: Table S4 presents quantitative results.

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

Revised: β − NLL with β = 0 .

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

Revised: Figure S6 shows qualitative examples.

Revised: Compared to

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

Revised: refer

Revised: [Equation]

Revised: [Table]

Revised: MNIST

Revised: Fashion-MNIST dataset

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

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

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

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

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

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

Revised: L MM 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)

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

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: VBEM* meanstd ∼ posterior ∼ N (0 , 1) [Figure] C D ATASETS AND T RAINING S ETTINGS Sinusoidal without heteroscedastic noise This dataset is created by taking 1 000 uniformly spaced points on the interval [0 , 12] as inputs x and applying the function y ( x ) = 0 .

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

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

Original: Sinusoidal  heteroscedastic noise We use the synthetic data as introduced in Detlefsen et al.
Revised: Sinusoidal 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 a MLP with one hidden layer with 50 units and tanh activation function (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 3 .
Revised: UCI Datasets We use the UCI datasets suite commonly used to benchmark uncertainty estimation, stemming from the UCI Machine Learning Repository.

Original:  In particular, we use the training-test protocol from (Hern´andez-Lobato & Adams, 2015; Gal & Ghahramani, 2016), and their data splits 4 .
Revised: 3 In particular, we use the training-test protocol from (Hern´andez-Lobato & Adams, 2015; Gal & Ghahramani, 2016), and their data 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 divide the training set into 80% training data and 20% validation data and search for an optimal learning rate from the set { 10 − 4 , 3 · 10 − 4 , 7 · 10 − 4 , 10 − 3 , 3 · 10 − 3 , 7 · 10 − 3 } by monitoring log-likelihood on the validation set.

Original: We train for a maximum of 20 000 updates, besides on the larger “kin8m”, “power plant”, “protein” and “naval” datasets, where we train for a maximum of 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, aside from “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: [Table] Following Stirn & Knowles (2020), for each method, we report the number of datasets for which the method  statistically indistinguishable from the respective best method in Table 1 (”Ties”).
Revised: 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  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 taskit 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 position and velocity, and the continuous action space is a one-dimensional movement command.
Revised: The continuous state space consists of 4 dimensions: agent and object 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  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.[Table]
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  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 3D-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 ofepochs with a batch size of 256 , and evaluate the model with the best validation log-likelihood on the test set afterwards. NYUv2 Depth Regression We use the dataset in the variant provided by Lee et al.

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

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

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

Original: C 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 parametrize the Gaussian distribution using two linear layers on top of shared features produced by a 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.

Revised: Some baselines use a Student’s t-distribution (Student’s t, xVAMP(*), VBEM(*)) as their predictive distribution.

Original: C. S INUSOIDAL R EGRESSION P ROBLEM
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). 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. In this case, the MLP has three outputs: mean, alpha, and beta.

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

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

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

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 .

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.

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

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

Revised: Following

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

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 with two hidden layers and 128 neurons each, tanh activation, and  optimized with  learning rate  5 · 10 − 4 and a batch size of 100 .
Revised: 1 results from a network 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 
Revised: For the experiment in Sec.

Original: Sec. 5.1, we scan over learning rates and architectures with different hidden layers and units per layer, as detailed in Table S3.
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: For each tested loss function, we performed a grid search for the ObjectSlide and Fetch-PickAndPlace datasets, We report the parameters  in Table S4.
Revised: For each tested loss function, we performed a grid search on the ObjectSlide and Fetch-PickAndPlace datasets. We report the parameters we scanned over in Table S6.

Original: For the results in Table 2, we retrained the model conﬁgurations with the best validation log-likelihood on the grid search, and evaluated them on the hold-out test set.
Revised: Table S7 reports the model conﬁgurations with the best validation log-likelihood on the grid search. 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: D. V ARIATIONAL A UTOENCODERS

Original: The best found conﬁguration on ObjectSlide coincided to 3 hidden layers with  128 neurons, relu activation and learning rate 0 .
Revised: We largely follow Stirn & Knowles (2020) for their training settings for the VAE experiment. In particular, we use an encoder with three layers of 512 , 256 , 128 neurons and a decoder with three [Table] [Table]

Original: The best found conﬁguration on Fetch-PickAndPlace coincided to 4 hidden layers with 128  neurons,  relu activation and  learning rates 0 .
Revised: layers of 128 , 256 , 512 neurons, all with relu activations. The latent space is 10 -dimensional for MNIST and 25 -dimensional for FashionMNIST. We train the VAEs for a maximum of 1 000 epochs, using Adam with a learning rate of 0 .

Original: 0003 for L NLL , L β − NLL
Revised: 0003 and a batch size of 256 . Early-stopping with a patience of 50 epochs is performed on the log-likelihood of the validation set. The validation set consists of 20% of the MNIST/FashionMNIST training set. D. D EPTH R EGRESSION

Original: [Table]
Revised: We use the ofﬁcial implementation of AdaBins (Bhat et al., 2021), 6 thereby reproducing their exact training settings and evaluation protocol. We remove the AdaBins/mini-ViT Transformer from the model. Instead, the feature map output by the U-Net is reduced to two channels using a 1 ×convolution, where we use the ﬁrst channel as the mean predictor and the second channel as the variance predictor. In this setting, both mean and variance are constrained to positive numbers by a softplus activation.

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 .

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