Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Sequential infinite-dimensional Bayesian optimal experimental design with derivative-informed latent attention neural operator

Jinwoo Go, Peng Chen

TL;DR

This work tackles sequential Bayesian optimal experimental design for PDE-constrained systems with infinite-dimensional parameters by introducing an adaptive terminal formulation and an equivalent conditional KL objective, enabling scalable global decision-making. It couples a Laplace-based and low-rank posterior framework with a novel derivative-informed latent attention neural operator (LANO) that compresses inputs/outputs with DIS and PCA, propagates dynamics through latent attention, and yields accurate PtO maps and Jacobians with automatic differentiation. Numerical experiments on tumor-growth MRI design demonstrate that LANO achieves high accuracy for MAP points and eigenvalues while delivering substantial online/offline speedups (for example, up to 388× faster for PtO evaluations and 1364× faster for eigenpairs), enabling amortized speedups around 180×. The proposed framework offers a practical, scalable pathway to high-dimensional SBOED and suggests extensions to variational inference and spatial sensor placement for broader predictive digital-twin applications.

Abstract

We develop a new computational framework to solve sequential Bayesian optimal experimental design (SBOED) problems constrained by large-scale partial differential equations with infinite-dimensional random parameters. We propose an adaptive terminal formulation of the optimality criteria for SBOED to achieve adaptive global optimality. We also establish an equivalent optimization formulation to achieve computational simplicity enabled by Laplace and low-rank approximations of the posterior. To accelerate the solution of the SBOED problem, we develop a derivative-informed latent attention neural operator (LANO), a new neural network surrogate model that leverages (1) derivative-informed dimension reduction for latent encoding, (2) an attention mechanism to capture the dynamics in the latent space, (3) an efficient training in the latent space augmented by projected Jacobian, which collectively leads to an efficient, accurate, and scalable surrogate in computing not only the parameter-to-observable (PtO) maps but also their Jacobians. We further develop the formulation for the computation of the MAP points, the eigenpairs, and the sampling from posterior by LANO in the reduced spaces and use these computations to solve the SBOED problem. We demonstrate the superior accuracy of LANO compared to two other neural architectures and the high accuracy of LANO compared to the finite element method (FEM) for the computation of MAP points and eigenvalues in solving the SBOED problem with application to the experimental design of the time to take MRI images in monitoring tumor growth. We show that the proposed computational framework achieves an amortized $180\times$ speedup.

Sequential infinite-dimensional Bayesian optimal experimental design with derivative-informed latent attention neural operator

TL;DR

This work tackles sequential Bayesian optimal experimental design for PDE-constrained systems with infinite-dimensional parameters by introducing an adaptive terminal formulation and an equivalent conditional KL objective, enabling scalable global decision-making. It couples a Laplace-based and low-rank posterior framework with a novel derivative-informed latent attention neural operator (LANO) that compresses inputs/outputs with DIS and PCA, propagates dynamics through latent attention, and yields accurate PtO maps and Jacobians with automatic differentiation. Numerical experiments on tumor-growth MRI design demonstrate that LANO achieves high accuracy for MAP points and eigenvalues while delivering substantial online/offline speedups (for example, up to 388× faster for PtO evaluations and 1364× faster for eigenpairs), enabling amortized speedups around 180×. The proposed framework offers a practical, scalable pathway to high-dimensional SBOED and suggests extensions to variational inference and spatial sensor placement for broader predictive digital-twin applications.

Abstract

We develop a new computational framework to solve sequential Bayesian optimal experimental design (SBOED) problems constrained by large-scale partial differential equations with infinite-dimensional random parameters. We propose an adaptive terminal formulation of the optimality criteria for SBOED to achieve adaptive global optimality. We also establish an equivalent optimization formulation to achieve computational simplicity enabled by Laplace and low-rank approximations of the posterior. To accelerate the solution of the SBOED problem, we develop a derivative-informed latent attention neural operator (LANO), a new neural network surrogate model that leverages (1) derivative-informed dimension reduction for latent encoding, (2) an attention mechanism to capture the dynamics in the latent space, (3) an efficient training in the latent space augmented by projected Jacobian, which collectively leads to an efficient, accurate, and scalable surrogate in computing not only the parameter-to-observable (PtO) maps but also their Jacobians. We further develop the formulation for the computation of the MAP points, the eigenpairs, and the sampling from posterior by LANO in the reduced spaces and use these computations to solve the SBOED problem. We demonstrate the superior accuracy of LANO compared to two other neural architectures and the high accuracy of LANO compared to the finite element method (FEM) for the computation of MAP points and eigenvalues in solving the SBOED problem with application to the experimental design of the time to take MRI images in monitoring tumor growth. We show that the proposed computational framework achieves an amortized speedup.
Paper Structure (29 sections, 1 theorem, 63 equations, 7 figures, 5 tables, 2 algorithms)

This paper contains 29 sections, 1 theorem, 63 equations, 7 figures, 5 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Problem formulation
  3. Bayesian inverse problem
  4. Sequential Bayesian optimal experimental design
  5. Adaptive terminal formulation of SBOED
  6. Scalable approximations for SBOED
  7. High-fidelity discretization
  8. Laplace approximation of the posterior distribution
  9. Low-rank approximation of the posterior distribution
  10. Computation of the optimality criteria of SBOED
  11. Adaptive optimization for SBOED
  12. Derivative-informed latent attention neural operator
  13. Derivative-informed dimension reduction
  14. Latent attention neural operator
  15. Data generation and derivative-informed training
  16. ...and 14 more sections

Key Result

Theorem 1

Let $\mu(m|{\boldsymbol{y}}_{1:i:K}, \mathop{\mathrm{\boldsymbol{\xi}}}\nolimits_{1:i:K})$ denote the posterior distribution for the observations ${\boldsymbol{y}}_{1:i:K}$ given experimental design $\mathop{\mathrm{\boldsymbol{\xi}}}\nolimits_{1:i:K}$, then the optimization problem eq:cumulativeBOE

Figures (7)

  • Figure 1: Left: Illustration of gray and white matter in a rat's brain. Middle: Mean of the prior distribution $m_\text{prior}$. Right: A random sample drawn from the prior distribution $m \sim {\mathcal{N}} (m_\text{prior}, {\mathcal{C}}_\text{prior})$.
  • Figure 2: Left: Initial tumor implantation at day $t_0 = 0$. The volume fraction of the tumor at day $t_{40} = 4$ (middle) and day $t_{90} = 9$ (right) at a random sample of the parameter.
  • Figure 3: Decay of the eigenvalues of DIS for input parameter dimension reduction (top left) and singular values (bottom left) by SVD for PCA output dimension reduction, and their corresponding modes.
  • Figure 4: Comparison of the approximation of the PtO map/state by our neural network (NN) surrogate and FEM computation. NN (top), FEM (middle), and their difference (bottom) on days 1, 2, 4, and 8.
  • Figure 5: MAP points computed by FEM (left) and our NN surrogate (middle), and their difference (right) for a random sample drawn from the prior. Top: daily observations, bottom: observations at day 2, 5, 8.
  • ...and 2 more figures

Theorems & Definitions (3)

  • Example 1
  • Theorem 1
  • proof