Scientists Just Made AI Smarter at Guessing Hidden Causes

Bayesian simulation-based inference (SBI) can struggle to recover increasingly concentrated posteriors as observation size grows. We consider learning sufficient summary statistics jointly with the posterior as a general principle for scalable inference in various SBI methods.

A common measure taken in SBI is dimension reduction through predetermined summaries, which may partially discard posterior information. Our approach instead imposes structure on how the posterior learner depends on a learned summary (e.g., affinity), and learns the summary simultaneously with the posterior.

We identify an important mechanism named two-way guidance: under a structural restriction on the summary, the posterior sufficiency of the learned summary simplifies posterior learning, whereas a suitable posterior objective guides the learned summary towards posterior sufficiency. This mechanism explains a counterintuitive finding that the structural restriction on the summary statistics improves posterior learning, rather than limit its performance. We establish theoretical guarantees on both posterior estimation error and the learned summary's departure from sufficiency, thus identifying conditions on the simulation budget for recovering increasingly concentrated posteriors as observation size grows.