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: The authors analyze the known phenomenon that the optimization of the Gaussian likelihood of a probabilistic neural network with heteroscedastic output variance using gradient ascent can get stuck at solutions with a suboptimal mean fit, compensated by large output variance: the likelihood amplifies non-uniform initial distributions of feature granularity due to variance weighting, i.e., the gradient w.r.t. the mean are scaled by the inverse variance leading to well-fitted samples dominating the gradient. The authors argue that this effect can be undesirable because it can prevent spending model expressivity on hard-to-fit regions, while it can be desirable in other situations because it allows the network to ignore outliers arising due to noise. This is complementary to the MSE objective which is dominated by badly-fitted samples (and thus hard to fit regions) but sensitive to outliers, and which does not allow to fit the output variance. The authors propose a modified loss function, $\beta$-NLL, which interpolates between the likelihood and the MSE loss functions with the incentive to retain the "best of both worlds". They compare $\beta$-NLL with standard NLL, MSE, and a moment-matching version on several datasets.