Neural-Inspired Posterior Approximation (NIPA)

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.

💡 Research Summary

The paper introduces Neural‑Inspired Posterior Approximation (NIPA), a novel Bayesian sampling framework that draws inspiration from three interacting neural systems observed in human learning: model‑based (goal‑directed planning), model‑free (habitual responding), and episodic memory (rapid recall of specific experiences). The authors map these systems onto three complementary sampling modules.

1. Model‑Based (MB) Module – Implements a conventional Hamiltonian Monte Carlo (HMC) or stochastic‑gradient HMC (SGHMC) step that evaluates the exact log‑posterior and its gradient on the full dataset. This yields high‑quality proposals but is computationally expensive, making it suitable for exploring regions where no reliable surrogate or memory exists.

2. Model‑Free (MF) Module – Constructs a surrogate for the log‑posterior using past MB samples. An autoencoder first compresses the high‑dimensional parameter vector θ into a low‑dimensional latent code u. A deep neural network (DNN) regression model g(u) is then trained to predict log π(θ) from u. Once trained, the surrogate can generate cheap approximations of the log‑posterior, enabling fast “habitual” proposals without full data passes.

3. Episodic Control (EC) Module – Stores previously evaluated samples and their exact log‑posterior values. For a new candidate θ̃, the algorithm computes the minimum standardized Euclidean distance d*(θ̃,P) to the current pool P. If this distance is below a low threshold t₁, the algorithm directly reuses the nearest stored log‑posterior value, performing a one‑shot accept‑reject decision without any additional computation.

A gating mechanism with two thresholds (t₁ < t₂) decides which module to activate at each iteration:

• d* > t₂ → MB (explore new territory),
• t₁ < d* ≤ t₂ → MF (use learned surrogate),
• d* ≤ t₁ → EC (recall memory).

The workflow begins by running SGHMC on the full posterior to collect an initial pool P₀ (≈100 samples) and their exact log‑posterior values. These samples train the autoencoder and the DNN surrogate. During sampling, the algorithm repeatedly proposes a Gaussian random‑walk candidate, evaluates d*, selects the appropriate module, and updates the pool with accepted samples.

The authors apply NIPA to Bayesian neural networks (BNNs) for both regression and classification tasks, focusing on uncertainty quantification (UQ). Empirical comparisons include standard HMC, SGHMC, variational Bayes (VB), and the Calibration‑Emulation‑Sampling (CES) framework. Results show that NIPA achieves a higher effective sample size per unit time than the baselines, especially in high‑dimensional settings where CES’s Gaussian‑process surrogate becomes computationally prohibitive (O(N³)). NIPA’s DNN surrogate scales to thousands of parameters with modest overhead. Moreover, predictive intervals from NIPA are better calibrated than those from VB, which typically under‑estimate posterior variance.

Key contributions are:

• A principled, biologically motivated decomposition of Bayesian sampling into three interacting modules, providing a flexible architecture that can incorporate various MB, MF, and EC implementations.
• Introduction of an autoencoder‑DNN surrogate that reduces the dimensionality of the parameter space and replaces expensive log‑likelihood evaluations with fast neural approximations.
• A distance‑based gating scheme that dynamically balances exploration (MB), exploitation of learned patterns (MF), and rapid recall (EC), thereby improving both computational efficiency and sampling accuracy.
• Demonstration of superior performance on large‑scale BNNs, highlighting the framework’s scalability and its ability to deliver reliable uncertainty estimates.

Limitations include sensitivity to the amount of initial MB samples needed to train a reliable surrogate, the choice of latent dimension for the autoencoder, and potential loss of diversity when EC dominates. Future work suggested by the authors involves meta‑learning of the gating thresholds, exploring alternative surrogate models (e.g., normalizing flows), and theoretical analysis of convergence properties under non‑stationary gating decisions.