LazyDINO: Fast, scalable, and efficiently amortized Bayesian inversion via structure-exploiting and surrogate-driven measure transport

TL;DR

The numerical results demonstrate that LazyDINO is highly efficient in cost amortization for Bayesian inversion, and outperforms Laplace approximation consistently using fewer than 1000 offline samples, while other amortized inference methods struggle and sometimes fail at 16,000 offline samples.

Abstract

We present LazyDINO, a transport map variational inference method for fast, scalable, and efficiently amortized solutions of high-dimensional nonlinear Bayesian inverse problems with expensive parameter-to-observable (PtO) maps. Our method consists of an offline phase in which we construct a derivative-informed neural surrogate of the PtO map using joint samples of the PtO map and its Jacobian. During the online phase, when given observational data, we seek rapid posterior approximation using surrogate-driven training of a lazy map [Brennan et al., NeurIPS, (2020)], i.e., a structure-exploiting transport map with low-dimensional nonlinearity. The trained lazy map then produces approximate posterior samples or density evaluations. Our surrogate construction is optimized for amortized Bayesian inversion using lazy map variational inference. We show that (i) the derivative-based reduced basis architecture [O'Leary-Roseberry et al., Comput. Methods Appl. Mech. Eng., 388 (2022)] minimizes the upper bound on the expected error in surrogate posterior approximation, and (ii) the derivative-informed training formulation [O'Leary-Roseberry et al., J. Comput. Phys., 496 (2024)] minimizes the expected error due to surrogate-driven transport map optimization. Our numerical results demonstrate that LazyDINO is highly efficient in cost amortization for Bayesian inversion. We observe one to two orders of magnitude reduction of offline cost for accurate posterior approximation, compared to simulation-based amortized inference via conditional transport and conventional surrogate-driven transport. In particular, LazyDINO outperforms Laplace approximation consistently using fewer than 1000 offline samples, while other amortized inference methods struggle and sometimes fail at 16,000 offline samples.

Figures (37)

• Figure 1: Overview of the RB-DINO construction.
• Figure 2: Overview of the latent representation lazy map construction.
• Figure 3: Overview of LazyDINO amortization procedure.
• Figure 4: Example I. Setup for inferring the diffusivity field in a nonlinear reaction--diffusion PDE detailed in \ref{['example1']}. For each BIP (#1--4), we show the data-generating synthetic parameter $m$ drawn from the prior, the synthetic data $\bm{y}$ placed on top of the PDE solution at $m$, and the MAP estimate $m_{\text{MAP}}^{\bm{y}}$, i.e., the solution of the deterministic inverse problem.
• Figure 5: Example II. Setup for inferring a heterogeneous hyperelastic material property detailed in \ref{['example2']}. For each BIP (#1--4), we visualize the synthetic parameter $m$ drawn from the prior, the corresponding deformed configuration, the synthetic displacement data $\bm{y}$, and the MAP estimate $m_{\text{MAP}}^{\bm{y}}$, i.e., the solution of the deterministic inverse problem.

Theorems & Definitions (13)

• Remark 1
• Proposition 2.1
• Remark 2
• Theorem 3.1: Posterior approximation through a ridge function surrogate
• Theorem 3.2: Surrogate approximation of the rKL objective gradient
• Corollary 3.3: Optimality gap for surrogate-driven LMVI
• Remark 3
• Remark 4
• Remark 5
• proof