Preconditioned training of normalizing flows for variational inference in inverse problems

Our aim is to perform approximate Bayesian inference for inverse problems characterized by computationally expensive forward operators, $F:\mathcal{X}\rightarrow\mathcal{Y}$, with a data likelihood, $\pi_{\text{like}}(\bm{y}\mid\bm{x})$:

where $\bm{x}\in\mathcal{X}$ is the unknown model, $\bm{y}\in\mathcal{Y}$ the observed data, and $\bm{\epsilon}\sim\mathrm{N}(\bm{0},\sigma^{2}\bm{I})$ the measurement noise. Given a prior density, $\pi_{\text{prior}}(\bm{x})$, variational inference (VI, Jordan et al., 1999) based on normalizing flows (NFs, Rezende and Mohamed, 2015) can be used where the Kullback-Leibler (KL) divergence is minimized between the predicted and the target—i.e., high-fidelity, posterior density $\pi_{\text{post}}(\bm{x}\mid\bm{y})$ (Liu and Wang, 2016; Kruse et al., 2019; Rizzuti et al., 2020; Siahkoohi et al., 2020; Sun and Bouman, 2020):

In the above expression, $T_{\theta}:\mathcal{Z}_{x}\to\mathcal{X}$ denotes a NF with parameters $\bm{\theta}$ and a Gaussian latent variable $\bm{z}\in\mathcal{Z}_{x}$. The above objective consists of the data likelihood term, regularization on the output of the NF, and a log-determinant term that is related to the entropy of the NF output. The last term is necessarily to prevent the output of the NF from collapsing on the maximum a posteriori estimate. For details regarding the derivation of the objective in Equation (2), we refer to Appendix A. During training, we replace the expectation by Monte-Carlo averaging using mini-batches of $\bm{z}$. After training, samples from the approximated posterior, $\pi_{\theta}(\bm{x}\mid\bm{y})\approx\pi_{\text{post}}(\bm{x}\mid\bm{y})$, can be drawn by evaluating $T_{\theta}(\bm{z})$ for $\bm{z}\sim\pi_{z}(\bm{z})$ (Kruse et al., 2019). It is important to note that Equation (2) trains a NF specific to the observed data $\bm{y}$. While the above VI formulation in principle allows us to train a NF to generate samples from the posterior given a single observation $\bm{y}$, this variational estimate requires access to a prior density, and the training calls for repeated evaluations of the forward operator, $F$, as well as the adjoint of its Jacobian, $\nabla{F}^{\top}$. As in multi-fidelity Markov chain Monte Carlo (MCMC) sampling (Peherstorfer and Marzouk, 2019), the costs associated with the forward operator may become prohibitive even though VI-based methods are known to have computational advantages over MCMC (Blei et al., 2017).

Aside from the above computational considerations, reliance on having access to a prior may be problematic especially when dealing with images of the Earth’s subsurface, which are the result of complex geological processes that do not lend themselves to be easily captured by hand-crafted priors. Under these circumstances, data-driven priors—or even better data-driven posteriors obtained by training over model and data pairs sampled from the joint distribution, $\widehat{\pi}_{y,x}(\bm{y},\bm{x})$—are preferable. More specifically, we follow Kruse et al. (2019), Kovachki et al. (2020), and Baptista et al. (2020), and formulate the objective function in terms of a block-triangular conditional NF, $G_{\phi}:\mathcal{Y}\times\mathcal{X}\to\mathcal{Z}_{y}\times\mathcal{Z}_{x}$, with latent space $\mathcal{Z}_{y}\times\mathcal{Z}_{x}$:

Thanks to the block-triangular structure of $G_{\phi}$, samples of the approximated posterior, $\pi_{\phi}(\bm{x}\mid\bm{y})\approx\pi_{\text{post}}(\bm{x}\mid\bm{y})$ can be drawn by evaluating $G_{\phi_{x}}^{-1}(G_{\phi_{y}}(\bm{y}),\bm{z})$ for $\bm{z}\sim\pi_{z}(\bm{z})$ (Marzouk et al., 2016). Unlike the objective in Equation (2), training $G_{\phi}$ does not involve multiple evaluations of $F$ and $\nabla F^{\top}$, nor does it require specifying a prior density. However, its success during inference heavily relies on having access to training pairs from the joint distribution, $\bm{y},\bm{x}\sim\widehat{\pi}_{y,x}(\bm{y},\bm{x})$. Unfortunately, unlike medical imaging, where data is abundant and variability among patients is relatively limited, samples from the joint distribution are unavailable in geophysical applications. Attempts have been made to address this lack of training pairs including the generation of simplified artificial geological models (Mosser et al., 2019), but these approaches cannot capture the true heterogeneity exhibited by the Earth’s subsurface. This is illustrated in Figure 1, which shows several true seismic image patches drawn from the Parihaka dataset. Even though samples are drawn from a single data set, they illustrate significant differences between shallow (Figures 1a and 1b) and deeper (Figures 1c and 1d) sections.

To meet the challenges of computational cost, heterogeneity and lack of access to training pairs, we propose a preconditioning scheme where the two described VI methods are combined to:

1. 1.

   take maximum advantage of available samples from the joint distribution $\widehat{\pi}_{y,x}(\bm{y},\bm{x})$, to pretrain $G_{\phi}$ by minimizing Equation (3). We only incur these costs once, by training this NF beforehand. As these samples typically come from a different (neighboring) region, they are considered as low-fidelity;

2. 2.

   exploit the invertibility of $G_{\phi_{x}}(\bm{y},\,\cdot\,)$, which gives us access to a low-fidelity posterior density, $\pi_{\phi}(\bm{x}\mid\bm{y})$. For a given $\bm{y}$, this trained (conditional) prior can be used in Equation (2);

3. 3.

   initialize $T_{\theta}$ with weights from the pretrained $G_{\phi_{x}}^{-1}$. This initialization can be considered as an instance of transfer learning (Yosinski et al., 2014), and we expect a considerable speed-up when solving Equation (2). This is important since it involves inverting $F$, which is computationally expensive.

In the context of variational inference for inverse problems with expensive forward operators, Herrmann et al. (2019) train a generative model to sample from the posterior distribution, given indirect measurements of the unknown model. This approach is based on an Expectation Maximization technique, which infers the latent representation directly instead of using an inference encoding model. While that approach allows for inclusion of hand-crafted priors, capturing the posterior is not fully developed. Like Kovachki et al. (2020), we also use a block-triangular map between the joint model and data distribution and their respective latent spaces to train a network to generate samples from the conditional posterior. By imposing an additional monotonicity constraint, these authors train a generative adversarial network (GAN, Goodfellow et al., 2014) to directly sample from the posterior distribution. To allow for scalability to large scale problems, we work with NFs instead, because they allow for more memory efficient training (Leemput et al., 2019; Putzky and Welling, 2019; Peters et al., 2020; Peters and Haber, 2020). Our contribution essentially corresponds to a reformulation of Parno and Marzouk (2018) and Peherstorfer and Marzouk (2019). In that work, transport-based maps are used as non-Gaussian proposal distributions during MCMC sampling. As part of the MCMC, this transport map is then tuned to match the target density, which improves the efficiency of the sampling. Peherstorfer and Marzouk (2019) extend this approach by proposing a preconditioned MCMC sampling technique where a transport-map trained to sample from a low-fidelity posterior distribution is used as a preconditioner. This idea of multi-fidelity preconditioned MCMC inspired our work where we setup a VI objective instead. We argue that this formulation can be faster and may be easier to scale to large-scale Bayesian inference problems (Blei et al., 2017).

Finally, there is a conceptional connection between our work and previous contributions on amortized variational inference (Gershman and Goodman, 2014), including an iterative refinement step (Hjelm et al., 2016; Krishnan et al., 2018; Kim et al., 2018; Marino et al., 2018). Although similar in spirit, our approach is different from these attempts because we adapt the weights of our conditional generative model to account for the inference errors instead of correcting the inaccurate latent representation of the out-of-distribution data.

For an observation $\bm{y}$, we define a NF $T_{\phi_{x}}:\mathcal{Z}_{x}\to\mathcal{X}$ as

where we obtain $\bm{\phi}=\left\{\bm{\phi}_{y},\bm{\phi}_{x}\right\}$ by training $G_{\phi}$ through minimizing the objective function in Equation (3). To train $G_{\phi}$, we use available low-fidelity training pairs $\bm{y},\bm{x}\sim\widehat{\pi}_{y,x}(\bm{y},\bm{x})$. We perform this training phase beforehand, similar to the pretraining phase during transfer learning (Yosinski et al., 2014). Thanks to the invertibility of $G_{\phi}$, it provides an expression for the posterior. We refer to this posterior as low-fidelity because the network is trained with often scarce and out-of-distribution training pairs. Because the Earth’s heterogeneity does not lend itself to be easily captured by hand-crafted priors, we argue that this NF can still serve as a (conditional) prior in Equation (2):

To train the high-fidelity NF given observed data $\bm{y}$, we minimize the KL divergence between the predicted and the high-fidelity posterior density, $\pi_{\text{post}}(\bm{x}\mid\bm{y})$ (Liu and Wang, 2016; Kruse et al., 2019)

where the prior density of Equation (5) is used. Notice that this minimization problem differs from the one stated in Equation (2). Here, the optimization involves “fine-tuning” the low-fidelity network parameters $\phi_{x}$ introduced in Equation (3). Moreover, this low-fidelity network is also used as a prior. While other choices exist for the latter—e.g., it could be replaced or combined with a hand-crafted prior in the form of constraints (Peters et al., 2019) or by a separately trained data-driven prior (Mosser et al., 2019), using the low-fidelity posterior as a prior (cf. Equation (5)) has certain distinct advantages. First, it removes the need for training a separate data-driven prior model. Second, use of the low-fidelity posterior may be more informative (Yang and Soatto, 2018) than its unconditional counterpart because it is conditioned by the observed data $\bm{y}$. In addition, our multi-fidelity approach has strong connections with online variational Bayes (Zeno et al., 2018) where data arrives sequentially and previous posterior approximates are used as priors for subsequent approximations.

In summary, the problem in Equation (6) can be interpreted as an instance of transfer learning (Yosinski et al., 2014) for conditional NFs. This formulation is particularly useful for inverse problems with expensive forward operators, where access to high fidelity training samples, i.e. samples from the target distribution, is limited. In the next section, we present two numerical experiments designed to show the speed-up and accuracy gained with our proposed multi-fidelity formulation.

We present two synthetic examples aimed at verifying the anticipated speed-up and increase in accuracy of the predicted posterior density via our multi-fidelity preconditioning scheme. The first example is a two-dimensional problem where the posterior density can be accurately and cheaply sampled via MCMC. The second example demonstrates the effect of the preconditioning scheme in a seismic compressed sensing (Candes et al., 2006; Donoho, 2006) problem. Details regarding training hyperparameters and the NF architectures are included in Appendix B. Code to reproduce our results are made available on GitHub. Our implementation relies on InvertibleNetworks.jl (Witte et al., 2020), a recently-developed memory-efficient framework for training invertible networks in the Julia programming language.

Inverse problems in fields such as seismology are challenging for several reasons. The forward operators are complex and expensive to evaluate numerically while the Earth is highly heterogeneous. To handle this situation and to quantify uncertainty, we propose a preconditioned scheme for training normalizing flows for Bayesian inference. The proposed scheme is designed to take full advantage of having access to training pairs drawn from a joint distribution, which for the reasons stated above is close but not equal to the actual joint distribution. We use these samples to train a normalizing flow via likelihood maximization leveraging the normalizing property. We use this pretrained low-fidelity estimate for the posterior as a prior and preconditioner for the actual variational inference on the observed data, which minimizes the Kullback-Leibler divergence between the predicted and the desired posterior density. By means of a series of examples, we demonstrate that our preconditioned scheme leads to considerable speed-ups compared to training a normalizing flow from scratch.