Quantifying Uncertainty in the Presence of Distribution Shifts

TL;DR

This paper introduces an adaptive Bayesian framework with amortized variational inference to improve uncertainty estimates in neural networks under covariate distribution shifts, validated on synthetic and real data.

Contribution

It proposes a novel adaptive prior conditioned on covariates and an efficient inference method to better quantify uncertainty during distribution shifts.

Findings

Significantly improves uncertainty estimation under covariate shifts

Effective in synthetic environments simulating various shifts

Outperforms traditional fixed-prior methods

Abstract

Neural networks make accurate predictions but often fail to provide reliable uncertainty estimates, especially under covariate distribution shifts between training and testing. To address this problem, we propose a Bayesian framework for uncertainty estimation that explicitly accounts for covariate shifts. While conventional approaches rely on fixed priors, the key idea of our method is an adaptive prior, conditioned on both training and new covariates. This prior naturally increases uncertainty for inputs that lie far from the training distribution in regions where predictive performance is likely to degrade. To efficiently approximate the resulting posterior predictive distribution, we employ amortized variational inference. Finally, we construct synthetic environments by drawing small bootstrap samples from the training data, simulating a range of plausible covariate shift using only the original dataset. We evaluate our method on both synthetic and real-world data. It yields substantially improved uncertainty estimates under distribution shifts.