[Paper Review] On the Pitfalls of Heteroscedastic Uncertainty Estimation with Probabilistic Neural Networks

Maximilian Seitzer, Arash Tavakoli+2 more|arXiv (Cornell University)|Mar 17, 2022

Gaussian Processes and Bayesian Inference37 citations

TL;DR

The paper analyzes why Gaussian negative log-likelihood (NLL) training can yield suboptimal mean fits in probabilistic regression and proposes β-NLL, a variance-weighted loss that interpolates between NLL and MSE to improve both mean accuracy and uncertainty estimation.

ABSTRACT

Capturing aleatoric uncertainty is a critical part of many machine learning systems. 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. In this work, we examine this approach and identify potential hazards associated with the use of log-likelihood in conjunction with gradient-based optimizers. First, we present a synthetic example illustrating how this approach can lead to very poor but stable parameter estimates. Second, we identify the culprit to be the log-likelihood loss, along with certain conditions that exacerbate the issue. 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. We show that using an appropriate $β$ largely mitigates the issue in our illustrative example. 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.

Motivation & Objective

• Motivate the importance of accurate aleatoric uncertainty estimation in neural networks.
• Diagnose why standard NLL training can produce stable yet poor mean fits.
• Propose a variance-weighted loss (β-NLL) to mitigate gradient bias and balance mean and uncertainty training.
• Demonstrate the approach across synthetic and real-world regression, dynamics modeling, and generative tasks.

Proposed method

• Formalize NLL for heteroscedastic Gaussian regression and identify gradient-variance dependency as a pitfall.
• Introduce β-NLL, where each data point’s loss contribution is weighted by σ^2^β with stop-gradient on the weighting term.
• Show that β=0 recovers NLL and β=1 yields MSE-like gradients for the mean while preserving uncertainty estimates.
• Provide guidance that β≈0.5 offers a good trade-off in practice.
• Empirically evaluate on synthetic sinusoidal regression, UCI regression, dynamics modeling, and generative/depth tasks to compare RMSE and NLL.

Experimental results

Research questions

• RQ1 Why can NLL-based training lead to suboptimal mean fits in probabilistic regression models?
• RQ2 Can a variance-weighted loss mitigate the underweighting of hard-to-predict regions and improve both mean accuracy and calibrated uncertainty?
• RQ3 How does β-NLL perform across synthetic and real-world datasets compared to NLL and alternative uncertainty losses?

Key findings

• NLL training can converge to suboptimal mean predictors due to gradient scaling by predicted variance, effectively undersampling hard regions.
• β-NLL weights per-sample gradients by σ^2^β, with β in [0,1], interpolating between NLL and MSE behavior.
• β≈0.5 generally provides the best trade-off between RMSE and log-likelihood across tasks.
• β-NLL achieves robust performance with less sensitivity to hyperparameters than NLL in dynamics modeling tasks.
• β-NLL provides well-calibrated uncertainty estimates and competitive or superior predictive accuracy on UCI, MNIST/Fashion-MNIST, NYUv2, and synthetic sine/dynamics datasets.