Normalizing flow with adaptive surrogate (NoFAS)

LINFA is designed to accommodate black-box models \(\boldsymbol{f}: \boldsymbol{\mathcal{Z}} \to \boldsymbol{\mathcal{X}}\) between the random inputs \(\boldsymbol{z} = (z_1, z_2, \cdots, z_d)^T \in \boldsymbol{\mathcal{Z}}\) and the outputs \((x_1, x_2,\cdots,x_m)^T \in \boldsymbol{\mathcal{X}}\), and assumes \(n\) observations \(\boldsymbol x = \{\boldsymbol x_i\}_{i=1}^n \subset \boldsymbol{\mathcal{X}}\) to be available.

Our goal is to infer \(\boldsymbol z\) and to quantify its uncertainty given \(\boldsymbol{x}\). We employ a variational Bayesian paradigm and sample from the posterior distribution \(p(\boldsymbol z\vert \boldsymbol x)\propto \ell_{\boldsymbol z}(\boldsymbol x,\boldsymbol{f})\,p(\boldsymbol z)\), with prior \(p(\boldsymbol z)\) via normalizing flows.

This requires the evaluation of the gradient of the ELBO (1) with respect to the NF parameters \(\boldsymbol{\lambda}\), replacing \(p(\boldsymbol x, \boldsymbol z_K)\) with \(p(\boldsymbol x\vert\boldsymbol z_K)\,p(\boldsymbol z)=\ell_{\boldsymbol z_K}(\boldsymbol{x},\boldsymbol{f})\,p(\boldsymbol z)\), and approximating the expectations with their MC estimates.

However, the likelihood function needs to be evaluated at every MC realization, which can be costly if the model \(\boldsymbol{f}(\boldsymbol{z})\) is computationally expensive. In addition, automatic differentiation through a legacy (e.g. physics-based) solver may be an impractical, time-consuming, or require the development of an adjoint solver.

Our solution is to replace the model \(\boldsymbol{f}\) with a computationally inexpensive surrogate \(\widehat{\boldsymbol{f}}: \boldsymbol{\mathcal{Z}} \times \boldsymbol{\mathcal{W}} \to \boldsymbol{\mathcal{X}}\) parameterized by the weigths \(\boldsymbol{w} \in \boldsymbol{\mathcal{W}}\), whose derivatives can be obtained at a relatively low computational cost, but intrinsic bias in the selected surrogate formulation, a limited number of training examples, and locally optimal \(\boldsymbol{w}\) can compromise the accuracy of \(\widehat{\boldsymbol{f}}\).

To resolve these issues, LINFA implements NoFAS, which updates the surrogate model adaptively by smartly weighting the samples of \(\boldsymbol{z}\) from NF thanks to a memory-aware loss function. Once a newly updated surrogate is obtained, the likelihood function is updated, leading to a new posterior distribution that will be approximated by VI-NF, producing, in turn, new samples for the next surrogate model update, and so on. Additional details can be found in [WLS22].