Practical Deep Heteroskedastic Regression

TL;DR

This paper introduces a practical method for deep heteroskedastic regression that improves uncertainty quantification by post-hoc variance modeling, addressing training challenges without sacrificing prediction accuracy.

Contribution

It proposes a simple, efficient post-hoc variance fitting approach that overcomes key training difficulties in deep heteroskedastic regression models.

Findings

Achieves state-of-the-art uncertainty quantification on molecular datasets.

Maintains mean prediction accuracy while improving uncertainty estimates.

Remains computationally cheap at inference time.

Abstract

Uncertainty quantification (UQ) in deep learning regression is of wide interest, as it supports critical applications including sequential decision making and risk-sensitive tasks. In heteroskedastic regression, where the uncertainty of the target depends on the input, a common approach is to train a neural network that parameterizes the mean and the variance of the predictive distribution. Still, training deep heteroskedastic regression models poses practical challenges in the trade-off between uncertainty quantification and mean prediction, such as optimization difficulties, representation collapse, and variance overfitting. In this work we identify previously undiscussed fallacies and propose a simple and efficient procedure that addresses these challenges jointly by post-hoc fitting a variance model across the intermediate layers of a pretrained network on a hold-out dataset. We demonstrate that our method achieves on-par or state-of-the-art uncertainty quantification on several molecular graph datasets, without compromising mean prediction accuracy and remaining cheap to use at prediction time.