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. This procedure leads to iterative refinement of the amortized posterior approximations without the need for extra training data. We validate our method in a controlled setting by applying it to a stylized problem, and observe improved posterior approximations with each iteration. Additionally, we showcase the capability of our method in tackling realistically sized problems by applying it to transcranial ultrasound, a high-dimensional, nonlinear inverse problem governed by wave physics, and observe enhanced posterior quality through better image reconstruction with the posterior mean.
翻译:我们提出了一种迭代框架,用于改进贝叶斯逆问题中后验分布的摊销近似。该框架受循环展开梯度下降法启发,且在理论上基于最大信息摘要统计量。摊销变分推断受限于所选变分分布的表达能力及联合数据和参数样本形式的训练数据可用性,常导致摊销差距等近似误差。为解决此问题,我们提出一种逐步骤优化当前摊销后验近似的迭代框架。该方法交替执行两个步骤:(1)构建由摘要数据残差与参数配对组成的训练数据集,其中摘要数据残差通过基于梯度的摘要统计量生成;(2)在此数据集上训练条件生成模型(本文示例中采用正则化流),以获取未知参数的概率更新。该过程可在无需额外训练数据的情况下迭代优化摊销后验近似。我们通过将方法应用于一个范式化问题进行受控验证,观察到每次迭代后后验近似质量均得到提升。此外,我们展示了该方法在处理实际规模问题中的能力:将其应用于经颅超声成像——一个受波动方程控制的高维非线性逆问题,并通过后验均值重构图像的质量提升验证了后验分布估计的优化效果。