On the Pitfalls of Heteroscedastic Uncertainty Estimation with Probabilistic Neural Networks

1 INTRODUCTION . Endowing models with the ability to capture uncertainty is of crucial importance in machine learning . Uncertainty can be usefully categorized into two main types : epistemic uncertainty and aleatoric uncertainty ( Kiureghian & Ditlevsen , 2009 ) . Epistemic uncertainty accounts for subjective uncertainty in the model , one that is reducible given sufficient data . By contrast , aleatoric uncertainty captures the stochasticity inherent in the observations and can itself be subdivided into homoscedastic and heteroscedastic uncertainty . Homoscedastic uncertainty is noise that stays constant across different inputs , whereas heteroscedastic uncertainty is one that varies depending on the inputs to the model . There are well-established benefits for modeling each type of uncertainty . For instance , capturing epistemic uncertainty enables effective budgeted data collection in active learning ( Gal et al. , 2017 ) , allows for efficient exploration in reinforcement learning ( Osband et al. , 2016 ) , and is indispensable in cost-sensitive decision making ( Amodei et al. , 2016 ) . On the other hand , quantifying aleatoric uncertainty enables learning dynamics models of stochastic processes ( e.g . for model-based or offline reinforcement learning ) ( Chua et al. , 2018 ; Yu et al. , 2020 ) , improves performance in semantic segmentation , depth regression and object detection ( Kendall & Gal , 2017 ; Harakeh & Waslander , 2021 ) , and allows for risk-sensitive decision making ( Dabney et al. , 2018 ; Vlastelica et al. , 2021 ) . Many modern applications of machine learning rely on neural networks and deep learning tools to achieve state-of-the-art performance . However , most neural network models are not readily equipped with a capacity for uncertainty estimation . To address this shortcoming , a multitude of approaches have been proposed in the past few years . Nevertheless , the majority of work in this domain has focused on epistemic uncertainty quantification ( Blundell et al. , 2015 ; Gal & Ghahramani , 2016 ) or uncertainty estimation for classification ( Hendrycks & Gimpel , 2017 ; Guo et al. , 2017 ; Ovadia et al. , 2019 ; Mukhoti et al. , 2021 ) . We examine a common approach for quantifying aleatoric uncertainty in neural network regression . By assuming that the regression targets follow a particular distribution , we can use a neural network to predict the parameters of that distribution , typically the input-dependent mean and variance when assuming a heteroscedastic Gaussian distribution . Then , the parameters of the network can be learned using maximum likelihood estimation ( MLE ) , i.e . by minimizing the negative log-likelihood ( NLL ) criterion using stochastic gradient descent . This simple procedure , which is the de-facto standard ( Nix & Weigend , 1994 ; Lakshminarayanan et al. , 2017 ; Kendall & Gal , 2017 ; Chua et al. , 2018 ; Kloss et al. , 2021 ) , is known to be subject to overconfident variance estimates . Whereas strategies have been proposed to alleviate this specific issue ( Detlefsen et al. , 2019 ; Hu et al. , 2020 ; Stirn & Knowles , 2020 ) , we argue that an equally important issue is that this procedure can additionally lead to subpar mean fits . In this work , we analyse and propose a simple modification to mitigate this issue . Summary of contributions We demonstrate a pitfall of optimizing the NLL loss for neural network regression , one that hinders the training of accurate mean predictors ( see Fig . 1 for an illustrative example ) . The primary culprit is the high dependence of the gradients on the predictive variance . While such dependence is generally known to be responsible for instabilities in joint optimization of mean and variance estimators ( Takahashi et al. , 2018 ; Stirn & Knowles , 2020 ) , we identify a fresh perspective on how this dependence can further be problematic . Namely , we hypothesize that the issue arises due to the NLL loss scaling down the gradient of poorly-predicted data points relative to the well-predicted ones , leading to effectively undersampling the poorly-predicted data points . We then introduce an alternative loss formulation that counteracts this by weighting the contribution of each data point to the overall loss by its β-exponentiated variance estimate , where β controls the extent of dependency of gradients on predictive variance . This formulation subsumes the standard NLL loss for β = 0 and allows to lessen the dependency of gradients on the variance estimates for 0 < β ≤ 1 . Interestingly , using β = 1 completely removes such dependency for training the mean estimator , yielding the standard mean squared error ( MSE ) loss – but with the additional capacity of uncertainty estimation . Finally , we empirically show that our modified loss formulation largely mitigates the issue of poor fits , achieving considerable improvements on a set of standard benchmarks while exhibiting more robustness to network architecture and learning rate configurations . 2 PRELIMINARIES . Let X , Y be two random variables describing the input and target , following the joint distribution P ( X , Y ) . We assume that Y is conditionally independent givenX and that it follows some probability distribution P ( Y | X ) . In the following , we use the common assumption that Y is normally distributed given X ; i.e . P ( Y | X ) = N ( µ ( X ) , σ2 ( X ) ) , where µ : RM 7→ R and σ2 : RM 7→ R+ are respectively the true input-dependent mean and variance functions.1 Equivalently , we can write Y = µ ( X ) + ( X ) , with ( X ) ∼ N ( 0 , σ2 ( X ) ) ; i.e . Y is generated fromX by µ ( X ) plus a zero-mean Gaussian noise with variance σ2 ( X ) . This input-dependent variance quantifies the heteroscedastic uncertainty , or input-dependent aleatoric uncertainty . 1For notational convenience , we focus on univariate regression but point out that the work extends to the multivariate case as well . To learn estimates µ̂ ( X ) , σ̂2 ( X ) of the true mean and variance functions , it is common to use a neural network fθ parameterized by θ . Here , µ̂ ( X ) and σ̂2 ( X ) can be outputs of the final layer ( Nix & Weigend , 1994 ) or use two completely separate networks ( Detlefsen et al. , 2019 ) .The variance output is hereby constrained to the positive region using a suitable activation function , e.g . softplus . The optimal parameters θ∗NLL can then be found using maximum likelihood estimation ( MLE ) by minimizing the negative log-likelihood ( NLL ) criterion LNLL under the distribution P ( X , Y ) : θ∗NLL = argmin θ LNLL ( θ ; D ) = argmin θ E X , Y [ 1 2 log σ̂2 ( X ) + ( Y − µ̂ ( X ) ) 2 2σ̂2 ( X ) + const ] . ( 1 ) In contrast , standard regression minimizes the mean squared error ( MSE ) LMSE : θ∗MSE = argmin θ LMSE ( θ ; D ) = argmin θ E X , Y [ ( Y − µ̂ ( X ) ) 2 2 ] . ( 2 ) In practice , Eq . 1 and Eq . 2 are optimized using stochastic gradient descent with batches of samples drawn from P ( X , Y ) . The gradients of LNLL with respect to µ̂ ( X ) , σ̂2 ( X ) are given by ∇µ̂LNLL ( θ ) = E X , Y [ µ̂ ( X ) − Y σ̂2 ( X ) ] , ∇σ̂2LNLL ( θ ) = E X , Y [ σ̂2 ( X ) − ( Y − µ̂ ( X ) ) 2 2 ( σ̂2 ( X ) ) 2 ] . ( 3 , 4 ) 3 ANALYSIS . We now return to the example of trying to fit a sinusoidal function from Sec . 1 . Recall from Fig . 1 that using the Gaussian NLL as the objective resulted in a suboptimal fit . In contrast , using MSE as the objective , the model converged without problems in reasonable time . We now analyze the reasons behind this surprising result . From Eq . 3 , we see that the true mean µ ( X ) is a minimizer of the NLL loss . It thus becomes clear that a ) the solution found in Fig . 1 is not the optimal one , and b ) the NLL objective should , in principle , drive µ̂ ( X ) to the optimal solution µ ( X ) . So why is the model not converging ? We identify two main culprits for the non-convergence of the Gaussian NLL objective : 1 . Initial flatness of the feature space can create an undercomplex but locally stable mean fit . This fit results from local symmetries and requires a form of symmetry breaking to escape . 2 . The NLL loss scales the gradient of badly-predicted points down relative to well-predicted points , effectively undersampling those points . This effect worsens as training progresses . These culprits and their effect on training are illustrated in Fig . 2 ( left ) . If the network can not fit a certain region yet , because its feature space ( spanned by the last hidden layer ) is too coarse , it results in a high effective data variance . This leads to down-weighting the data from these regions , fueling a vicious circle of self-amplifying the increasingly imbalanced weighting . In the following , we analyse these effects and their reasons in more detail . 3.1 SYMMETRY AND FEATURE NON-LINEARITY . It is instructive to see how the model evolves over training as shown in Fig . 3 . The network first learns essentially the best linear fit while adapting the variance to match the residuals . The situation is locally stable . That is , due to the symmetries of errors below and above the mean fit , there is no incentive to change the situation . Symmetry breaking is required for further progress . One form of symmetry breaking comes from the stochasticity of mini-batch sampling in SGD , or natural asymmetries contained in the dataset , e.g . outlier points . In addition , we hypothesize that the local non-linearity of the feature space plays an important role in creating the necessary non-linear fit . Thus , let us consider the non-linearity of the feature space . This quantity is not easy to capture . To approximate it , we compute how much the Jacobian Jf of the features f ( x ) with respect to the input varies in an L2-ball with radius r around a point x , denoted as the Jacobian variance:2 V ( x ) = 1 |B| ∑ x′∈S ( Jf ( x ′ ) − 1|S| ∑ x′′∈B Jf ( x ′′ ) ) 2 , B = { x′ : ‖x− x′‖2 ≤ r } . ( 5 ) Figure 4 visualizes the Jacobian variance over the input space as a function of the training progress . Although initially relatively flat , it becomes more fine-granular in parts of the input space – the parts which are later well fit . The region with low Jacobian variance remains stuck in this configuration , see Fig . 1 . This provides evidence that the non-linearity of the feature space is important for the success or failure of learning . However , why does gradient descent not break out of this situation ? 0 4 8 12 Input X 0 1 2 3 U pd at e St ep s × 10 6 10−4 100 Ja co bi an V ar ia nc e Figure 4 : Jacobian variance V ( x ) ( see Eq . 5 ) over training time . 0 4 8 12 Input X 0 1 2 3 U pd at e St ep s × 10 6 10−4 10−2 Sa m pl in g Pr ob ab ili ty Figure 5 : Probability of sampling a data point at x over training time .

Summary: **Summary**  For regression tasks with heteroscedastic noise, a Gaussian model whose mean and variance are a function of the input is frequently used. Neural networks are then used as function approximators for the mean and variance and are fit using the negative log-likilhood. The paper identifies a problem with this approach: optimization can get stuck in configurations where the predicted mean is far from the true mean, as this is compensated for by a high predicted variance that also stalls learning. By adjusting the NLL objective, the authors compensate for the problem.