Refining Amortized Posterior Approximations using Gradient-Based Summary Statistics

Understand

We present an iterative framework to improve the amortized approximations of posterior distributions in the context of Bayesian inverse problems, which is inspired by loop-unrolled gradient descent methods and is theoretically grounded in maximally informative summary statistics.

• Amortized variational inference is restricted by the expressive power of the chosen variational distribution and the availability of training data in the form of joint data and parameter samples, which often lead to approximation errors such as the amortization gap.
• To address this issue, we propose an iterative framework that refines the current amortized posterior approximation at each step.
• Our approach involves alternating between two steps: (1) constructing a training dataset consisting of pairs of summarized data residuals and parameters, where the summarized data residual is generated using a gradient-based summary statistic, and (2) training a conditional generative model -- a normalizing flow in our examples -- on this dataset to obtain a probabilistic update of the unknown parameter.