PaperlayerView exact revision on arXiv

Source-linked AI summary

On the Pitfalls of Heteroscedastic Uncertainty Estimation with Probabilistic Neural Networks

Maximilian Seitzer, Arash Tavakoli, Dimitrije Antic, Georg Martius

arXiv:2203.09168v2cs.LGstat.ML

TL;DR

The paper examines why standard NLL training for heteroscedastic probabilistic neural networks can yield poor, stable mean fits. It introduces β-NLL to control variance-dependent gradient weighting and reports better, more hyperparameter-robust performance across domains and tasks.

  • Problem

    Standard NLL training for heteroscedastic Gaussian regression can hinder accurate mean predictors because poorly predicted points receive reduced gradient weight.

  • Method

    β-NLL weights each data point’s loss contribution by its β-exponentiated variance estimate, interpolating between NLL and MSE while retaining uncertainty estimation.

  • Results

    β-NLL generally outperforms NLL across UCI regression, dynamics learning, generative modeling, and depth prediction, while being considerably easier to tune.

  • Takeaways & Limitations

    The β parameter lets practitioners control NLL’s self-regularization strength and reduce the tendency to ignore hard-to-fit regions.

  • Takeaways & Limitations

    β-NLL is designed to produce meaningful gradients, so its loss value should not be used to evaluate model quality; original NLL and optionally RMSE should be monitored instead.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Aleatoric uncertainty modeling is useful across machine-learning tasks, but the standard heteroscedastic Gaussian NLL approach can produce poor mean fits. The paper identifies this failure and proposes β-NLL, which reduces variance-dependent gradient effects and improves robustness.

  • Aleatoric uncertainty modeling supports stochastic dynamics learning, perception tasks, and risk-sensitive decisions.
  • Heteroscedastic Gaussian regression commonly predicts input-dependent means and variances by minimizing NLL with stochastic gradient descent.
  • The standard NLL procedure can produce subpar mean fits in addition to overconfident variance estimates.
  • NLL optimization can hinder accurate mean prediction because its gradients depend strongly on predictive variance, effectively undersampling poorly predicted data points.
  • β-NLL weights each data point by its β-exponentiated variance estimate, reducing gradient dependence on variance and largely improving performance across domains and tasks.

2 PRELIMINARIES

The paper formulates heteroscedastic Gaussian regression through input-dependent mean and variance functions estimated by a neural network. It contrasts NLL-based training with MSE and notes that the presentation focuses on univariate regression while extending to the multivariate case.

  • The analysis assumes Y given X is Gaussian with input-dependent mean µ(X) and variance σ^2(X).
  • A neural network estimates the mean and variance functions, with the variance constrained positive using an activation such as softplus.
  • The discussion focuses on univariate regression but states that the approach extends to the multivariate case.
  • Maximum likelihood estimation trains these estimates by minimizing the Gaussian negative log-likelihood, whereas standard regression minimizes mean squared error.
  • In practice, both objectives are optimized with stochastic gradient descent using mini-batches sampled from P(X,Y).

3 ANALYSIS

NLL training can settle into a stable but suboptimal mean fit because feature-space flatness and inverse-variance weighting suppress progress on poorly fit regions. As training proceeds, this creates self-reinforcing sampling imbalance, while local feature nonlinearity and symmetry breaking determine whether fitting progresses.

  • Training behavior: NLL can produce a suboptimal yet locally stable mean fit, even though the true mean minimizes the objective.The sinusoidal example contrasts this behavior with MSE, which converges to the optimal mean fit in reasonable time.
  • Symmetry and feature non-linearity: Initial feature-space flatness creates an undercomplex fit, while symmetry breaking is required to escape it.Mini-batch stochasticity, dataset asymmetries, and local feature nonlinearity can provide symmetry breaking.
  • Symmetry and feature non-linearity: Low feature-space nonlinearity remains associated with regions that stay poorly fit, whereas more granular Jacobian structure appears in regions later fit well.Figure 4 tracks Jacobian variance over training and supports a connection between feature nonlinearity and learning success.
  • Inverse-variance weighting: NLL scales mean gradients by inverse predicted variance, amplifying contributions from low-error points relative to high-error points.When well-fit regions have much smaller variance than badly fit regions, or are more numerous, progress on difficult regions can be hindered.
  • Inverse-variance weighting: The effective sampling probability of high-error points can fall from 10^-3 under uniform sampling to 10^-5 during training.This modified training distribution explains how NLL can effectively undersample difficult regions; the same behavior also carries over to a real-world dataset.

4 METHOD

β−NLL modifies NLL’s variance-dependent weighting to interpolate between NLL and MSE while retaining calibrated uncertainty estimates. The choice of β controls how difficult data points influence training, with β = 0.5 offering a favorable accuracy–log-likelihood trade-off in the experiments.

  • Variance-weighting the gradients of the NLL: β−NLL introduces a variance-weighting factor that modifies NLL’s effective data-point sampling distribution, interpolating between NLL and uniform importance.The parameter β controls this interpolation, with β = 0 corresponding to NLL and larger β reducing gradient dependence on predictive variance.
  • Variance-weighting the gradients of the NLL: For β = 0, β−NLL recovers NLL; for β = 1, its mean gradient is equivalent to MSE, while intermediate values provide other loss interpolations.The variance gradient at β = 1 remains distinct, with 2σ^2 in its denominator.
  • Variance-weighting the gradients of the NLL: β = 0.5 generally achieves the best trade-off between accuracy and log-likelihood in the experiments.At this value, data points are weighted by inverse standard deviation rather than inverse variance.
  • Variance-weighting the gradients of the NLL: The β−NLL loss is designed to produce meaningful gradients rather than to serve as a performance-evaluation objective.Training should be monitored with the original NLL objective and optionally RMSE for evaluating mean-fit quality.
  • Allocation of function approximator capacity: NLL can ignore difficult but important data points by assigning them low weight, whereas β−NLL makes this behavior adjustable through β.On ObjectSlide, NLL produces a long tail of ignored difficult points, while β−NLL fits easy points accurately without ignoring difficult ones.

5 EXPERIMENTS

Across synthetic and real-world tasks, β−NLL generally produces accurate predictions with better convergence and lower hyperparameter sensitivity than standard NLL. Comparisons with alternative losses reveal trade-offs in likelihood, accuracy, stability, and computational cost.

  • Cross-domain evaluation: β−NLL generally performs better than NLL and is considerably easier to tune across UCI regression, dynamics modeling, generative modeling, and depth prediction.The experiments evaluate RMSE and predictive NLL across these domains.
  • Synthetic datasets: Standard NLL fails to obtain reasonable mean fits on the sinusoidal regression problem, whereas β−NLL with β = 0.5 fits the function successfully.MSE and β−NLL with β = 1 also achieve the expected good mean fit.
  • Synthetic datasets: All β−NLL variants produce well-calibrated uncertainty estimates inside the training range on the heteroscedastic sine example.Moment matching significantly underestimates the variance everywhere.
  • Real-world datasets: β−NLL achieves predictive log-likelihoods on par with or better than NLL while clearly improving predictive accuracy on most UCI datasets.Table 1 reports predictive log-likelihood and RMSE across 12 datasets.
  • Real-world datasets: On dynamics tasks, NLL is highly hyperparameter-sensitive and has poor, high-variance predictive performance, while β−NLL performs well across configurations.Among β values, β = 0.5 generally achieves the best trade-off between accuracy and log-likelihood.
  • Comparison to other loss functions: Moment matching gives high accuracy but is unstable and has poor likelihoods, whereas Student’s t-based methods often improve likelihood at higher implementation and computational cost.β−NLL outperforms the Student’s t-based approaches in predictive accuracy, with exceptions on some limited-data UCI datasets.

6 CONCLUSION

The paper identifies premature convergence to suboptimal fits as a recurring NLL problem and proposes β−NLL as a simple, calibrated interpolation between NLL and MSE. Across experiments, β = 0.5 is generally a useful starting point.

  • 6 CONCLUSION: NLL training can get stuck in suboptimal function fits because badly fit regions receive progressively less loss weight.β−NLL rescales the original NLL gradient using the β-exponentiated per-sample variance.
  • 6 CONCLUSION: β−NLL interpolates between NLL and MSE while retaining well-behaved uncertainty estimates.The parameter β controls the strength of NLL’s self-regularization of high-noise or difficult-to-predict points.
  • 6 CONCLUSION: β = 0.5 is generally recommended as a starting point for balancing accuracy and log-likelihood.The conclusion frames this as a practical starting point rather than a universal optimum.

REPRODUCIBILITY STATEMENT

The paper makes its code and data publicly available, supporting reproduction of the described settings and experiments.

  • REPRODUCIBILITY STATEMENT: Full code and data are available at the authors’ beta-nll GitHub repository.The paper directs readers to Sections C and D for detailed settings.

A THE MOMENT MATCHING LOSS

The moment matching loss estimates Gaussian moments directly, using MSE for the mean and an analogous loss for variance. It can avoid premature convergence but may underestimate variances and become unstable.

  • A THE MOMENT MATCHING LOSS: Moment matching generally fixes NLL’s premature-convergence problem but underestimates variances and exhibits considerable training instabilities.These drawbacks limit it as a complete replacement for NLL-based training.
  • A THE MOMENT MATCHING LOSS: Moment matching estimates the first two Gaussian moments directly rather than relying on the likelihood loss.A Gaussian is fully described by its mean and variance.
  • A THE MOMENT MATCHING LOSS: The mean-fitting component is the standard MSE loss, while variance fitting uses an analogous loss for the second central moment.The combined formulation sums the mean and variance losses.
  • A THE MOMENT MATCHING LOSS: Moment matching makes balancing the mean and variance losses easier.This follows from expressing the total objective as the sum of separately defined losses.
  • A THE MOMENT MATCHING LOSS: β−NLL subsumes moment matching when β = 1 is used for the mean and β = 2 for variance estimation, producing the same gradients.The paper notes that this connection was not investigated further.

B.2 SYNTHETIC DATASET

On a sinusoidal toy regression problem, NLL training remains stably trapped in poor fits across trials and optimization settings. β−NLL improves convergence for many settings, especially when β exceeds 0.5.

  • B.2 SYNTHETIC DATASET: None of the tested optimizer configurations reached RMSE below 0.1, although the optimal value is RMSE 0.01.The result indicates that the poor fit is not specific to one optimizer.
  • B.2 SYNTHETIC DATASET: The poor NLL behavior remained stable across independent trials.The experiment repeated the sinusoidal training procedure across random seeds.
  • B.2 SYNTHETIC DATASET: For the NLL criterion, β > 0.5 produced good fits across many hyperparameter settings, including β = 1.At β = 1, good uncertainty predictions were also obtained with β−NLL, as indicated by low NLL scores.

B.3 ANALYSIS OF SAMPLING PROBABILITIES ON FETCH-PICKANDPLACE

The paper examines effective sampling probabilities and evaluates β−NLL across robotics, UCI regression, generative modeling, and depth prediction. Across these settings, β−NLL generally improves accuracy or robustness while preserving competitive likelihood.

  • B.3 ANALYSIS OF SAMPLING PROBABILITIES ON FETCH-PICKANDPLACE: On Fetch-PickAndPlace, NLL’s effective sampling distribution mismatched the MSE-based reference and drastically undersampled some data points.Some points were effectively never sampled during training.
  • B.3 ANALYSIS OF SAMPLING PROBABILITIES ON FETCH-PICKANDPLACE: On UCI datasets, β−NLL achieved predictive log-likelihoods on par with or better than NLL while improving predictive accuracy on most datasets.Student’s t baselines often had better predictive log-likelihood, while β−NLL was often on par with MSE for RMSE.
  • B.3 ANALYSIS OF SAMPLING PROBABILITIES ON FETCH-PICKANDPLACE: For MNIST and FashionMNIST VAEs, β−NLL with β > 0 provided a good compromise between reconstruction error and log-likelihood.It also produced meaningful uncertainties and semantically meaningful prior samples.
  • B.3 ANALYSIS OF SAMPLING PROBABILITIES ON FETCH-PICKANDPLACE: On NYUv2 depth regression, β−NLL with β > 0 achieved better RMSE than NLL, while β = 0.5 achieved similar RMSE to MSE and better performance on some other metrics.Predicted depth maps were noticeably sharper than those from NLL.

C DATASETS AND TRAINING SETTINGS

The experiments cover synthetic regression, UCI uncertainty benchmarks, robotic dynamics prediction, generative modeling, and NYUv2 depth regression. Training uses standardized dataset protocols, validation-based hyperparameter selection, and neural networks with probabilistic outputs.

  • C DATASETS AND TRAINING SETTINGS: The synthetic benchmarks include homoscedastic sinusoidal data with noise standard deviation 0.01 and heteroscedastic data generated from y = x sin(x) + xξ1 + ξ2.The heteroscedastic dataset contains 500 uniformly spaced samples on [0, 10].
  • C DATASETS AND TRAINING SETTINGS: The UCI evaluation uses established benchmark protocols and data splits, with custom random splits for carbon, energy, naval, superconductivity, and wine-white.Each dataset is evaluated across 20 randomly sampled 80%-20% train-test splits.
  • C DATASETS AND TRAINING SETTINGS: Learning rates are selected from six candidate values using validation log-likelihood, with training lasting up to 20,000 updates for the UCI experiments.Inputs and targets are whitened on the training set, while metrics are reported on the original scale.
  • C DATASETS AND TRAINING SETTINGS: The reported metrics for uncertainty estimation are RMSE and negative log-likelihood.These metrics are used to evaluate prediction quality quantitatively.
  • C DATASETS AND TRAINING SETTINGS: The robotic benchmarks include ObjectSlide and Fetch-PickAndPlace, where models predict changes in object position from states and actions.Fetch-PickAndPlace uses a 25-dimensional state space and four-dimensional action space.
  • C DATASETS AND TRAINING SETTINGS: For NYUv2 depth regression, models are trained for 25 epochs with batch size 16 and selected using the validation REL metric.The evaluation protocol follows the cited AdaBins setup.

D HYPERPARAMETER SETTINGS AND IMPLEMENTATION DETAILS

The experiments use neural networks with shared features and distribution-specific output parameterizations, alongside stated optimization, architecture, and hyperparameter-search procedures. The implementation also defines the β-NLL loss and its weighting parameter.

  • Gaussian parameterization and optimization: Gaussian predictors use two linear output layers over shared MLP features, with softplus-constrained variance, a 10^-8 floor, and a maximum variance of 1 000.All experiments use Adam with β1 = 0.9 and β2 = 0.999.
  • Alternative predictive distributions: Student’s t baselines integrate out Gaussian variance using a learned Gamma prior whose data-dependent alpha and beta parameters come from linear layers over shared features.The Student’s t predictive distribution is used by Student’s t, xVAMP(*), and VBEM(*).
  • Alternative predictive distributions: xVAMP and VBEM additionally predict categorical mixture logits, use initialized mixture-of-Gamma priors, and approximate the resulting KL divergence with 20 Monte-Carlo samples.The objective includes a KL divergence between a Gamma distribution and a mixture of Gamma distributions.
  • Sinusoidal regression: The sinusoidal regression network has two hidden layers with 128 neurons each, tanh activations, learning rate 5 · 10^-4, and batch size 100.The Sec. 5.1 experiment scans learning rates and architectures with different hidden layers and units per layer.
  • Hyperparameter search and evaluation: ObjectSlide and Fetch-PickAndPlace use 96 configurations per loss function in grid search, select models by validation log-likelihood, and retrain the best configuration with five random seeds for testing.The reported test results use a hold-out test set.
  • VAE experiment: The VAE experiment follows specified encoder and decoder architectures, uses 10- or 25-dimensional latent spaces, and trains with Adam, learning rate 0.0003, batch size 256, and early stopping.The encoder and decoder use three layers with the stated neuron widths and relu activations, with a maximum of 1 000 epochs and patience 50.
  • Depth-estimation implementation: The AdaBins-based model removes the mini-ViT Transformer and uses a two-channel U-Net output for mean and variance prediction, with both constrained positive by softplus activation.The first channel predicts the mean and the second predicts the variance.
  • β-NLL implementation: The β-NLL implementation computes Gaussian loss from squared error divided by variance plus log variance, then uses β in [0, 1] to control relative weighting across data points.β = 0 gives high weight to low-error points, whereas β = 1 gives equal weighting; the loss is returned per batch element.
Loading 2203.09168v2…