DeepGEM: Generalized Expectation-Maximization for Blind Inversion

1 Introduction . Physics-based inversion methods typically recover an unknown source from indirect measurements by assuming that the source and measurements are related via a known forward model [ 15 , 7 , 30 ] . For example , non-blind deconvolution algorithms often assume that a measured blurry image is related to its true sharp image via a known spatially-invariant blur kernel [ 9 ] ; and traditional seismic inversion methods assume that the spatially-varying velocity of the Earth ’ s interior is known a priori when solving for an earthquake ’ s hypocenter [ 30 ] . However , these “ known ” forward models are generally idealized and ignore intricacies of the systems that are either hard to capture or simply unknown . Inversion algorithms with forward model mismatch result in biased reconstructions . For instance , bias is regularly seen in non-blind deconvolution results , where reconstruction artifacts are often present due to the use of an incorrect blur kernel [ 9 ] . In order to reduce the effects of model mismatch , in this paper we tackle the problem of blind inversion : solving an inverse problem without knowledge of the underlying forward model parameters . When considering how learning can help , a natural first idea might be to learn a direct map from measurements to the desired source via supervised learning . However , such a “ model-free ” approach is generally not practical , due to the lack of available ground truth training data . For example , the blur kernel caused by handheld camera shake can not be reproduced to get a training set of 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) , . sharp-blurry pairs to train a deconvolution approach . In seismic tomography , synthetic earthquakes can not be placed densely throughout the interior of the Earth to measure the ground response as a function of hypocenter location . An alternative approach , which we adopt , is to develop unsupervised methods that treat the true forward model as something that is unobserved and must be inferred . An additional consideration is that we must solve for the true forward model from a single dataset , without knowledge of the true source that produces the measurements In this paper , we propose Deep Generalized Expectation-Maximization ( DeepGEM ) for solving blind inverse problems . Using the indirect measurements as input , DeepGEM jointly estimates the source and forward model that together produce the observed measurements . DeepGEM is a variational inference based framework that makes use of deep learning machinery to easily capture and optimize complex probabilistic distributions that can not be easily integrated in analytic ExpectationMaximization ( EM ) solutions . Our proposed framework is generic and can be applied to blind inversion problems described by differentiable forward models . In Section 4 we showcase the effectiveness of our DeepGEM approach on the challenging problem of blind seismic tomography , where we significantly outperform methods used in seismology . We also demonstrate the generality of DeepGEM by applying it to blind deconvolution in Section 5 . 2 Background and related work . The joint optimization of a forward model with source recovery is a very challenging ill-posed problem , leading to many possible solutions that are hard to disambiguate . For example , in blind deconvolution the blurry image observed can be equivalently explained by a sharp source image convolved with an extended kernel or a blurry source image convolved with a impulse kernel ; to prefer one solution over another , additional information , such as image priors , must be considered . Previous work on inversion in poorly characterized systems focused on limited contexts , such as spatially-invariant blind deconvolution [ 21 , 22 , 14 , 10 ] and CT with a simple rotational error [ 35 ] . These methods tend to be highly specialized to each application domain , and can not be easily generalized . In contrast , we propose a flexible model-based Bayesian framework that can be applied across multiple differentiable blind inversion problems . 2.1 Model-based Bayesian inversion . In model-based inversion , unobserved sources x and observed measurements y are related through a forward model : y = f ( x ) . When the parameters of the underlying model are unknown , one can parameterize the forward model as fθ ( x ) and then solve for the true model parameters θ∗ . A common approach is to solve a maximum a posteriori ( MAP ) objective : either MAPθ , x or MAPθ . MAPθ , x solves for the optimal point estimate of the pair { θ̂ , x̂ } that maximizes a joint objective : { θ̂ , x̂ } = arg max θ , x [ log p ( θ , x|y ) ] = arg max θ , x [ log p ( y|θ , x ) + log p ( θ ) + log p ( x ) ] . ( 1 ) In practice , formal probabilistic definitions of p ( θ ) and p ( x ) are often unknown and replaced with regularization terms ( e.g. , total variation , sparsity ) [ 22 , 21 ] . Although MAPθ , x provides a straightforward approach to solve for θ , the energy landscape of Eq . 1 is typically poorly behaved due to the ill-posed nature of the problem [ 22 ] . As a result , optimization is likely to get stuck in a ( bad ) local minimum . MAPθ attempts to smooth the energy landscape by reducing the number of parameters that must be optimized . This is done by solving for parameters of the forward model , θ , that perform best under the full volume of possible x interpretations : θ̂ = arg max θ [ log p ( θ|y ) ] = arg max θ [ log ( ∫ x p ( y|θ , x ) p ( x ) dx ) + log p ( θ ) ] . ( 2 ) Since this marginalization integral is often intractable , Expectation-Maximization ( EM ) algorithms have long been used for solving MAPθ efficiently [ 11 ] . EM is an iterative algorithm that alternates between : “ E ” -Step ) calculating the posterior of x conditioned on the current estimated forward model parameters θ ( t−1 ) ; and M-Step ) updating θ to maximize the expected value of the log likelihood : θ ( t ) = arg max θ [ Ex∼p ( x|y , θ ( t−1 ) ) [ log p ( y|θ , x ) ] + log p ( θ ) ] . ( 3 ) The advantage of MAPθ over MAPθ , x optimization for the blind deconvolution problem was described in [ 22 ] and its success demonstrated via EM optimization in [ 21 ] . However , it is important to note that evaluating the expectation in of Eq . ( 3 ) over complex distributions is often intractable . For instance , the authors of [ 21 ] were forced to restrict the posterior distribution to a Gaussian distribution . Alternatively , stochastic EM methods [ 5 , 8 ] bypasses the need to evaluate the expectation directly , approximating it by sampling the distribution . In this paper , we solve MAPθ using complex distributions parameterized by deep neural networks . Forward model parameterization Model-based inversion requires that the parametric form of fθ ( x ) is well matched with the true forward model , which is not always known . Alternatively , neural networks can be used to approximate the forward model , where θ represents the network weights . This setup is flexible , in that it can be used to approximate arbitrarily complex forward models , with the downside that it is often not interpretable when the parameters have no physical meaning . In this work , we develop and make use of interpretable , physically-motivated deep neural networks to parameterize fθ ( · ) for the problems of blind seismic tomography and blind deconvolution . 2.2 Deep variational distributions . We are interested in solving blind inverse problems using the EM algorithm to optimize the MAPθ objective . Optimizing in this fashion requires the use of the posterior distribution p ( x|y , θ ( t−1 ) ) in evaluating Eq . 3 . As inverse problems are often ill-posed , we expect that the posterior distribution of source x conditioned on the forward model parameters θ is likely to be complex and perhaps even multi-modal . Therefore , it is important to be able to parameterize a flexible family of distributions to best estimate this conditional posterior . Deep Probabilistic Imaging ( DPI ) [ 32 ] uses a normalizing flow-based generative model , Gφ ( · ) , to solve for the uncertainty of an unknown source x given a fixed forward model f ( x ) and measurements y. DPI solves the variational objective : φ̂ = argmin φ KL ( qφ ( x ) ||p ( x|y ) ) = argmin φ Ex∼qφ ( x ) [ − log p ( y|x ) − log p ( x ) + log qφ ( x ) ] , ( 4 ) where qφ ( x ) is the implicit distribution defined by Gφ ( z ) for z ∈ R|x| ∼ N ( 0 , 1 ) 1 , and log p ( y|x ) ∝ ||y − f ( x ) ||2 + c when there exists i.i.d . Gaussian noise on the measurements , y . After inference , qφ ( x ) can be efficiently sampled by evaluating Gφ ( z ) for a Gaussian sample z . The DPI variational objective is equivalent to the Variational Autoencoder [ 17 , 28 ] objective , except with a fixed decoder , f ( x ) . In practice , vanilla VAEs constrain the posterior to be a Gaussian 1|x| is the dimension of x distribution , relying on the reparameterization trick for tractable optimization . Alternatively , DPI uses flow-based networks to efficiently sample and directly compute qφ ( x ) [ 12 , 13 , 18 ] . DPI ’ s use of a flow-based network allows for complicated and multi-modal posterior distributions constrained only by the space of possible invertible network architectures . Our proposed DeepGEM approach utilizes similar tools to model flexible distributions over x , while simultaneously learning the forward model parameters θ . 3 Methods . We propose a deep variational EM approach ( DeepGEM ) that optimizes the MAPθ objective in Eq . 2 to recover the parameters of a forward model fθ ( x ) using only a single static set of measurements y . Once learned , the updated forward model can then be used to estimate the posterior distribution of the unknown source , x. DeepGEM iterates between two stages that are inspired by the standard EM algorithm : ( 1 ) an “ E ” -step that learns a variational distribution , qφ ( t ) ( x ) , to approximate the posterior distribution of x given the current forward model parameters θ ( t−1 ) , and ( 2 ) an M-step ( refer to Eq . 3 ) that solves for θ ( t ) that maximizes the expected value of the log likelihood function of θ , with respect to the posterior distribution qφ ( t ) ( x ) estimated in the prior “ E ” -step . Each step is alternated and solved to convergence . Since this is an EM algorithm ( up to the variational approximation ) , our method inherits properties of EM , including convergence to a local minimum of Eq . 2 . 3.1 “ Expectation ” step ( “ E ” -step ) . The goal of DeepGEM ’ s “ E ” -step is to solve for the posterior distribution p ( x|y , θ ( t−1 ) ) that facilitates optimizing Eq . 3 . Because this posterior distribution can be very complex , and even multi-modal , we propose to use a flexible variational approach to learn the parameters φ of an approximate posterior distribution qφ ( x ) . The variational distribution qφ ( x ) can then be used to evaluate Eq . 3 via efficient sampling . Using DPI we solve for a flexible variational distribution , qφ ( x ) that well approximates the posterior distribution p ( x|y , θ ( t−1 ) ) . DPI uses a normalizing flow network , Gφ ( z ) , with input z ∈ R|x| , where x = Gφ ( z ) ∼ qφ ( x ) when z ∼ N ( 0,1 ) . Normalizing flow networks allow for exact computation of the log-likelihood log qφ ( x ) needed to solve φ ( t ) = arg min φ KL ( qφ ( x ) ||p ( x|y , θ ( t−1 ) ) ) ≈ arg min φ 1 N N∑ n=1 [ − log p ( y|θ ( t−1 ) , xn ) − log p ( xn ) + log qφ ( xn ) ] for xn = Gφ ( zn ) , zn ∼ N ( 0,1 ) , ( 5 ) ( as derived from Eq . 4 ) for a batch size ofN where log p ( x ) is a prior on the source and log p ( y|x , θ ( t ) ) is the data likelihood . When assuming the measurements y experience i.i.d additive Gaussian noise with standard deviation σy , log p ( y|θ ( t ) , xn ) = 12σ2y ‖y − fθ ( t ) ( xn ) ) ‖ 2 + c .

Summary: This paper considers the problem of inverting a signal from a given forward model $y=f_\theta(x^*)+\epsilon$, while accounting model mismatch from true model $f_{\theta^*}$. The authors propose to solve this problem by first approximating the distribution of $x$ with $q_\phi(x) \sim G_\phi (z)$ using a deep model and carrying out an expectation over $x\sim q_\phi(x)$ of the log-likelihood $p(y|\theta,x)$. Model parameters $\theta$ are approximated via maximization over this expectation (EM framework). They call this framework DeepGEM.  The authors specifically apply this mathematical model to 1) blind seismic tomography, where the task is to estimate the velocity structure of earthquake waves. The velocity computation depends on travel time $T$ (this is the variable observed) parameterized by source location $x$ and receiver location $r$.  In this paper $T(x,s) = f_\theta(x,s) $ and $y \approx T(x,s=r)$. 2) Blind deconvolution with $y=x * k_\theta +\epsilon$.