Preconditioned training of normalizing ﬂows for variational inference in inverse problems
Abstract
Obtaining samples from the posterior distribution of inverse problems with expensive forward operators is challenging especially when the unknowns involve the strongly heterogeneous Earth. To meet these challenges, we propose a preconditioning scheme involving a conditional normalizing ﬂow (NF) capable of sampling from a low-ﬁdelity posterior distribution directly.
This conditional NF is used to speed up the training of the high-ﬁdelity objective involving minimization of the Kullback-Leibler divergence between the predicted and the desired high-ﬁdelity posterior density for indirect measurements at hand.
To minimize costs associated with the forward operator, we initialize the high-ﬁdelity NF with the weights of the pretrained low-ﬁdelity NF, which is trained beforehand on available model and data pairs.
Our numerical experiments, including a 2D toy and a seismic compressed sensing example, demonstrate that thanks to the preconditioning considerable speed-ups are achievable compared to training NFs from scratch.
1. Introduction
Our aim is to perform approximate Bayesian inference for inverse problems characterized by computationally expensive forward operators, F : X → Y , with a data likelihood, π like ( y | x ) :
y = F ( x ) + (cid:15) , (1) where x ∈ X is the unknown model,
y ∈ Y the observed data and (cid:15) ∼ N( 0 , σ 2 I ) the measurement noise.
Given a prior density, π prior ( x ) , variational inference (VI, Jordan et al., 1999) based on normalizing ﬂows (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-ﬁdelity , posterior density π post ( x | y ) (Liu and Wang, 2016; Kruse et al., 2019; Rizzuti et al., 2020; Siahkoohi et al., 2020; Sun and Bouman, 2020):
[Equation]
In the above expression, T θ : Z x → X denotes a NF with parameters θ and a Gaussian latent variable z ∈ 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 MonteCarlo averaging using mini-batches of z .
After training, samples from the approximated posterior, π θ ( x | y )
≈ π post ( x | y ) , can be drawn by evaluating T θ ( z ) for z ∼ π z ( z ) (Kruse et al., 2019).
It is important to note that Equation (2) trains a NF speciﬁc to the observed data y . While the above VI formulation in principle allows us to train a NF to generate samples from the posterior given a single observation 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, ∇ F (cid:62) .
As in multi-ﬁdelity 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, (cid:98) π y,x ( y , x ) —are preferable.
More speciﬁcally, 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 φ :
Y × X → Z y × Z x , with latent space
Z y × Z x :
[Equation]
where
[Equation]
Thanks to the block-triangular structure of G φ , samples of the approximated posterior, π φ ( x | y )
≈ π post ( x | y ) can be drawn by evaluating G − 1 φ x ( G φ y ( y ) , z ) for z ∼ π z ( z ) (Marzouk et al., 2016).
Unlike the objective in Equation (2), training G φ does not involve multiple evaluations of F and ∇ F (cid:62) , 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, y , x ∼ (cid:98) π y,x
( y , 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 simpliﬁed artiﬁcial 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 signiﬁcant diﬀerences between shallow (Figures 1a and 1b) and deeper (Figures 1c and 1d) sections. [Figure]
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. take maximum advantage of available samples from the joint distribution (cid:98) π y,x ( y , x ) , to pretrain G φ by minimizing Equation (3).
We only incur these costs once, by training this NF beforehand. As these samples typically come from a diﬀerent (neighboring) region, they are considered as low-ﬁdelity; 2. exploit the invertibility of G φ x ( y , · ) , which gives us access to a low-ﬁdelity posterior density, π φ ( x | y ) .
For a given y , this trained (conditional) prior can be used in Equation (2); 3.
initialize T θ with weights from the pretrained G − 1 φ x .
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.
2. Related work
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 eﬃcient 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 eﬃciency 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-ﬁdelity posterior distribution is used as a preconditioner.
This idea of multi-ﬁdelity 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 reﬁnement step (Hjelm et al., 2016; Krishnan et al., 2018;
Kim et al., 2018; Marino et al., 2018).
Although similar in spirit, our approach is diﬀerent 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.
3. Multi-ﬁdelity preconditioning scheme
For an observation y , we deﬁne a NF T φ x : Z x → X as
[Equation]
(4) where we obtain φ = { φ y , φ x } by training G φ through minimizing the objective function in Equation (3).
To train G φ , we use available low-ﬁdelity training pairs y , x ∼ (cid:98) π y,x
( y , 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 φ , it provides an expression for the posterior. We refer to this posterior as low-ﬁdelity 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):
[Equation]
To train the high-ﬁdelity NF given observed data y , we minimize the KL divergence between the predicted and the high-ﬁdelity posterior density, π post ( x | y )
(Liu and Wang, 2016; Kruse et al., 2019)
[Equation]
where the prior density of Equation (5) is used.
Notice that this minimization problem diﬀers from the one stated in Equation (2). Here, the optimization involves “ﬁne-tuning” the low-ﬁdelity network parameters φ x introduced in Equation (3). Moreover, this low-ﬁdelity 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-ﬁdelity 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-ﬁdelity posterior may be more informative (Yang and Soatto, 2018) than its unconditional counterpart because it is conditioned by the observed data y . In addition, our multi-ﬁdelity 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 ﬁdelity 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-ﬁdelity formulation.
4. Numerical experiments
We present two synthetic examples aimed at verifying the anticipated speed-up and increase in accuracy of the predicted posterior density via our multi-ﬁdelity preconditioning scheme.
The ﬁrst example is a two-dimensional problem where the posterior density can be accurately and cheaply sampled via MCMC.
The second example demonstrates the eﬀect 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-eﬃcient framework for training invertible networks in the Julia programming language.
4.1. 2D toy example
To illustrate, the advantages of working with our multi-ﬁdelity scheme, we consider the 2D
Rosenbrock distribution, π prior ( x ) ∝ exp (cid:16) − 12 x 21 − (cid:0) x 2 − x 21 (cid:1) 2 (cid:17) , plotted in Figure 2a.
Highﬁdelity data y ∈ R 2 are generated via y =
A x + (cid:15) , where (cid:15) ∼ N( 0 , 0 .
4 2 I ) and A ∈ R 2 × 2 is a forward operator.
To control the discrepancy between the low- and high-ﬁdelity samples, we set A equal to ¯ A / ρ ( ¯ A ) , where ρ ( · ) is the spectral radius of ¯
A = Γ + γ I , Γ ∈ R 2 × 2 has independent and normally distributed entries, and γ = 3 .
By choosing smaller values for γ , we make A more dissimilar to the identity matrix, therefore increasing the discrepancy between the low- and high-ﬁdelity posterior.
Figure 2b depicts the low- (purple) and high-ﬁdelity (red) data densities.
The dark star represents the unknown model.
Low-ﬁdelity data samples are generated with the identity as the forward operator. During the pretraining phase conducted beforehand, we minimize the objective function in Equation (3) for 25 epochs.
The pretrained low-ﬁdelity posterior is subsequently used to precondition the minimization of (6) given observed data y . The resulting low- and high-ﬁdelity estimates for the posterior as plotted in Figure 2c diﬀer signiﬁcantly.
In Figure (2d), the accuracy of the proposed method is veriﬁed by comparing the approximated high-ﬁdelity posterior density (orange contours) with the approximation (in green) obtained by minimizing the objective of Equation (2). The overlap between the orange contours and the green shaded area conﬁrms consistency between the two methods.
To assess the accuracy of the estimated densities themselves, [Figure] we also include samples from the posterior (dark circles) obtained via stochastic gradient Langevin dynamics (Welling and Teh, 2011), an MCMC sampling technique.
As expected, the estimated posterior densities with and without the preconditioning scheme are in agreement with the MCMC samples.
Finally, to illustrate the performance our multi-ﬁdelity scheme, we consider the convergence plot in Figure 2e where the objective values of Equations (2) and (6) are compared. As explained in Appendix A, the values of the objective functions correspond to the KL divergence (plus a constant) between the posterior given by Equation (2) and the posterior distribution obtained by our multi-ﬁdelity approach (Equation (6)). As expected, the multiﬁdelity objective converges much faster because of the “warm start”. In addition, the updates of T φ x via Equation (6) succeed in bringing down the KL divergence within only ﬁve epochs (see orange curve), whereas it takes 25 epochs via the objective in Equation (2) to reach approximately the same KL divergence.
This pattern holds for smaller values of γ too as indicated in Table 1.
According to Table 1, the improvements by our multi-ﬁdelity method become more pronounced if we decrease the γ . This behavior is to be expected since the samples used for pretraining are more and more out of distribution in that case. We refer to Appendix C for additional ﬁgures for diﬀerent values of γ .
[Table] 4.2. Seismic compressed sensing example
This experiment is designed to show challenges with geophysical inverse problems due to the
Earth’s strong heterogeneity.
We consider the inversion of noisy indirect measurements of image patches x ∈ R 256 × 256 sampled from deeper parts of the Parihaka seismic dataset.
The observed measurements are given by y
=
Ax + (cid:15) where (cid:15) ∼ N( 0 , 0 . 2 2 I ) .
For simplicity, we chose A = M T M with M a compressing sensing matrix with 66 .
66 % subsampling. The measurement vector y corresponds to a pseudo-recovered model contaminated with noise.
To mimic a realistic situation in practice, we change the likelihood distribution by reducing the standard deviation of the noise to 0 .
01 in combination with using image patches sampled from the shallow part of the Parihaka dataset. As we have seen in Figure 1, these patches diﬀer in texture. Given pairs y , x ∼ (cid:98) π y,x ( y , x ) , we pretrain our network by minimizing
Equation (3).
Figures 3a and 3b contain a pair not used during pretraining. Estimates for the conditional mean and standard deviation obtained by drawing 1000 samples from the pretrained conditional NF for the noisy indirect measurement (Figure 3b) are included in Figures 3c and 3d.
Both estimates exhibit the expected behavior because the examples in
Figure 3a and 3b are within the distribution.
As anticipated, this observation no longer holds if we apply this pretrained network to indirect data depicted in Figure 3f, which is sampled from the deeper part. However, these results are signiﬁcantly improved when the pretrained network is ﬁne-tuned by minimizing Equation (6).
After ﬁne tuning, the ﬁne details in the image are recovered (compare Figures 3g and 3h).
This improvement is conﬁrmed by the relative errors plotted in Figures 3i and 3j, as well as by the reduced standard deviation (compare Figures 3k and 3l).
5. Conclusions
Inverse problems in ﬁelds 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 ﬂows 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 [Figure] distribution.
We use these samples to train a normalizing ﬂow via likelihood maximization leveraging the normalizing property. We use this pretrained low-ﬁdelity 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 ﬂow from scratch. Appendix A. Mathematical derivations
Let f : Z → X be a bijective transformation that maps a random variable z ∼ π z ( z ) to x ∼ π x ( x ) .
We can write the change of variable formula (Villani, 2009) that relates probability density functions π z and π x in the following manner:
[Equation] This relation serves as the basis for the objective functions used throughout this paper. A.1. Derivation of Equation (2) In Equation 2, we train a bijective transformation, denoted by T θ : Z x → X , that maps the latent distribution π z x
( z x ) to the high-ﬁdelity posterior density π post ( x | y ) .
We optimize the parameters of T θ by minimizing the KL divergence between the push-forward density (Bogachev, V.I., 2006), denoted by π θ ( · | y ) :
= T (cid:93)
π z x , and the posterior density:
[Equation]
In the above expression, we can rewrite the expectation with respect to π θ ( x | y ) as the expectation with respect to the latent distribution, followed by a mapping via T θ —i.e.,
[Equation]
The last term in the expectation above can be further simpliﬁed via the change of variable formula in Equation (7).
If x = T θ ( z x ) , then:
[Equation]
The last equality in Equation (10) holds due to the invertibility of T θ and the diﬀerentiability of its inverse (inverse function theorem).
By combining Equations (9) and (10), we arrive at the following objective function for training T θ :
[Equation]
Finally, by ignoring the log π z x ( z x ) term, which is constant with respect to θ , using Bayes’ rule, and inserting our data likelihood model from Equation (1), we derive Equation (2): min θ
[Equation]
Next, based on this equation, we derive the objective function used in the pretraining phase.
A.2. Derivation of Equation (3)
The derivation of objective in Equation (3) follows directly from the change of variable formula in Equation 7, applied to a bijective map, G − 1 φ : Z y ×
Z x → Y × X , where Z y and
Z y are Gaussian latent spaces.
That is to say:
[Equation]
Given (low-ﬁdelity) training pairs, y , x ∼ (cid:98) π y,x ( y , x ) , the maximum likelihood estimate of φ is obtained via the following objective:
[Equation]
that is the objective function in Equation (3).
The NF trained via the objective function, given samples from the latent distribution, draws samples from the low-ﬁdelity joint distribution, (cid:98) π y,x .
By construction, G φ is a block-triangular map—i.e.
[Equation]
Kruse et al.
(2019) showed that after solving the optimization problem in Equation (3), G φ approximates the well-known triangular Knothe-Rosenblat map (Santambrogio, 2015).
As shown in Marzouk et al.
(2016), the triangular structure and G φ ’s invertibility yields the following property,
[Equation]
where (cid:98) π post denotes the low-ﬁdelity posterior probability density function.
The expression above means we can get access to low-ﬁdelity posterior distribution samples by simply evaluating G − 1 φ x ( G φ y ( y ) , z x ) for z x ∼ π z x ( z x ) for a given observed data y .
Appendix B. Training details and network architectures
For our network architecture, we adapt the recursive coupling blocks proposed by Kruse et al.
(2019), which use invertible coupling layers from Dinh et al.
in a hierarchical way.
In other words, we recursively divide the incoming state variables and apply an aﬃne coupling layer.
The ﬁnal architecture is obtained by composing several of these hierarchical coupling blocks.
The hierarchical structure leads to dense triangular Jacobian, which is essential in representation power of NFs (Kruse et al., 2019).
For all examples in this paper, we use the hierarchical coupling blocks as described in
Kruse et al.
The aﬃne coupling layers within each hierarchal block contain a residual block as described in He et al. Each residual block has the following dimensions:input, 128 hidden, and 64 output channels, except for the ﬁrst and last coupling layer where we have 4 input and output channels, respectively.
We use the Wavelet transform and its transpose before feeding seismic images into the network and after the last layer of the NFs.
Below, we describe the network architectures and training details regarding the two numerical experiments described in the paper.
Throughout the experiments, we use the Adam optimization algorithm (Kingma and Ba, 2014).
B.1. 2D toy example
We use 8 hierarchal coupling blocks, as described above for both G φ x and G φ y (Equation (3)).
As a result, due to our proposed method in Equation (4), we choose the same architecture for T φ x (Equation (2)).
For pretraining G θ according to Equation (3), we use 5000 low-ﬁdelity joint training pairs, y , x ∼ (cid:98) π y,x
( y , x ) .
We minimize Equation (3) for 25 epochs with a batch size of 64 and a (starting) learning rate of 0 .
We decrease the learning rate each epoch by a factor of 0 . 9 .
For the preconditioned step—i.e., solving Equation 6, we use 1000 latent training samples.
We train for 5 epochs with a batch size of 64 and a learning rate 0 .
Finally, as a comparison, we solve the objective in Equation 6 for a randomly initialized NF with the same 1000 latent training samples for 25 epochs.
We decrease the learning rate each epoch by 0 . 9 .
B.2. Seismic compressed sensing
We use 12 hierarchal coupling blocks, as described above for both G φ x , G φ y , and we use the same architecture for T φ x as G φ x .
For pretraining G θ according to Equation (3), we use 5282 low-ﬁdelity joint training pairs, y , x ∼ (cid:98) π y,x
( y , x ) .
We minimize Equation (3) for 50 epochs with a batch size of 4 and a starting learning rate of 0 .
Once again, we decrease the learning rate each epoch by 0 . 9 .
For the preconditioned step—i.e., solving Equation 6, we use 1000 latent training samples.
We train for 10 epochs with a batch size 4 and a learning rate of 0 . 0001 , where we decay the step by 0 .
9 in every 5 th epoch.
Appendix C. 2D toy example—more results
Here we show the eﬀect γ on our proposed method in the 2D toy experiment.
By choosing smaller values for γ , we make ¯ A / ρ ( ¯ A ) with ¯ A =
Γ + γ I less close to the identity matrix, hence enhancing the discrepancy between the low- and high-ﬁdelity posterior.
The ﬁrst row of Figure 4 shows the low- (purple) and high-ﬁdelity (red) data densities for decreasing values of γ from 2 down to 0 . The second row depicts the predicted posterior densities via the preconditioned scheme (orange contours) and the low-ﬁdelity posterior in green along with MCMC samples (dark circles). The third row compares the preconditioned posterior densities to samples obtained via the low-ﬁdelity pretrained NF—i.e., Equation (3). Finally, the last row shows the objective function values during training with (orange) and without (green) preconditioning.
We observe that by decreasing γ from 2 to 0 , the low-ﬁdelity posterior approximations become worse.
As a result, the objective function for the preconditioned approach (orange) at the beginning start from a higher value, indicating more mismatch between low- and high-ﬁdelity posterior densities.
Finally, our preconditioning method consistently improves upon low-ﬁdelity posterior by training for 5 epochs.
Appendix D. Seismic compressed sensing—more results
Here we show more examples to verify the pretraining phase obtained via solving the objective in Equation (2).
Each row in Figure 5 corresponds to a diﬀerent testing image. The ﬁrst column shows the diﬀerent true seismic images used to create low-ﬁdelity compressive sensing data, depicted in the second column.
The third and last columns correspond to the conditional mean and pointwise STD estimates, respectively.
Clearly, the pretrained network is able to successfully recover the true image, and consistently indicates more uncertainty in areas with low-amplitude events.