Learned imaging with constraints and uncertainty quantification

We outline new approaches to incorporate ideas from convolutional networks into wave-based least-squares imaging. The aim is to combine hand-crafted constraints with deep convolutional networks allowing us to directly train a network capable of generating samples from the posterior. The main contributions include combination of weak deep priors with hard handcrafted constraints and a possible new way to sample the posterior. The seismic imaging problem In least-squares imaging, we are interested in inverting the following inconsistent mildly ill-conditioned linear inverse problem:


The seismic imaging problem
In least-squares imaging, we are interested in inverting the following inconsistent mildly ill-conditioned linear inverse problem: In this expression, the unknown vector x represents the image, y i , i = 1, . . . , N the observed data from N experiments and A i the discretized linearized forward operator for the ith source experiment. Despite being overdetermined, the above least-squares imaging problem is challenging. The linear systems A i are large, expensive to evaluate, and inconsistent because of noise and/or linearization errors.
As in many inverse problems, solutions of problem 1 benefit from adding prior information in the form of penalties or preferentially in the form of constraints, yielding with C representing a single or multiple (convex) constraint set(s). While this approach offers flexibility to include multiple handcrafted constraints, several key issues remain, namely; (i) we can not afford to work with all N experiments when computing gradients for the above data-misfit objective; (ii) constrained optimization problems converge slowly; (iii) handcrafted priors may not capture complexities of natural images; (iv) it is non-trivial to obtain UQ information.

Stochastic linearized Bregman
To meet the computational challenges that come with solving problem 2 for non-differentiable structure promoting constraints, such as the 1 -norm, we solve problem 2 with Bregman iterations for a batch size of one. The kth iteration readsx with A k the adjoint of A k and P C (x) = argmin being the projection onto the (convex) set and t k = A k x − y k 2 2 / A k (A k x − y k ) 2 2 the dynamic steplength. Contrary to the Iterative Shrinkage Thresholding Algorithm (ISTA), we iterate on the dual variablex. Moreover, to handle more general situations and to ensure we are for every iteration feasible (= in the constraint set) we replace sparsity-promoting thresholding with projections that ensure that each model iterate remains in the constraint set. As reported in Witte et al. [2019], iterations 3 are known to converge fast for pairs {y k , A k } that are randomly drawn from iteration to iteration. As such, Equation 3 can be interpreted as stochastic gradients on the dual variable.
To illustrate the improved convergence of Equation 3, we consider two linearized imaging examples. The first involves sparsity-promoting least-squares migration with curvelet domain thresholding, which solves where C is the curvelet transform matrix and λ the threshold level. The results of 50 thresholded Bregman iterations (we stop early do not overfit) are included in Figure 1. This example clearly shows the advantage of the redrawing pairs {y i , A i }. Similarly, we also improve the convergence when imposing the TV-norm [Esser et al., 2018, Peters et al. [2018] at each Bregman iteration. For this specific type of image, TV-norm minimization gives superior results compared to curvelets (cf. Figure 1). We also assumed knowledge on the size of the TV-ball τ , which we imposed as a constraint when solving

Stochastic deep prior with constraints
Handcrafted priors, encoded in the constraint set C, in combination with stochastic optimization, where we randomly draw a different source experiment for each iteration of Equation 3, have proven to be scalable allowing us to create high-fidelity images by only working with random subsets of the data. While encouraging, this approach relies on handcrafted priors encoded in the constraint set C. Motivated by recent successes in machine learning and deep convolutional networks (CNNs) in particular, we follow Van Veen et al. [2018] propose to incorporate CNNs as deep priors on the model. Compared to handcrafted priors, deep priors defined by CNNs are less biased since they only require the model to be in the range of the CNN, which includes natural images and excludes images with unnatural noise. In its most basic form, this involves solving problems of the following type [Van Veen et al., 2018]: In this expression, g(z, w) is a deep CNN parameterized by unknown weights w and z ∼ N(0, 1) is a fixed random vector in the latent space. In this formulation, we replaced the unknown model by a neural net. This makes this formulation suitable for situations where we do not have access to data-image training pairs but where we are looking for natural images that are in the range of the generator. In recent work by Van Veen et al. [2018], it is shown that solving problem 7 can lead to good estimates for x via the generator g(z, w) where w is the minimizer of problem 7 highly suitable for situations where we only have access to data. In this approach, the parameterization of the network by w for a fixed z plays the role of a non-linear redundant transform. While using neural nets as strong constraints-i.e. problem 7 is equivalent to may offer certain advantages, there are no guarantees that the model iterates remain physically feasible, which is a prerequisite if we want to solve non-linear imaging problems that include physical parameters [Esser et al., 2018, Peters et al. [2018]. Unless we pre-train the network, early iterations solving problems 7 or 8 will be unfeasible. Moreover, as mentioned by Van Veen et al. [2018], results from solving inverse problems with deep priors may benefit from additional types of regularization. We accomplish this by combining hard handcrafted constraints with a weak constraint for the deep prior resulting in a reformulation of the program 8 into minimize x∈C, w In this expression, the deep prior appears as a penalty term weighted by the trade-off parameter λ > 0. In this weak formulation, x is a slack variable, which by virtue of the hard constraint will be feasible throughout the iterations.
The above formulation offers flexibility to impose constraints on the model that can be relaxed during the iterations as the network is gradually "trained". We can do this by either relaxing the constraint set (eg. by increasing the size of the TV-norm ball) or by increasing the trade-off parameter λ.

Learning posterior statistics through inversion
While the combination of hard and soft constraints offers flexibility to get the best of handcrafted and deep priors, we have not exploited the fact that generative models are when trained properly capable of generating samples from a distribution. After training, samples from the generator are produced by via nonlinear mappings of random realizations of the latent variable z ∼ p z (z) to g(z, w * ).
Unfortunately, in inverse problems, it is often difficult, if not impossible, to have access to training pairs {x i , y i } to obtain estimates for w, as outlined in recent work by Adler and Öktem [2018] on deep Bayesian inversion. These authors argue that information on the posterior distribution can be obtained by a direct method, which estimates different moments (e.g. mean over standard deviation) of the posterior distribution, or by estimating the posterior during training by minimizing the Wasserstein p = 1 distance between the statistics of the generator and the posterior. The latter approach leads to a supervised training method, wich involves a Generative Adversarial Network, training pairs {x i , y i }, and random draws of the latent space variable z ∼ p z (z).
Since we do not have access to training pairs, we instead propose to solve problem 9 for different pairs {x, y i }, where x represents a single Earth model within the survey area and y i data obtained from multiple source experiments with that same area. We solve this problem by taking a single (expensive) Bregman step followed by multiple (cheap) neural network updates during which we redraw the latent variable z ∼ p z (z). Strictly speaking this approach, summarized in Algorithm 1, corresponds to solving an inverse problem but because we have multiplicity in the data, the argument can be made we are at the same time training the generator to produce random samples from the posterior via g(z, w * ) and z ∼ p z (z) where the w * is found by carrying out the iterations in Algorithm 1.

Example
To test our speculation that solving problem 10 with stochastic gradient descent on the dual not only yield in image but also reaps information on the posterior, we conduct an experiment where the source experiments contain incoherent noise and coherent linearization errors, e = (F k (m + δm) − F k (m) − ∇F k (m)δm), where A k = ∇F k is the Jacobian and F k (m) is the nonlinear forward operator. The results of this experiment are included in Figure 2 from which we make the following observations. First, as expected the models generated from g(z, w * ) are smoother than the primal Bregman variable. Second, there are clearly variations amongst the different g(z, w * )'s and these variations average out in the mean, which has fewer imaging artifacts.
Because we were able to train the g(z, w) as a "biproduct" of the inversion, we are able to compute statistical information from the trained generator that may give us information about uncertainty. In Figure 3, we Algorithm 1 Stochastic weak deep prior imaging with constraints. input: Latent space distribution, p z (z) N Pairs of observed data and discretized linearized forward operators, Trade-off parameter, λ > 0 Projection operator, P C (·) Maximum training steps, M Inner loop generator optimization steps, T Generator learning rate, η initialization Randomly initialize generator parameters, w Initialize primal variable (unknown vector) with zeros Initialize dual variable with zeros end for end for output: w included a plot of the pointwise standard deviation and two examples of sample "prior" (before training) and "posterior" distribution functions. As expected, the pointwise standard deviations show a reasonable sharpening of the probabilities before and after training through inversion. We also argue that the areas of high pointwise standard deviation coincide with regions that are difficult to image because of the linearization error and noise.

Discussion and Conclusions
Solving large-scale inverse problems with expensive to evaluate forward models is challenging by itself. For this reason, the problem of obtaining uncertainty information remains elusive except for certain special situations [see e.g. Fang et al., 2018]. While we do not yet offer formal proof, we argue that the presented method, which combines stochastic optimization on the dual variable with on-the-fly estimation of the generator's network weights, reaps information on the posterior distribution leveraging multiplicity in the data and the fact that this data is acquired over one and the same Earth model. Our preliminary results seem consistent with a behavior to be expected from a posterior distribution. Figure 2: Imaging according to Algorithm 1. Top left: the Bregman primal variable x * obtained after 401 iterations. Top right: the mean of g(z, w * ) obtained by generating 2000 random realizations of z ∼ p z (z) and averaging the corresponding g(z, w * )'s. Middle row: two examples of generated images from g(z, w * ) for different z's. Bottom row: the differences between images in the middle row with another generated image. Figure 3: Statistics of imaging according to Algorithm 1. Top row: the pointwise standard deviation among samples generated by evaluating g(z, w * ) over 2000 random realizations of z ∼ p z (z). Bottom row: sample "prior" (before training) and "posterior" distribution functions for two points in the model.