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 proposes a general variational Expectation-Maximization framework aided by a normalizing flow generative network to solve inverse problems for which a forward physical model is known, but important model mismatch exists. The approach is evaluated on a blind seismic tomographic task and showed to quantitatively and qualitatively outperform a state-of-the art baseline through various experiments, that include ablation studies. The versatility of the framework is also illustrated by a small case-study on image blind deconvolution, where the authors report convincing qualitative results on two images.