Function Space Bayesian Pseudocoreset for Bayesian Neural Networks
Balhae Kim Hyungi Lee Juho Lee
A Bayesian pseudocoreset is a compact synthetic dataset summarizing essential
information of a large-scale dataset and thus can be used as a proxy dataset
for scalable Bayesian inference. Typically, a Bayesian pseudocoreset is
constructed by minimizing a divergence measure between the posterior
conditioning on the pseudocoreset and the posterior conditioning on the full
dataset. However, evaluating the divergence can be challenging, particularly
for the models like deep neural networks having high-dimensional parameters. In
this paper, we propose a novel Bayesian pseudocoreset construction method that
operates on a function space. Unlike previous methods, which construct and
match the coreset and full data posteriors in the space of model parameters
(weights), our method constructs variational approximations to the coreset
posterior on a function space and matches it to the full data posterior in the
function space. By working directly on the function space, our method could
bypass several challenges that may arise when working on a weight space,
including limited scalability and multi-modality issue. Through various
experiments, we demonstrate that the Bayesian pseudocoresets constructed from
our method enjoys enhanced uncertainty quantification and better robustness
across various model architectures.