Neural-Inspired Posterior Approximation (NIPA)
Babak Shahbaba, Zahra Moslemi
TL;DR
The paper tackles the computational bottleneck of exact posterior inference in high-dimensional Bayesian models by introducing Neural-Inspired Posterior Approximation ($MB$, $MF$, $EC$), a framework that unifies model-based sampling, surrogate modeling, and episodic-memory recall. It uses an initial pool from SGHMC, a latent-space autoencoder with a DNN surrogate, and a memory-based episodic component to decide between exact, surrogate-based, and memory-driven proposals via a gating rule $d^{*}(\tilde{\theta},\mathcal{P})$ and thresholds $t_1,t_2$. The approach yields substantial speedups over BNN-HMC while maintaining or improving uncertainty quantification across regression and classification tasks, demonstrating practical scalability for Bayesian deep learning. This framework offers a general template for combining MB/MF/EC modules in probabilistic inference and potentially extends to other Bayesian settings such as Gaussian processes and inverse problems.
Abstract
Humans learn efficiently from their environment by engaging multiple interacting neural systems that support distinct yet complementary forms of control, including model-based (goal-directed) planning, model-free (habitual) responding, and episodic memory-based learning. Model-based mechanisms compute prospective action values using an internal model of the environment, supporting flexible but computationally costly planning; model-free mechanisms cache value estimates and build heuristics that enable fast, efficient habitual responding; and memory-based mechanisms allow rapid adaptation from individual experience. In this work, we aim to elucidate the computational principles underlying this biological efficiency and translate them into a sampling algorithm for scalable Bayesian inference through effective exploration of the posterior distribution. More specifically, our proposed algorithm comprises three components: a model-based module that uses the target distribution for guided but computationally slow sampling; a model-free module that uses previous samples to learn patterns in the parameter space, enabling fast, reflexive sampling without directly evaluating the expensive target distribution; and an episodic-control module that supports rapid sampling by recalling specific past events (i.e., samples). We show that this approach advances Bayesian methods and facilitates their application to large-scale statistical machine learning problems. In particular, we apply our proposed framework to Bayesian deep learning, with an emphasis on proper and principled uncertainty quantification.
