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

TL;DR

This paper introduces a novel neural operator framework called LANO for efficient Bayesian optimal experimental design involving large-scale PDEs with infinite-dimensional parameters, achieving high accuracy and significant computational speedup.

Contribution

It develops a derivative-informed latent attention neural operator (LANO) that accelerates SBOED solutions with improved accuracy and scalability, incorporating new adaptive and approximation techniques.

Findings

LANO outperforms other neural architectures in accuracy.

LANO achieves 180x speedup over traditional methods.

The framework effectively designs MRI experiments for tumor growth monitoring.

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.