Table of Contents
Fetching ...

Differentiable Gaussianization Layers for Inverse Problems Regularized by Deep Generative Models

Dongzhuo Li

TL;DR

The paper tackles ill-posed inverse problems regularized by deep generative models, where latent codes can drift away from the standard Gaussian prior under noise and model error. It introduces optimization-based differentiable Gaussianization layers that reparameterize latent tensors into an in-distribution Gaussian, via a patch-based ICA stage followed by 1D Gaussianization (Yeo-Johnson and Lambert W × F_X) and standardization, enabling unconstrained optimization of latent variables. Across three tasks—compressive-sensing MRI, image deblurring, and eikonal tomography—the approach with StyleGAN2 and Glow achieves state-of-the-art accuracy and consistency, showing robustness to noise and forward-model mismatch; ablations identify ICA as the major contributor and larger patch sizes as beneficial. The work provides a plug-and-play regularization mechanism for DGMs in nonlinear inverse problems, with potential broad impact in imaging and geophysics, while noting computational costs and limitations from training-data bias and synthetic evaluations.

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.

Differentiable Gaussianization Layers for Inverse Problems Regularized by Deep Generative Models

TL;DR

The paper tackles ill-posed inverse problems regularized by deep generative models, where latent codes can drift away from the standard Gaussian prior under noise and model error. It introduces optimization-based differentiable Gaussianization layers that reparameterize latent tensors into an in-distribution Gaussian, via a patch-based ICA stage followed by 1D Gaussianization (Yeo-Johnson and Lambert W × F_X) and standardization, enabling unconstrained optimization of latent variables. Across three tasks—compressive-sensing MRI, image deblurring, and eikonal tomography—the approach with StyleGAN2 and Glow achieves state-of-the-art accuracy and consistency, showing robustness to noise and forward-model mismatch; ablations identify ICA as the major contributor and larger patch sizes as beneficial. The work provides a plug-and-play regularization mechanism for DGMs in nonlinear inverse problems, with potential broad impact in imaging and geophysics, while noting computational costs and limitations from training-data bias and synthetic evaluations.

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.
Paper Structure (44 sections, 2 theorems, 49 equations, 21 figures, 11 tables, 3 algorithms)

This paper contains 44 sections, 2 theorems, 49 equations, 21 figures, 11 tables, 3 algorithms.

Key Result

Theorem 1

For an n-dimensional standard Gaussian, for any $\beta\leq \sqrt{n}$, all but at most $3e^{-c\beta^2}$ of the probability mass lies within the annulus $\sqrt{n} - \beta \leq |x| \leq \sqrt{n} + \beta$, where $c$ is a fixed positive constant.

Figures (21)

  • 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\times 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$\sqrt{\text{vec dim}}$ because of the Gaussian Annulus Theorem (App. \ref{['append:gaussian_typical_set']}) and the fact that Glow works best with a temperature smaller than one (see Fig. \ref{['fig:various_radii']}). Additional examples for StyleGAN2 and Stable DiffusionRombach_2022_CVPR can be found in App. \ref{['append:add_latent_deviations']}.
  • Figure 2: Illustration of the proposed inversion process. Gradient computation in the Gaussianization layers is enabled by the implicit function theorem and automatic differentiation (App. \ref{['append:opt_layer_backward']}). We use the L-BFGS optimizer nocedal2006numerical to update ${\mathbf{v}}$.
  • Figure 3: The nonlinear activation functions from (a) the power transformation (Yeo-Johnson) layer and (b) the Lambert $W \times F_X$ layer.
  • Figure 4: Comparison of compressive sensing MRI inversion results (Accl=8x, SNR=20 dB).
  • Figure 5: Comparison of deblurring results from different methods.
  • ...and 16 more figures

Theorems & Definitions (4)

  • Theorem 1: Gaussian Annulus Theorem blum2020foundations
  • Definition 1: cover2012elements
  • Definition 2: Gaussian Typical Set
  • Theorem 2: cover2012elements