Deep Heteroskedastic Regression: Post-Hoc Variance Estimation from Latent Representations

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 parametrises the mean and the variance of the predictive distribution. Yet to this day, training deep heteroskedastic regression models poses severe practical challenges in the trade-off between uncertainty quantification and mean prediction, such as variance overfitting, optimization difficulties and representation collapse. In this work, we identify the core issues and propose a simple and efficient procedure that addresses these jointly by post-hoc fitting a variance model across the intermediate layers of a pretrained network on a hold-out dataset. We demonstrate that this method is competitive with end-to-end trained mean-variance networks in heteroskedastic UQ on several data modalities. The method retains mean prediction accuracy, is cheap at both train and prediction time, and requires no additional data compared to existing methods. Finally, we show this method works on large scale foundation models for chemistry, paving the way for cheap heteroskedastic UQ in large scale neural networks without having to retrain such models end-to-end.