GNN-Guided Block Selection in Gibbs MCMC

Benjamin Dayan, Bhaskar Mishra, Justin Svegliato, Stuart Russell

Keywords: Bayesian Networks, probabilistic graphical models, posterior inference, approximate inference, Bayesian inference, probabilistic inference, inference, Markov Chain Monte Carlo, MCMC, Gibbs sampling, blocked Gibbs sampling, block selection, spectral gap, mixing time, Graph Neural Networks, GNNs

TL;DR: Amortise inference time with a GNN trained on a novel coupling detection heuristic to propose good blocks for speedier Blocked Gibbs sampling

Abstract:

Exact inference in large Bayesian Networks (BNs) is computationally intractable, limiting its practical application.
Markov Chain Monte Carlo (MCMC) methods like Gibbs sampling offer a scalable alternative but can be arbitrarily slowed by highly coupled variables--- addressable by jointly sampling some variables as a block.
We propose an automated block detection method to amortise inference time: training a Graph Neural Network (GNN) to propose blocks directly from the BN structure.
We further introduce a novel coupling heuristic based on the Markov chain's spectral gap, which we show can be more robust than existing heuristics.
Our GNN, trained on a dataset of small, randomly generated BNs, generalizes well to larger networks, accelerating MCMC sample efficiency in our experiments.

Submission Number: 115