Quantifying Uncertainty in the Presence of Distribution Shifts

Yuli Slavutsky, David Blei

Keywords: Uncertainty estimation, distribution shift, variational inference, bayesian learning, out-of-distribution generalization

TL;DR: This work proposes a Bayesian framework with variational inference that adapts the prior and posterior to covariate shifts, improving uncertainty estimates by capturing predictive reliability rather than just input dissimilarity.

Abstract: Neural networks make accurate predictions but often fail to provide reliable uncertainty estimates, especially when test-time covariates differ from those seen during training, as occurs with selection bias or shifts over time. To address this, we propose a Bayesian framework for uncertainty estimation that explicitly accounts for covariate shifts. Unlike conventional approaches that rely on fixed priors, a 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 shifts using only the original dataset. We evaluate our method on both synthetic and real-world data, demonstrating that it yields substantially improved uncertainty estimates under distribution shift compared to existing approaches.

Primary Area: Probabilistic methods (e.g., variational inference, causal inference, Gaussian processes)

Submission Number: 13647