Optimal Experimental Design for Bayesian Inverse Problems using Energy-Based Couplings

Authors: paula  cordero encinar
tobias  schröder
andrew b duncan
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-10 12:16:58

Abstract: Bayesian Experimental Design (BED) is a robust model-based framework for optimising experiments  but faces significant computational barriers, especially in the setting of inverse problems for partial differential equations (PDEs). In this paper, we propose a novel approach,  modelling the joint posterior distribution with an energy-based model, trained on simulation data.   Unlike existing simulation-based inference approaches, we leverage implicit neural representations to learn a functional representation of parameters and data.  This is used as a resolution-independent plug-and-play surrogate for the posterior, which can be conditioned over any set of design-points, permitting an efficient approach to BED.

Main paper: 
 introduction: 
 Mathematical models based on partial differential equations (PDEs) play a crucial role in quantitative analysis of complex systems arising in science, engineering and socio-economics (Burger et al., 2014) . These often involve large numbers of parameters which must be calibrated. In many settings, the unknown parameters are functions, e.g. spatially-varying coefficient fields. Identifying parameters based on measurement data involves solving an inverse problem: Given a possibly stochastic forward operator G : A → U between a parameter space A and a solution space U, and empirical observations of the solution y ∈ R n at locations x = (x 1 , . . . , x n ), x i ∈ Ω ⊂ R d , we aim to infer the parameter a ∈ A such that y = G(a)(x) + η, where η denotes mean-zero observational noise. The Bayesian approach to inverse problems (Stuart, 2010 ) is a systematic method for uncertainty quantification of parametric estimates, where a and y are viewed as coupled random variables. Equipping a ∈ A with a prior distribution, Bayes' rule yields a posterior distribution for a given y.A significant challenge is determining which measurement positions x yield the most information about an unknown parameter. This question can be tackled using Bayesian Experimental Design (BED) (Chaloner & Verdinelli, 1995) , a versatile framework that guides the process of data collection. Typical approaches to BED involve a nested Monte Carlo approach, with Markov Chain Monte Carlo (MCMC) methods used to estimate a utility for any given candidate design, which is subsequently optimised over admissible designs (Rainforth et al., 2023) . In the setting of a PDEgoverned inverse problem, each MCMC step necessitates at least one computation of the underlying PDE solution, meaning that the process quickly becomes computationally demanding (Alexanderian, 2021). Traditional surrogate model approaches (Gramacy, 2020) are not readily applicable if the parameters are functional, or if the forward map exhibits non-Gaussian noise. In the latter case, the likelihood becomes intractable due to the introduction of auxiliary variables, and one must resort to expensive pseudo-marginal MCMC methods (Andrieu & Roberts, 2009) , simulation based inference approaches (Glaser et al., 2022) or synthetic likelihoods (Price et al., 2018) .Our work offers a computationally efficient alternative to these methods by learning a generative model for the joint posterior distribution over (a, G(a)), i.e., a joint distribution over the parameters of interest and the value of the (possibly stochastic) solution map. This extends the scope of classical surrogate methods, since no deterministic relationship between the parameters and solutions of the model is assumed. Crucially for BED, the proposed method models the association between the parameter a and the functional-form of the solution G(a), and thus is not dependent on a fixed set of evaluation points x or a particular discretisation of the domain. To achieve this, we embed the functions in a finite-dimensions using implicit neural representations (Sitzmann et al., 2020) and use an energy-based model (Lecun et al., 2006) to model the distribution of latent representations. 
energy-based coupling on function spaces: 
 We propose to approximate the solution operator G using a probabilistic model p θ over the joint distribution of (a, G(a)) such that high likelihood regions of p θ correspond to solutions u = G(a).To do so, we proceed in two steps: Using training data consisting of pairs {(a i (x j i ), u i (x j i )) Ni j=1 } M i=1 with u i (x j i ) = G(a i )(x j i ), x j i∈ Ω, we encode the functions a and u into finite-dimensional latent codes z a , z u . On the finite-dimensional representation space, we learn the joint distribution of the latent codes using a so-called energy-based model, where p θ (z a , z u ) ∝ exp(-E θ (z a , z u )) . Our workflow is visualised in Figure 4 . Through the learnt latent representations, the model is independent of any pre-specified evaluation grid. This not only allows inference of solutions from observations at arbitrary points, but it also enables Bayesian experimental design of sensor placement as a downstream task.Following COIN++ (Dupont et al., 2022) , we compress the functional data points (a i ) M i=1 and (u i ) M i=1 into implicit neural representations defined by a i (•) = g ψ (•, z ai ) and u i (•) = f ϕ (•, z ui ). We generate training data by sampling parameters a ∈ A from a pre-specified prior distribution and simulating the forward dynamics G(a). We then train the neural representations by minimising the mean square error between the function value and the neural network prediction at random evaluation points. Each layer of g ψ and f ϕ takes the form of a SIREN layer sin(ω 0 (W h + b + β)) (Sitzmann et al., 2020) . While the weights and biases W, b of each layer are shared among all data points, the shifts β depend on the data point specific latent code z and are trained individually for each functional data point. This produces highly compressed latent representations (z ai , z ui ) M i=1 , which can be decoded efficiently by applying the maps g ψ and f ϕ , respectively.Energy-based models (EBMs) (Lecun et al., 2006) are unnormalised statistical models of the form exp(-E θ )), where the energy-function E θ is typically modelled with a scalar-valued neural network. We assume lossless compression in the INR and learn the joint distribution of a and u = G(a) as an energy-based model over tuples (z ai , z ui ). The architecture of the energy function first embeds the latent codes into an embedding space of shared dimension and processes the embedded vectors using MLPs. The suggested architecture is resilient to vanishing signals, ensuring that information flows in a stable manner (for details, see Appendix B.2). Since unnormalised models are not amenable to optimisation with maximum likelihood estimation, we explore training the model with contrastive divergence (Hinton et al., 2006) and energy discrepancy (Schröder et al., 2023) , achieving the best results with the latter approach.At inference time, we are given noisy observations of the system at finitely many evaluation pointsD = {(x i , y i )} n i=1 with y i = u(x i ) + η i , where η i ∼ N (0, σ 2 ).The posterior distribution of the latent representations (z a , z u ) conditioned on the observed data is given byEQUATIONThe noise assumption, together with the PDE model, results in p(x i ,y i |z a , z u ) ≈ N (y i ; f ϕ (x i , z u ), σ 2 ).The desired parameter solution pair (a, u = G(a)) can now be sampled from the posterior using stochastic gradient Langevin dynamics (Welling & Teh, 2011) .One of the main motivations of our approach lies in Bayesian experimental design. Specifically, we seek to determine optimal sparse sensor placement positions d = {d 1 , d 2 , . . . , d D } for the inference of (a, u = G(a)) based on y = u(d) + η. We measure the utility of a sensor placement position by calculating the expected information gain over the prior as measured by relative entropy, i.e. we use the utility functionU (d) := E p(y|d) D KL (p(z a , z u |y, d) ∥ p(z a , z u )).Using Bayes theorem and Monte Carlo estima-tion we obtain the following biased utility estimatorEQUATIONwhere (z ai , z ui ), (z aj , z uj ) ∼ p(z a , z u ) and y i ∼ N (f ϕ (d, z ui ), σ 2 ). The utility is sequentially optimised using Bayesian optimisation. For details, see Appendix C. 
numerical experiments: 
 Our training data consists of M pairs of parameters and their corresponding solutions,{(a i , u i )} M i=1for a i ∈ A and u i ∈ U. We assume access to only N i point observations of them, where the set of N i locations varies across the M function realizations and need not be the same for a and u. While the method can handle functional parameters through the INR encoding, we assume for simplicity in the first presented example that a is parameterised by a real-valued vector of finite dimension. We emphasise that PDE solutions are only required to train the INR and EBM models. Once trained, these models are reused for inference leading to high savings in terms of computational cost.Consider the boundary value problem (BVP) on the interval [-1, 1] given by the non-linear PDE The training data consists of pairs (a, b) and their corresponding solution for a realisation of X a and X b . Figure 1 compares the numerical solutions of the BVP and samples from the learnt energybased coupling. One can see that the sampled solutions resemble the numerical solutions and respect the boundary conditions. As a quantitative test we perform dimensionality reduction via t-SNE (van der Maaten & Hinton, 2008) which demonstrates that training data and generated data do not form clusters, i.e. the t-SNE test cannot distinguish between training and generated data.u ′′ (x) -u 2 (x)u ′ (x) = f (x), u(-1) = X a , u(1) = X b , f (x) = -π 2 sin(πx) -π cos(πx) sin 2 (πx), X a ∼ N (a, 0.3 2 ), X b ∼ Unif(b -0.3, b + 0.4), a, b ∼ Unif(-3, 3).We perform inference on the parameters a and b when true values are set at random to a = 2.9 and b = -0.11, based on sparse observations of the PDE solution D = {(x i , y i )} i=1,...,10 , where y i = u(x i ) + η i , and η i are iid N (0, 0.1 2 ). Inference results for 10 3 posterior samples from p θ (z a , z u |D) are summarised in Figure 2 . The left and right panels display histograms for sampled a and b, respectively, while the central plot shows the generated solutions f ϕ (•, z u ), together with the true solution and the observations D.We observe a strong agreement between posterior samples and the ground truth. The relative L 2 error norm of the true solution and the posterior mean is ∥u -u truth ∥ 2 /∥u truth ∥ 2 = 0.065 and the MSEs for a and b are 0.37 and 0.25, respectively. It is important to mention that we observe comparable values in all performance metrics regardless of the randomly chosen PDE parameters (a, b) and the set of noisy observations D.Next we consider learning the diffusion coefficient κ of the 2D Darcy flow equation defined by the PDE -∇ • κ(x)∇u(x) = f (x) with domain x ∈ Ω = [0, 1] 2 and Dirichlet boundary conditions u| ∂Ω = 0. We assume that κ and f take the form Although the diffusion coefficient is characterised by just two parameters, we treat it as a function and learn an INR as outlined before. Using the learnt energy-based coupling, we perform an experimental design task in which, based on three initial observations (intentionally chosen in non informative places) of a solution with parameters a = 1 and b = 1, we maximise the estimated utility function to find optimal locations for ten additional measurement sites. To benchmark the efficacy of the new design locations, we generate the same number of points from a Sobol sequence (Sobol', 1967) within our domain and compare the posterior means along with the ground truth. We observe that results inferred from optimally chosen points closely match the ground truth unlike those inferred from Sobol points, see Figure 3 . Corresponding plots for the diffusion coefficient can be found in the Appendix, Figure 6 . In particular, the relative L 2 error norms between the ground truth and the posterior mean for the solution and the diffusion coefficient have lower values when performing inference based on optimal location measurements (see Table 1 ), with performance remaining constant over the entire domain of a and b. It should be noted that the inference for Sobol points is also conducted using the learnt energy-based coupling. κ((x 1 , x 2 )) = exp[(1.5 + a • cos(πx 1 )(-1.5 + b • cos(πx 2 )], f = 0.5 + ε where ε ∼ N (0, σ 2 ).∥ u -u tr ∥ 2 /∥u tr ∥ 2 ∥ log κ -log κ tr ∥ 2 /∥ 
discussion and related work: 
 The proposed approach is strongly related to the problem of Neural Operator Learning, a class of models that learn mappings between function spaces and solve partial differential equations. Examples include Fourier Neural Operators (Li et al., 2021) and Deep O-Nets (Lu et al., 2019) . A common thread across these methods is their resolution invariance, i.e. their ability to generate predictions at any resolution. Existing work has largely focused on learning the deterministic relationship between the input and output, although recent work has sought to extend these approaches to developing generative models on function space (Rahman et al., 2022; Lim et al., 2023) . See also (Salvi et al., 2022) . Our approach leverages similar resolution-invariant paradigms, integrating them into an EBM architecture to provide a generative associative map between function spaces. The results above already demonstrate the effectiveness of the proposed methodology, but a detailed exploration of the impact of different neural operator architectures is left for future work. These approaches holds promise for accelerating scientific discovery (Azizzadenesheli et al., 2024) , such as enhancing the pathway to sustainable nuclear energy (Gopakumar et al., 2023; Kobayashi & Alam, 2024) .Boundary value problem The dataset consists of 20000 pairs of boundary conditions (a, b) and associated solutions. Each solution is evaluated at N i points, where N i is a randomly chosen integer between 30 and 40.The dataset consists of 20000 pairs of diffusion coefficients and solutions of the PDE. Each solution and diffusion coefficient is evaluated at 784 points.We have implemented all experiments with PyTorch (Paszke et al., 2019) .The training is done in two steps (as shown in Figure 4 ). First, we train the modulated INRs to represent the data. Once, the INRs have been fitted, we obtain the latent representations of the functions of interest. The latent modulations of the INRs are concatenated, in the cases where we seek for a prior over different functional parameters. Furthermore, as the data is encoded using only a few steps of gradient descent (for details, see Dupont et al. (2022) ), the resulting standard deviation of the codes is very small, falling within the range of [10 -3 , 10 -1 ]. Therefore, these raw latent representations are not appropriate for subsequent processing. To address this, we standardise the codes by subtracting the mean and dividing by the standard deviation. The normalised latent embeddings are then used to train the EBM model. The means and standard deviations are stored to unnormalised the samples generated by the EBM before using them as modulations for the INR to recover the functions in real space. It is important to remark that, unlike the INR representation, the use of the truncated coefficient expansion for a basis, such as the Fourier basis, to represent the function of interest does not scale well with dimension. The criterion to determine the value of the latent dimension is that the mean relative L 2 error norm between the true function (either solution or parameter) and the one reconstructed from the latent embedding across all data points remains below a certain threshold. At the same time, we want the dimension to be low to ensure that the resulting distribution of the latent representations has a positive probability density that can be easily modelled by the EBM. For the diffusion coefficient, the relative L 2 error norm is of the order of 10 -8 for a latent dimension of 2 and the error grows as the dimension increases or decreases. In particular, for a dimension greater than 2 the additional dimensions are redundant, with values repeating across dimensions. For the solutions the mean relative L 2 error norms for different latent dimensions computed on a validation set not seen during training for the BVP and the Darcy flow equation are shown in Tables 2 and 3 , respectively. 2020), so that it can take arbitrary point evaluations of the functions of interest and not just random points of a fixed grid. We have also followed their initialisation scheme.In all our experiments, each training point for the EBM consists of two parts, the PDE solution and a functional or vector-valued coefficient, and our goal is to understand the connection between the two by learning their joint probability density. To do this, we first uplift each part of the input vector into a latent space (using an encoder) so that they have the same dimension (equal to 128) and then propagate and merge them, with shared connections between the two branches. The architecture of the network is illustrated in Figure 5 .The specific structure of each element of the architecture is the following• Encoder block Encoder(x) = Linear(σ • Linear(y) + y), y = σ • Linear(x),where σ is a GELU (Hendrycks & Gimpel, 2016 ) activation function.Figure 5 : Energy-based model neural network architecture, where z a and z u the finite dimensional latent embeddings of two functions, such as, the solution and a coefficient of the PDE.Block 1 1 k (x) = • Linear(x), where σ is a RELU activation and k the dimension. • ∆ ∆(x, (1 -z) • x + z • y, where • denotes point-wise multiplication. Note that the vectors x, y and z need to have the same dimension and the operation is just an interpolation.• + operation is a point-wise addition.The dimension remains constant at 128 in the 2 branches of the architecture. Therefore, when both vectors are concatenated, it becomes 256. The last two Block 1's reduce the dimension from 256 to 128 and from 128 to 64, respectively.As presented above, the experimental design task consists in optimally selecting d to maximise the information gain about the solution. In our proposed framework, this PDE solution is approximated by f ϕ (•, z u ), where z u is the associated latent embedding. Mathematically, the utility function for d needs to maximise the expected information gain over the prior p(z a , z u ), as measured by relative entropy. This is equivalent to maximising the expected KL-divergence (d) = E D KL p(z|d, y) ∥ p(z) = dz dy log p(z|d, y) -log p(z) p(z, where the expectation is computed over the predictive distribution of the new (yet unobserved) p(y|d) and sake of notation simplicity, we z = (z a , z u ). Applying Bayes theorem, we rewrite the above in a form amenable to estimation (3)Notice that gradient-based optimisation of Eq. ( 3) with respect to d may suffer from high of the gradients arising from the need to differentiate through the probability density p(z, y|d). This can be avoided by using Bayesian optimisation, wherein a differentiable surrogate function, such as a Gaussian process, is to U (d). Employing Monte Carlo to compute the predictive distribution p(y|d), we derive the following estimator for the utility function  
 
 Appendix:  


Decision: Authors propose using implicit neural representations coupled with energy-based models to provide an efficient surrogate for Bayesian experimental design. The high-level idea is interesting; I suggest authors make the following revisions to the paper for the camera-ready version. 1. clearly explain the contributions of the paper 2. better explain how the proposed method is different from prior work 3. provide more implementation details 4. quantify efficiency gains over baselines, and 5. demonstrating benefits on real-world problems for the camera-ready version.

Reviews: review: title: title: Good approach to Bayesian experimental design coupling representations and energy-based model
main_review: main_review: Evaluation
-------------------------
**Quality:** The paper introduces a novel approach, Bayesian Experimental Design (BED), for optimizing experiments in the setting of inverse problems for partial differential equations (PDEs). It presents a thorough explanation of the proposed method, including mathematical formulations and implementation details. The approach is evaluated through numerical experiments, demonstrating its effectiveness in various scenarios. Additionally, the paper discusses the implications of the results and compares them with existing methods, providing insights into the strengths and limitations of the proposed approach. Overall, it is a high-quality work. 

**Clarity:** The paper is generally well-written and organized, with clear sections outlining the problem statement, methodology, experiments, and discussion. Technical terms and concepts are defined and explained adequately, enhancing understanding for readers with varying levels of expertise. Figures and tables illustrate key points and present experimental results. However, some sections could benefit from additional clarification, particularly regarding the training process of implicit neural representations and energy-based models. This is not a problem for this short piece, especially because a very good supplementary material is provided for further clarification & details.

**Originality:** The paper introduces a novel approach to Bayesian Experimental Design for optimizing experiments in inverse problems for PDEs. It leverages implicit neural representations and energy-based models to model the joint posterior distribution efficiently, offering a unique perspective compared to existing simulation-based inference approaches. Integrating these techniques for experimental design in the context of PDEs is innovative, in my opinion, and extends the current SOTA in the field. 

**Significance:** The proposed approach addresses significant computational barriers in optimizing experiments for inverse problems in PDEs, offering a computationally efficient alternative to existing methods. The potential applications of the proposed methodology are broad, spanning various domains where PDEs are used for modeling complex systems. Besides, the paper opens avenues for future research in neural operator learning and Bayesian experimental design, with implications for accelerating scientific discovery and solving real-world problems.

Pros
-------------------------
1. Introduces a novel approach to BED for optimizing experiments in the setting of inverse problems for PDEs.
2. Leverages implicit neural representations and energy-based models to model the joint posterior distribution efficiently.
3. Demonstrates effectiveness through numerical experiments on boundary value problems and diffusion equations in 1D and 2D.
4. Offers a computationally efficient alternative to existing simulation-based inference approaches.
5. Provides insights into the potential applications and future research directions in neural operator learning and BED

Cons
-------------------------
1. Some sections could benefit from additional clarification, particularly regarding the training process of implicit neural representations and energy-based models.
2. While the numerical experiments demonstrate effectiveness, further validation on more diverse datasets and real-world applications may be needed.
4. The computational cost of training implicit neural representations and energy-based models could be high, especially for larger-scale problems.

Typos and grammar issues
-------------------------
- "coupled random variables" - consider adding a hyphen for clarity ("coupled-random variables").
- "Equipping a with a prior distribution" ---> "Equipping 'a' with a prior distribution"
- "Each MCMC step necessitates at least one solve of the underlying PDE" - "solve" should be "solving"?
- "Design points kb" - unclear heading, needs clarification or correction.
- "experimental design task in which, based on three initial observations" --> "experimental design task involves using three initial observations"
- "results inferred from optimally chosen points are visually more similar to the ground truth" --> "results inferred from optimally chosen points closely match the ground truth"
- clean the bib; missing entries, capital letters. This is a nice tool btw: https://flamingtempura.github.io/bibtex-tidy/.
title: title: The authors present and showcase an interesting idea for Bayesian Experimental Design, combining implicit neural representations and energy-based models with promising results.
main_review: main_review: Summary:
The work showcases a novel and intriguing approach to Bayesian Experimental Design (BED) aimed at overcoming computational barriers in inverse problems for partial differential equations (PDEs). The key idea is to learn implicit neural representations of the parameters $a$ and the respective solutions $u$ and to use an energy-based model to learn a joint distribution over these latent encodings. This yields an efficient and resolution-independent surrogate for the joint posterior distribution, enabling more effective optimization of experiments within the BED framework. It is a highly relevant problem with many applications in complex systems.

Pros:
- The manuscript is well-written, and the authors successfully provide condensed information on the broad range of concepts they bring together in their approach without compromising completeness.
- Relevant related work is discussed sufficiently.
- They neatly showcase the capabilities of their approach in two examples.

Cons:
- A comparison to some other more recent method would be desirable.

Further comments:
- It would help the reader if in the supplement B equation (3), the reformulation were spelled out.
- It would be helpful to explain briefly how the compressed latent representations are precisely retrieved from the COIN++ framework.
- As they are discussed in the text, the t-SNE plots should be added to the appendix.
title: title: Bayesian experimental design using energy-based coupling and neural operators
main_review: main_review: Summary: The authors propose using implicit neural representations with energy-based coupling to provide an efficient surrogate for Bayesian experimental design. 

Strengths and Weaknesses:

* It was slightly confusing to have Figure 2 come before Figure 1

[+] The paper is well-motivated and well-written. The application of neural operators to Bayesian experimental design (BED) is novel, and the authors make it clear why existing approaches would be inefficient compared to the proposed approach. 

[-] Given the author's claim to prior inefficiencies for BED, i.e., having to solve the PDE for each MCMC step, it would be nice to see some baseline comparison highlighting how the author's approach is much faster and more efficient than prior approaches. I see no reporting of training or inference times. 

[-] Does the author's INR model differ at all from CORAL (Serrano et. al., 2023)? It's stated the workflow is based on it, so I think the authors should state it by name, i.e., CORAL, like they do FNOs and DeepONets. This clarifies that the contribution is not CORAL itself, which has already been published. In the author's INR section, this is not referred to at all, while it is mentioned in the appendix, so it's a bit confusing what the authors are claiming as their contribution. This could be a misunderstanding of the author's modification to CORAL if that is the case, but I still think it should be named and clarified in any case. 

[-] I fail to see a strong contribution with the problem solved in section 3.1. It appears the authors are just using the CORAL or, alternatively, FNO/DeepONet architectures, which they propose as future work, to perform inference. It seems to me the main contribution of the paper is the application of these function to function map neural architectures to BED. In the case of the problem in 3.2, this contribution is clearer, and the application to BED is shown to be better than Sobol points. The abstract is about this application to BED, so I fail to see how section 3.1 adds to that beyond standard operator learning results, albeit with energy-based coupling. Even in that case, the results of the learned coupling (using EBMs) have more spread and include high-frequency oscillations compared to the numerical simulations. 

Conclusion: While the application of neural operators to BED is novel, the exact contributions and benefits of the author's approach are not clear and concise.