DIFFERENTIABLE GAUSSIANIZATION LAYERS FOR INVERSE PROBLEMS REGULARIZED BY DEEP GENERA-TIVE MODELS

Abstract

Deep generative models such as GANs, normalizing flows, and diffusion models are powerful regularizers for inverse problems. They exhibit great potential for helping reduce ill-posedness and attain high-quality results. However, the latent tensors of such deep generative models can fall out of the desired high-dimensional standard Gaussian distribution during inversion, particularly in the presence of data noise and inaccurate forward models, leading to low-fidelity solutions. To address this issue, we propose to reparameterize and Gaussianize the latent tensors using novel differentiable data-dependent layers wherein custom operators are defined by solving optimization problems. These proposed layers constrain inverse problems to obtain high-fidelity in-distribution solutions. We validate our technique on three inversion tasks: compressive-sensing MRI, image deblurring, and eikonal tomography (a nonlinear PDE-constrained inverse problem) using two representative deep generative models: StyleGAN2 and Glow. Our approach achieves state-of-the-art performance in terms of accuracy and consistency.

1. INTRODUCTION

Inverse problems play a crucial role in many scientific fields and everyday applications. For example, astrophysicists use radio electromagnetic data to image galaxies and black holes (Högbom, 1974; Akiyama et al., 2019) . Geoscientists rely on seismic recordings to reveal the internal structures of Earth (Tarantola, 1984; Tromp et al., 2005; Virieux & Operto, 2009) . Biomedical engineers and doctors use X-ray projections, ultrasound measurements, and magnetic resonance data to reconstruct images of human tissues and organs (Lauterbur, 1973; Gemmeke & Ruiter, 2007; Lustig et al., 2007) . Therefore, developing effective solutions for inverse problems is of great importance in advancing scientific endeavors and improving our daily lives. Solving an inverse problem starts with the definition of a forward mapping from parameters m to data d, which we formally write as d = f (m) + , where f stands for a forward model that usually describes some physical process, denotes noise, d the observed data, and m the parameters to be estimated. The forward model can be either linear or nonlinear and either explicit or implicitly defined by solving partial differential equations (PDEs). This study considers three representative inverse problems: Compressive Sensing MRI, Deblurring, and Eikonal (traveltime) Tomography, which have important applications in medical science, geoscience, and astronomy. The details of each problem and its forward model are in App. A. The forward problem maps m to d, while the inverse problem estimates m given d. Unfortunately, inverse problems are generally under-determined with infinitely many compatible solutions and intrinsically ill-posed because of the nature of the physical system. Worse still, the observed data are usually noisy, and the assumed forward model might be inaccurate, exacerbating the ill-posedness. These challenges require using regularization to inject a priori knowledge into inversion processes to obtain plausible and high-fidelity results. Therefore, an inverse problem is usually posed as an optimization problem: where R(m) is the regularization term. Beyond traditional regularization methods such as the Tikhonov regularization and Total Variation (TV) regularization, deep generative models (DGM), such as VAEs (Kingma & Welling, 2013 ), GANs (Goodfellow et al., 2014) , and normalizing flows (Dinh et al., 2014; 2016; Kingma et al., 2016; Papamakarios et al., 2017; Marinescu et al., 2020) , have shown great potential for regularizing inverse problems (Bora et al., 2017; Van Veen et al., 2018; Hand et al., 2018; Ongie et al., 2020; Asim et al., 2020; Mosser et al., 2020; Li et al., 2021; Siahkoohi et al., 2021; Whang et al., 2021; Cheng et al., 2022; Daras et al., 2021; 2022) . Such deep generative models directly learn from training data distributions and are a powerful and versatile prior. They map latent vectors z to outputs m distributed according to an a priori distribution: m = g(z) ∼ p target , z ∼ N (0, I), for example. The framework of DGM-regularized inversion (Bora et al., 2017) is arg min m (1/2) d -f (m) (c) arg min z (1/2) d -f • g (z) 2 2 + R (z), where the deep generative model g, whose layers are frozen, reparameterizes the original variable m, acting as a hard constraint. Instead of optimizing for m, we now estimate the latent variable z and retrieve the inverted m by forward mappings. Since the latent distribution is usually a standard Gaussian, the new (optional) regularization term R (z) can be chosen as β z 2 2 for GANs and VAEs, where β is a weighting factor. See App. J.1 for more details on a similar formulation for normalizing flows. Since the optimal β depends on the problem and data, tuning β is highly subjective and costly. However, this formulation of DGM-regularized inversion still leads to unsatisfactory results if the data are noisy or the forward model is inaccurate, as shown in Fig. 20 , even if we fine-tune the weighting parameter β. To analyze this problem, first recall that a well-trained DGM has a latent space (usually) defined on a standard Gaussian distribution. In other words, a DGM either only sees standard Gaussian latent



+ R(m),(2)



Figure 1: Comparison of images generated by a deep generative model (DGM), Glow, using latent tensors that deviate from a spherical Gaussian distribution (left) and those after corresponding corrections (right). The visual effects highlight the necessity of keeping the latent tensor within such a distribution during inversion. The second column shows the characteristics of deviated latent tensors: (a) histogram: i.i.d. components but the distribution is skewed; (b) histogram: i.i.d. components but the distribution is heavy-tailed; (c) latent tensor image: non-i.i.d. entries. The first column shows the corresponding outputs of a Glow network. The third column shows latent tensors corrected by (a) the Yeo-Johnson layer (YJ), (b) the Lambert W × F X layer (LB), and (c) the full set of our Gaussianization layers (G layers). Those corrected latent tensors map to realistic images shown in the fourth column. All latent tensors have a norm of 0.7 √ vec dim because of the Gaussian Annulus Theorem (App. I) and the fact that Glow works best with a temperature smaller than one (see Fig. 9). Additional examples for StyleGAN2 and Stable Diffusion (Rombach et al., 2022) can be found in App. B.

