A Novel Framework Using Variational Inference with Normalizing Flows to Train Transport Reversible Jump Proposals
This work addresses computational bottlenecks in Bayesian model selection for statisticians and machine learning practitioners, but it is incremental as it builds on prior transport reversible jump methods.
The authors tackled the problem of inefficient trans-dimensional Bayesian inference by proposing a framework using variational inference with normalizing flows to train reversible jump proposals, resulting in faster mixing compared to existing baselines.
We propose a unified framework that employs variational inference (VI) with (conditional) normalizing flows (NFs) to train both between-model and within-model proposals for reversible jump Markov chain Monte Carlo, enabling efficient trans-dimensional Bayesian inference. In contrast to the transport reversible jump (TRJ) of Davies et al. (2023), which optimizes forward KL divergence using pilot samples from the complex target distribution, our approach minimizes the reverse KL divergence, requiring only samples from a simple base distribution and largely reducing computational cost. Especially, we develop a novel trans-dimensional VI method with conditional NFs to fit the conditional transport proposal of Davies et al. (2023). We use RealNVP flows to learn the model-specific transport maps used for constructing proposals so that the calculation is parallelizable. Our framework also provides accurate estimates of marginal likelihoods, which may facilitate efficient model comparison and help design rejection-free proposals. Extensive numerical studies demonstrate that the TRJ method trained under our framework achieves faster mixing compared to existing baselines.