D IFFERENTIABLE G AUSSIANIZATION L AYERS FOR I NVERSE P ROBLEMS R EGULARIZED BY D EEP G ENERATIVE M ODELS
A BSTRACT
Deep generative models such as GANs, normalizing ﬂows, 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-ﬁdelity solutions.
To address this issue, we propose to reparameterize and Gaussianize the latent tensors using novel differentiable data-dependent layers wherein custom operators are deﬁned by solving optimization problems.
These proposed layers constrain inverse problems to obtain high-ﬁdelity 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.
I NTRODUCTION
Inverse problems play a crucial role in many scientiﬁc ﬁelds 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 scientiﬁc endeavors and improving our daily lives.
Solving an inverse problem starts with the deﬁnition of a forward mapping from parameters m to data d , which we formally write as
[Equation]
where f stands for a forward model that usually describes some physical process, (cid:15) 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 deﬁned 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 inﬁnitely 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-ﬁdelity results.
Therefore, an inverse problem is usually posed as an optimization problem:
[Equation]
[Figure]
Latent
[Figure]
properties the works best smaller than one (
Fig. 9)
Additional examples for StyleGAN2 and Stable Diffusion (Rombach et al., 2022) can be found in
App.
B. 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 ﬂows (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
[Equation]
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 (cid:48) ( z ) can be chosen as β (cid:107) z (cid:107) 22 for GANs and VAEs, where β is a weighting factor.
See App.
J.1 for more details on a similar formulation for normalizing ﬂows.
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 ﬁne-tune the weighting parameter β .
To analyze this problem, ﬁrst recall that a well-trained DGM has a latent space (usually) deﬁned on a standard Gaussian distribution.
In other words, a DGM either only sees standard Gaussian latent tensors during training ( e.g. , GANs) or learns to establish 1-1 mappings between training examples and typical samples from the standard Gaussian distribution ( e.g. , normalizing ﬂows, App.
J.2).
As a result, the generator may map out-of-distribution latent vectors to unrealistic results.
We show in
Fig.
1 the visual effects of several types of deviations of latent vectors from a spherical Gaussian (with a temperature of 0.7).
It can be seen that (1) the latent tensor should have i.i.d.
entries, and (2) these entries should be distributed as a 1D standard Gaussian in order to generate plausible images.
Since the traditional DGM-regularized inversion lacks such Gaussianity constraint, we conjecture 1 that the latent tensor deviates from the desired high-dimensional standard Gaussian distribution during inversion, leading to poor results.
Our observations and reasoning motivate us to propose a set of differentiable Gaussianization layers to reparameterize and Gaussianize the latent vectors of deep generative models ( e.g. , StyleGAN2 (Karras et al., 2020) and Glow (Kingma & Dhariwal, 2018)) for inverse problems.
The implementation is available here .
M ETHOD 2. O VERVIEW - REPARAMETERIZATION USING G AUSSIANIZATION LAYERS
Our solution is based on a necessary condition of the standard Gaussian prior on latent tensors.
Let us deﬁne a partition of a latent tensor z ∈ R n as the collection of non-overlapping patches
P ( z ) =
{ z i } i =1 , ··· ,N , where the patches z i ∈
R D are of the same dimension and can be assembled as z , i.e. , n = N × D .
If the latent tensor z ∼ N ( 0 , I ) , then for all z i from any partition P ( z ) , we have z i ∼ N ( 0 , I ) .
Note that z is a symbolic representation of the latent tensor.
In a speciﬁc DGM, z can be either a 2D/3D tensor or a list of such tensors corresponding to a multi-scale architecture (App. D).
Even though there are numerous partition schemes, such as random grouping of components, we choose the simplest: partitioning z based on neighboring components, which works well in practice.
To constrain z i ∼ N ( 0 , I ) during inversion, we reparameterize z i by constructing a mapping h : v i
→ z i , such that z i ∼ N ( 0 , I ) .
The new variables v i are of the same dimension as z i .
Suppose that we have constructed the patch-level mapping h , we can obtain a mapping h † at the tensor level, such that z = h † ( v ) , where v is assembled from v i in the same way as z from z i .
For example, h † = diag ( h, · · · , h ) if patches are extracted from neighboring components and are concatenated into a vectorized v .
Hence, the original DGM-regularized inversion 3 becomes
[Equation]
Since we are imposing a constraint through reparameterization, there is no need to include the regularization term R (cid:48)
( z ) .
The new formulation is still an unconstrained optimization problem, enabling us to use highly efﬁcient unconstrained optimizers, such as L-BFGS (Nocedal & Wright, 2006) and ADAM (Kingma & Ba, 2015).
The remaining critical piece is to construct h , and it leads to our Gaussianization layers.
First, we translate the constraint of z i
= h ( v i ) ∼ N ( 0 , I ) into the following optimization problem:
[Equation]
where p Z is the probability density function (PDF) of z i .
Second, we adopt the framework proposed by precursor works of normalizing ﬂows on Gaussianization (Chen & Gopinath, 2000; Laparra et al., 2011) to solve this optimization problem.
The KL-divergence can be decomposed as the sum of the multi-information I ( z i ) and the marginal negentropy J m ( z i ) (Chen & Gopinath,
[Equation]
Meng et al., 2020):
[Equation]
where
[Equation]
Here z ( j ) i denotes the j -th component of patch vector z i , and p j stands for the marginal PDF for that component.
The multi-information I ( z i ) quantiﬁes the independence of the components of z i , while the marginal negentropy J m ( z i ) describes how close each component is to a 1D standard
Gaussian.
With this decomposition, the optimization procedure depends on the facts: (1) the KL divergence and a standard Gaussian in Eq.
6 are invariant to an orthogonal transformation, and (2) the multi-information term is invariant to a component-wise invertible differentiable transformation (App. J.3).
As a result, we perform Gaussianization in two steps: 1.
Minimize the multi-information I ( z i ) .
This is done by an orthogonal transformation that keeps the overall KL divergence the same but increases the negentropy J m ( z i ) .
We achieve this by using our independent component analysis (ICA) layer.
ICA is the optimal choice since it maximizes non-Gaussianity so that the subsequent marginal Gaussianization step removes it and results in a large decrease in D KL ( p Z ( h ( v i ))
(cid:107)N ( 0 , I )) .
Minimize
the marginal negentropy
J m
( z i
) by component-wise operations that perform 1D
Gaussianization of marginal distributions p j , j =1 , ··· ,D .
The multi-information does not change under component-wise invertible operations.
Therefore, the overall KL divergence between z i and the
Gaussian distribution decreases.
The Gaussianization steps are well-aligned with the motivating example (Fig. 1).
To constrain DGM outputs to be plausible, one should make components within latent tensor patches independent (or destroy the patterns) (Fig. 1(c)) and shape the 1D distribution as close as possible to Gaussian (Fig. 1(a)(b)).
We will see that h is parameterized by an orthogonal matrix and two scalar parameters in 1D Gaussianization layers.
Unlike conventional neural network layers, the input-data-dependent Gaussianization layers are not deﬁned by learning from a dataset ahead of time but by solving certain optimization problems on the job (Fig. 18).
Special care should be taken to implement the gradient computation correctly and ensure that they pass the ﬁnite-difference convergence test (App. G).
As an overview, the composition of our proposed layers is:
[Equation]
where whitening and ICA belong to the ﬁrst step and the rest belong to the second step.
We will discuss in App.
K some possible simpliﬁcations of the layers in practice after our ablation studies.
The overall proposed inversion process (one iteration) is illustrated in Fig.
[Figure]
d
2. R EDUCING MULTI - INFORMATION – ICA LAYER
The orthogonal matrix W is constructed by the independent component analysis (ICA).
The input is the patch vectors { v i | i =1 , ··· ,N } .
The ICA algorithm ﬁrst computes the input-dependent orthogonal matrix W and then computes p i
= W (cid:62) v
i , i =1 , ··· ,N as the output.
The orthogonal matrix W makes the entries of each p i independent random variables.
We use the FastICA algorithm (Hyvarinen, 1999; Hyvärinen & Oja, 2000), which employs a ﬁxedpoint algorithm to maximize a contrast function Φ ( e.g. , the logcosh function), for our ICA layer to reduce multi-information.
The FastICA algorithm typically requires that the data are pre-whitened.
We adopt the ZCA whitening method or the iterative whitening method introduced in Hyvarinen (1999)
(App. E).
With the whitened data, we compute W using a damped ﬁxed-point iteration scheme:
[Equation]
where 1 is an all-one vector, the column vectors of V are { v i | i =1 , ··· ,N } , φ ( · ) =
Φ (cid:48) ( · ) , α ∈ (0 , 1) , and we use α = 0 .
8 throughout our experiments.
To save computation time, we only perform a maximum of 10 iterations.
The details of the whole algorithm can be found in App.
E.
We set the initial W as an identity matrix.
If the input vectors are already standard Gaussian (indistribution), the computed W will still be an identity matrix, which maps the input to the same output.
In practice, the empirical distribution from ﬁnite samples is not a standard Gaussian, so W is not an identity matrix but another orthogonal matrix, which still maps standard Gaussian input vectors to standard Gaussian vectors as output.
For this reason, we sample starting { v i | i =1 , ··· ,N } from the standard Gaussian distribution to start inversions with plausible outputs.
2. R EDUCING MARGINAL NEGENTROPY
For 1D Gaussianization, we choose a combination of the Yeo-Johnson transformation that reduces skewness and the Lambert W × F X transformation that reduces heavy-tailedness.
Both are layers based on optimization problems with only one parameter, which is cheap to compute and is easy to back-propagate the gradient.
Eq.
7 requires us to perform such 1D transformations for each component of the random vectors.
In other words, we need to solve the same optimization problem for D times, which imposes a substantial computational burden.
Instead, we empirically ﬁnd it acceptable to share the same optimization-generated parameter across all components.
In other words, we perform only a single 1D Gaussianization, treating all entry values in the latent vector as the data simultaneously.
Power transformation layer We propose to use the power transformation or Yeo-Johnson transformation (Yeo & Johnson, 2000) to reduce the skewness of distributions.
As shown in Fig. 3(a), the form of the Yeo-Johnson activation function depends on parameter λ .
If λ = 1 , the mapping is an identity mapping.
If λ > 1 , the activation function is convex, compressing the left tail and extending the right tail, reducing the left-skewness.
If λ < 1 , the activation function is concave, which oppositely reduces the right-skewness.
We refer the readers to App.
E for details.
[Figure]
Lambert W × F X layer Due to noise and inaccurate forward models, we observe that the distribution of latent vector values tends to be shaped as a heavy-tailed distribution during the inversion process.
To reduce the heavy-tailedness, we adopt the Lambert W × F X method detailed in Goerg (2015).
We use the parameterized Lambert W × F X distribution family to approximate a heavy-tailed input and solve an optimization to estimate an optimal parameter δ (App. E), with which the inverse transformation maps the heavy-tailed distribution towards a Gaussian.
Fig. 3(b) shows that the Lambert W × F X layer acts as a nonlinear squashing function.
As δ increases, it compresses more the large values and reduces the heavy-tailedness.
Intuitively, the
Lambert W × F X layer can also be interpreted as an intelligent way of imposing constraints on the range of values instead of a simple box constraint.
We refer the readers to App.
E for more details about the optimization problem and implementation.
Standardization
with temperature Since the Lambert W × F X layer output may not necessarily have a zero mean and a unit (or a prescribed) variance, we standardize the output using
[Equation]
where γ is the temperature parameter suggested in Kingma & Dhariwal (2018).
R ELATED WORK
End-to-end NNs for inverse problems There are numerous end-to-end neural networks designed for inverse problems, using CNNs (Chen et al., 2017; Jin et al., 2017; Sriram et al., 2020), GANs (Mardani et al., 2018; Lugmayr et al., 2020; Wang et al., 2018), invertible networks (Ardizzone et al., 2018), and diffusion models (Kawar et al., 2022).
The general idea is simple: train a neural network that directly maps observed data to estimated parameters.
Even though such methods seem effective in a few applications, DGM-regularized inversion with our Gaussianization layers is preferable for the following reasons.
First, the forward modeling can be so expensive computationally that it is infeasible to collect a decent training datasets for some applications.
For example, one large-scale ﬂuid mechanics or wave propagation simulation can take hours, if not days.
Second, the relationship between parameters and data can be highly nonlinear, and multiple solutions may exist.
An end-to-end network may map data to interpolated solutions that are not realistic.
In comparison, our method can start from different initializations or even employ sampling techniques to address this issue.
Third, the conﬁguration of data collection can change from experiment to experiment.
It is impractical, for example, to re-train the network each time we change the number and locations of sensors.
In contrast, not only can our method deal with this situation, but it can even use the same DGM for different forward models, as we can see in the compressive sensing MRI and eikonal tomography examples.
While almost all end-to-end methods are only applied to linear inverse problems,
our
Gaussianization layers are also effective in nonlinear problems.
Other techniques to improve DGM-regularized inversion In high-dimensional space, the probability mass of a standard Gaussian distribution concentrates within the so-called Gaussian typical set (App. I).
To be in the Gaussian typical set, one necessary but not sufﬁcient condition is to be within an annulus area around a high-dimensional sphere (App. I).
Utilizing this geometric property, DGMregularized inversion methods like Bojanowski et al.
(2017) and Liang et al.
(2021) force updated latent vectors to stay on the sphere.
This strategy is closely related to spherical interpolation (White, 2016).
We call this strategy the spherical constraint for inversion.
In the original StyleGAN2 paper, the authors also noticed that in image projection tasks, the noise maps tend to have leakage from signals (Karras et al., 2020) – the same phenomenon we discussed.
They proposed a multi-scale noise regularization term ( NoiseRlg ) to penalize spatial correlation in noise maps.
We extend the same technique to our inverse problems for comparison.
Note that we use a whitening layer before the ICA layer.
The whitening layer can be used alone, similar to Huang et al.
(2018) and Siarohin et al.
(2018), whose performance will be reported in ablation studies.
Also, Wulff & Torralba (2020) observed that for StyleGAN2, the Leaky ReLU function can “Gaussianize" latent vectors in the W space.
CSGM-w (Kelkar & Anastasio, 2021) utilizes this a priori knowledge to improve DGMregularized compressive sensing problems.
To further compare with the Gaussianizaion layers, we in addition propose an alternative idea that reparameterizes latent vectors using learnable orthogonal matrices (Cayley parameterization) and ﬁxed latent vectors, which is closely related to the work of orthogonal over-parameterized training (Liu et al., 2021)
(App. H ).
Recently, score-based generative models have started to show promise for solving inverse problems (Jalal et al., 2021; Song et al., 2021).
However, they have been mainly applied to linear inverse problems and challenged by noisy data (Kawar et al., 2021).
Besides, scored-based methods might not work for certain physics-based inverse problems, since Gaussian noise parameters may break the physics simulation engine 2 .
E XPERIMENTS
We consider three representative inversion problems for testing: compressive sensing MRI, image deblurring, and eikonal traveltime tomography.
For MRI and eikonal tomography, we used synthetic brain images as inversion targets and used the pre-trained StyleGAN2 weights from Kelkar &
Anastasio (2021) (trained on data from the databases of fastMRI (Zbontar et al., 2018; Knoll et al., 2020), TCIA-GBM (Scarpace et al., 2016), and OASIS-3 (LaMontagne et al., 2019)) for regularization.
We used the test split of the CelebA-HQ dataset (Karras et al., 2018) for deblurring, and the DGM is a Glow network trained on the training split.
We refer readers to App.
F for details on datasets and training.
We tested each parameter conﬁguration in each inversion on 100 images (25 in the eikonal tomography due to its expensive forward modeling).
Since the deep generative models are highly nonlinear, the results may get stuck in local minima.
Thus, we started inversion using three different randomly initialized latent tensors for each of the 100 or 25 images, picked the best value among the three for each metric, and reported the mean and standard deviation of those metrics, except for CSGM-w,
TV, and NoiseRlg, where the initialization is ﬁxed.
The metrics we used were PSNR, SSIM (Wang et al., 2004), and an additional LPIPS (Zhang et al., 2018) for the CelebA-HQ data.
We used the
LBFGS (Nocedal & Wright, 2006) optimizer in all experiments except TV, noise regularization, and CSGM-w, which use FISTA (Beck & Teboulle, 2009) or ADAM (Kingma & Ba, 2015).
The temperature was set to 1.0 for StyleGAN2 and 0.7 for Glow.
4. C OMPRESSIVE SENSING MRI USING S TYLE GAN
The mathematical model of compressive sensing MRI is
[Equation] [Equation]
) where A ∈ C M × N is the sensing matrix, which consists of FFT and subsampling in the k-space (frequency domain).
Eq.
10 is an under-determined system, and we use Accl = N/M to denote the acceleration ratio.
We also added i.i.d.
complex Gaussian noise with a signal-to-noise ratio (SNR) of 20 dB or 10 dB to the measured data.
See App. A.1 for more background information.
Table 1 compares the results from total variation regularization (TV), noise regularization (NoiseRlg) (Karras et al., 2020), spherical constraint/reparameterization: z = v / (cid:107) v (cid:107) 2 ∗ (cid:112) dim ( v ) ,
CSGM-w (Kelkar & Anastasio, 2021), our proposed orthogonal reparameterization (Orthogonal), and our proposed Gaussianization layers (G layers).
Fig.
4 shows examples of inversion results.
In the base case where Accl=8x and SNR=20 dB, the Gaussianization layers gives the best scores, and this advantage gets more signiﬁcant when data SNR decreases to 10 dB. Interestingly, the scores from all methods improve signiﬁcantly if we make the system better determined ( i.e. , Accl=2x), and the performance of TV, spherical constraint, and Gaussianization layers become more similar in this scenario.
We conclude that our proposed Gaussianization layers are effective and more robust than other methods, particularly in low-SNR scenarios.
[Figure]
Ablation study Table 2 summarizes the ablation study on the components of the Gaussianization layers.
We kept the standardization layer on for all cases.
The conclusions are as follows: 1.
The ICA layer plays the most signiﬁcant role in improving result scores; 2.
The whitening/ZCA layer alone is not effective; 3.
The Yeo-Johnson (YJ) and the Lambert W × F X (Lambt) layers are more effective when the noise level is higher ( e.g. , SNR=10 dB vs. 20 dB).
Their performance seems data-dependent: when SNR=20 dB, YJ seems more effective, while Lambt gives the best scores when SNR=10 dB.
But the overall improvement from these two layers is marginal.
In practice, one may only use ICA and one of the 1D Gaussianization layers.
Additionally, we tested the effect of patch size on the z (style) vectors (App. C).
We ﬁnd that the largest possible patch size gives the best results.
[Table] [Table]
4. I MAGE DEBLURRING USING G LOW
The mathematical model for deblurring is
[Equation]
11) where H is a smoothing ﬁlter and ∗ denotes convolution.
We used a Gaussian smoothing ﬁlter with a standard deviation of 3, and added noise (cid:15) ∼ N ( 0 , 50 2 I ) to the observed data.
Although the system may not be under-determined, the high-frequency information is lost due to low-pass ﬁltering; hence this is also an ill-posed problem.
Though we only tested on facial images, deblurring has wide applications in astronomy and geophysics.
See App.
A.2 for more background information.
[Figure]
Table 3 and Fig.
5 show that the Gaussianization layers are also effective in Glow, better than using the spherical constraint or the orthogonal reparameterization.
We also demonstrate the efﬁcacy of
Gaussianization layers when the forward model is inaccurate, in which case the induced error in data is not Gaussian (App. C).
Ablation study Table 3 also shows that the two parameterization schemes (App. D) for Glow using the Gaussianization layers have similar performance.
Besides, we report the ablation study on the components of the Gaussianizaiton layers for Glow (App. C).
4. E IKONAL TOMOGRAPHY USING S TYLE GAN
In acoustic wave imaging (e.g., ultrasound tomography), we excite waves using sparsely distributed sources one at a time at the boundary of the object.
Then we reconstruct its internal structures (the spatial distribution of wave speed) given the ﬁrst-arrival travel time recorded on the boundary.
The following eikonal equation approximately describes the shortest travel time T ( x ; x s ) that the acoustic wave emerging from the source location x s takes to reach location x inside the target object (Yilmaz, 2001): |∇ T ( x ; x s ) | = 1 /c ( x ) ,
T ( x s ; x s ) = 0 , (12) where c ( x ) is the wave propagation speed at each location.
Both this eikonal PDE and the implicitly deﬁned forward mapping c ( x ) → T ( x ) are nonlinear, and there has been little research on DGMregularized inverse problems with such nonlinear characteristics.
The inverse problem is severely illposed, which is equivalent to a curved-ray tomography problem.
See App.
A.3 for more background information.
We added noise to the recorded travel time using the following formula: T noisy ( x r ; x s ) =
T ( x r ; x s )(1 + (cid:15) ) , where (cid:15) ∼ N (0 , 0 .
001 2 ) and x r denotes any receiver location.
In other words, longer traveltime corresponds to larger uncertainties.
We show that the Gaussianization layers outperformed other methods in this tomography task in
Table 4 and Fig. 6.
Note that this is a statistical conclusion.
We also report an example where the spherical constraint works better than the Gaussianization layers (bottom right).
[Table] [Table]
D ISCUSSION AND CONCLUSIONS
We provide insights on the experiment results and further discuss this work’s limitations, computational cost, and broader impact in App.
K.
In summary, we have identiﬁed a critical problem in DGM-regularized inversion: the latent tensor can deviate from a typical example from the desired high-dimensional standard Gaussian distribution, leading to unsatisfactory inversion results.
To address the problem, we have introduced the differentiable Gaussianization layers that reparameterize and Gaussianize latent tensors so that the inversion results remain plausible.
In general, our method has achieved the best scores in terms of mean values and standard deviations compared with other methods, demonstrating our method’s advantages and high performance in terms of accuracy and consistency.
(Regarding the standard deviation of scores, we only compare with other methods with competitive mean scores, such as orthogonal reparameterization and spherical constraint.)
Our proposed layers are plug-and-play, require minimal parameter tuning, and can be applied to various deep generative models and inverse problems.
A CKNOWLEDGEMENTS
The author would like to thank Huseyin Denli, Ashutosh Tewari, Myun-Seok Cheon, Di Du, Stuart
Harwood, Yu Fan and Qiuzi Li for helpful discussions and the anonymous reviewers for their constructive feedback, which greatly improved the paper
A B ACKGROUND OF FORWARD MODELS A. C OMPRESSIVE SENSING MRI
The MRI process essentially samples the spatial frequency components of a target, following some trajectories in the spatial frequency space (k-space) according to the design of the physical system.
For example, a system may sample the k-space line-by-line horizontally/vertically or in radial directions.
If the k-space has been fully sampled on a Cartesian grid, one can directly use inverse FFT to reconstruct the image.
For various practical reasons, however, it is necessary to speed up the data collection process, usually by skipping data points in the k-space, which can be mathematically represented by a masking operation.
In addition, there can be multiple coils with different sensitivity maps collecting data simultaneously.
The mathematical formulation reads
[Equation]
) where d i ∈ C N is the k-space data corresponding to the i -th coil, m ∈ R N is the target object,
P ∈ R N × N is the mask, F ∈ C N × N is the Fourier transform operator, S i ∈ C N × N is the point-wise sensitivity map (a diagonal matrix) corresponding to the i -th coil, and (cid:15) denotes noise.
To ensure a fair comparison with prior work and reproducibility, we used the same masks from the repository of Kelkar & Anastasio (2021) (Fig. 7).
In addition, we also used the same single-coil setup as in Kelkar & Anastasio (2021), where the sensitivity matrix is an identity matrix.
[Figure]
We condense the effects of all operators into a linear operator A ∈ C M × N and arrive at the underdetermined system
[Equation]
) from the main text, where we use Accl = N/M to denote the acceleration ratio.
A
. D EBLURRING
The mathematical model behind deblurring is
[Equation]
[Equation]
where H is a smoothing ﬁlter, ∗ denotes convolution, and (cid:15) is noise.
The purpose of deburring is to recover the original sharp image m given a noisy blurred observation d .
In this study, we showed deblurring examples for natural images.
In scientiﬁc applications, deblurring or deconvolution is also a powerful tool.
For example, in astronomy, d is a blurred image from a telescope, H is a point-spread function (PSF) constructed from the physics model of the telescope, and we want to obtain a sharper image from the observation (Starck et al., 2002).
In geophysics, d can be the seismic data, H is a calibrated wavelet, and we want to obtain sharp images of reﬂectivities deﬁning the boundaries of subsurface strata (Lines & Treitel, 1984; Zhang & Castagna, 2011).
In general, H is a low-pass or band-pass ﬁlter, so certain frequency contents are lost in the forward process.
The deblurring or deconvolution process needs to recover such missing information.
In addition, the noise makes the inversion process unstable.
The deblurring or deconvolution problem is thus ill-posed.
[Figure]
A. E IKONAL TOMOGRAPHY
In acoustic wave imaging (e.g., ultrasound tomography), we excite waves using sparsely distributed sources one at a time at the boundary of the object.
Then we reconstruct its internal structures (the spatial distribution of wave speed) given the ﬁrst-arrival travel time recorded on the boundary.
The following eikonal equation describes the shortest travel time T ( x ; x s ) that the acoustic wave emerging from the source location x s takes to reach location x inside the target object in the high-frequency limit (Yilmaz, 2001):
[Equation]
where c ( x ) is the wave propagation speed at each location.
Both this eikonal PDE and the implicitly deﬁned forward mapping c ( x ) → T ( x ) are nonlinear.
To be more speciﬁc, we show the setup of our experiment in Fig. 8(a).
For the convenience of numerical testing, we put sources and receivers on the boundaries of a square box that contains the object, suggesting that the object is immersed in a square box ﬁlled with water.
In reality, one can put sources and receivers directly on the target.
The dimension of the square area is 25 .
6 cm × 25 .
6 cm with a grid interval of 0.001 m. There are eight sources located on each side, and receivers are located at each grid point on the boundary.
Fig. 8(b) shows the shortest travel time ﬁeld of the generated wave from the indicated source location.
For each source, we only use receivers on the three other sides, indicated by the red lines, excluding the one on which the source is located.
Fig. 8(c) shows the proﬁle of the shortest wave travel time recorded at the receiver array.
The receiver index starts from 0 in the top-left corner and increases in the clockwise direction.
We can interpret the nonlinearity and ill-posedness of eikonal tomography from another perspective.
As shown in Fig.
8(b), we plot contours T ( x ) =
const representing the wavefronts.
Under the high-frequency approximation, the propagation of waves can be viewed as curved rays traveling in directions perpendicular to such wavefronts.
Since the wavefronts depend on velocity ﬁeld c ( x ) , the rays are also functionals of the parameter c ( x ) to be estimated, contrary to straight-ray tomography such as CT.
The eikonal inversion problem is thus nonlinear.
Besides, the rays carry information about the medium averaged along their paths.
One property of curved-ray tomography is that the ray coverage is uneven inside the object.
In fact, curved rays tend to avoid low-velocity areas, giving us little information about such regions, making the inverse problem intrinsically ill-posed.
We solved the eikonal equation using the fast sweeping method (Zhao, 2005) and computed the gradient using the discrete adjoint-state method, using the code from the repository of Li et al.
Our eikonoal tomography uses the same StyleGAN2 network as the compressive sensing MRI experiments.
The value range of StyleGAN2 is [ − 1 , 1] .
In the forward model, we map the StyleGAN output to c ( x ) in two steps.
First, we convert its values to the range of [0 , 1] by using m ← ( m +1) / 2 .
Second, we convert the values to the range of acoustic wave velocity using m ← 100 × m + 1500 .
This relationship is purely manufactured for our synthetic tests.
One should use a more realistic relationship in practice.
B A DDITIONAL MOTIVATING EXAMPLES
Glow We varied the (cid:96) 2 -norm of the latent tensors for Glow and reported the outputs in Fig.
The images are getting increasingly unrealistic as we increase the norm from 1.0 √ d , where d is the latent tensor dimension.
This phenomenon demonstrates why we need to use a temperature < 1 for Glow.
[Figure]
various norms: 0.1, 0.7, 1.0, and 2.0 of √ d (the square root of the latent tensor dimension).
The images are getting increasingly unrealistic as we increase the norm from 1.0 √ d .
This phenomenon demonstrates why we need to use a temperature < 1 for Glow.
StyleGAN2 StyleGAN2 is a well-designed generator with a built-in spherical transformation for the style vector.
It seems very robust to random vectors drawn from the distributions similar to those in Fig.
However, we were able to ﬁnd challenging examples for it from inversion tasks.
In Fig.
The ﬁrst example is the direct output only using the pathologic style vector.
The second example is the generator output using the style vector after whitening only.
As we can see, whitening alone is ineffective in improving the image quality, which is conﬁrmed by our ablation studies.
On the contrary, if we use the ICA layer, we see a huge improvement in visual quality, and our interpretation is that getting rid of higher-order dependencies is very important.
Note that there are eyeglasses in the third ﬁgure, meaning that if we relax the 1D Gaussian requirement, we can sample some rare examples.
Finally, we use the full G layers to get another plausible image.
Stable diffusion Stable Diffusion (Rombach et al., 2022) is a state-of-the-art deep generative model with text-to-image synthesis capability, which maps a Gaussian latent tensor to a high-resolution image.
[Figure]
Before
After
S k
e w e d
H ea vy - t a il e d
N on - i .
i .
d .
Before
After
Before
After
[Figure]
We repeated the tests for Glow and StyleGAN2 on Stable Diffusion and reported the results in
Fig.
The distribution for skewed case was the exponential-gamma distribution log Γ(1 , 1) (or the “log-gamma distribution" in scipy), the distribution for the heavy-tailed case was a Lambert W × F X distribution with parameter δ = 0 .
5 based on a standard Gaussian (Goerg, 2015).
To generate the latent tensor with non-i.i.d.
entries, we ﬁrst sampled z 0 ∈ R 4 × 64 × 64 ∼ N ( 0 , I ) .
Then we added 0 .
5 ∗ sin(2 π/ 64 x ) ⊗ cos(2 π/ 64 y ) to each channel of z 0 for the latent tensor z , where ⊗ stands for outer product.
The number of denoising steps, classiﬁer-free guidance, and the output dimensions for
Stable Diffusion are 50, 7.5, and 512 × 512 , respectively.
The outputs of Stable Diffusion are completely out-of-range if the latent tensors deviate from standard
Gaussian, with much poorer quality than those from Glow and StyleGAN2.
For example, Stable
Diffusion picks up and ampliﬁes the sparse large-amplitude points in the heavy-tailed case and the 2D sinusoidal pattern in the non-i.i.d.
case.
However, with our Gaussianization layers, we are still able to constrain the outputs within the plausible range.
Interestingly, we show in Fig.
12 that Stable
Diffusion is much more sensitive to the norm of latent tensors than the other DGMs.
Therefore, we anticipate that our Gaussianization layers are critical for potential inversion applications using Stable
Diffusion such as text-guided inversion.
C A DDITIONAL EXPERIMENTS
First, we report in Table 5 an ablation study on the effects of various combinations of regularization techniques on style vectors and noise maps in StyleGAN2.
The observations are (1) If we turn off the update of noise maps, the Gaussianization layers lead to the best result; (2) If we turn on the update of noise maps, applying the Gaussianization layers to both the style vectors and noise maps gives the best results.
[Table]
Second, we studied the effect of patch size on the performance of Gaussianization layers.
To be more speciﬁc, we tested the effects of various patch sizes on 1D style vectors.
The largest patch size is 64 since the number of extracted patches should not be smaller than the dimension of the patches.
We only turned on the ICA layer and the standardization layer in this experiment.
One can observe that both the PSNR and the SSIM increase as the patch size increases.
We would advise using the largest possible patch size in Gaussianization layers, although we only used a patch size of 32 for style vectors in all other experiments.
In all experiments, we ﬁxed the patch size for Glow as 3 × 8 × 8 and the patch size for the noise maps in StyleGAN2 as 1 × 8 × 8 .
In the experiments with Glow, the image dimension was 3 × 128 × 128 .
If we chose a larger patch size, we could not obtain enough patch vectors.
As for StyleGAN2, if we look at the parameterization illustrated in Fig.
14, there is already a 4 × 4 noise map unconstrained if we choose the patch size as 8 × 8 .
If we increase the patch size to 1 × 16 × 16 , there will be two additional 8 × 8 noise maps unconstrained.
To minimize the inﬂuence of unconstrained parameters, we chose the patch size for noise maps as 1 × 8 × 8 .
[Table]
Third, we did a parameter sweep on the weighting parameter β in the Glow-regularized deblurring problem (formulation 3 or 44).
Consistent with Fig.
20, the conventional Glow-based regularization is not as effective as our Gaussianization layers.
If β is too large ( e.g. , 10 or 100), the inversion results become unrealistic with huge errors.
[Table]
Fourth, we investigated the performance of Gaussianization layers when the forward model is inaccurate (Table 8).
Our Gaussianization layers still outperform the conventional Glow-regularized inversion.
deblurring with an inaccurate ﬁlter.
The ground-truth standard deviation of the Gaussian ﬁlter is 3, which is used to generate the observed data, but we used 5 for inversion.
[Table]
In addition, we did an ablation study on the components of Gaussianization layers in Glow-regularized deblurring (Table 9).
We adopted the second parameterization scheme (App. D).
The observations are similar to those from Table 2: 1.
The ICA layer is the most signiﬁcant part in improving result scores, especially LPIPS – the score that matches human perception the best; 2.
The whitening/ZCA layer is not effective when used alone; 3.
There is no clear winner among combinations of 1D
Gaussianization layers, and their difference in performance is marginal.
So we picked the Lambert
W × F X layer, one of the best-performing options, in all other deblurring experiments as the 1D
Gaussianization layer.
Inversion using out-of-distribution images Finally, we tested the limit of DGM-regularized inversion by using real-world out-of-distribution target images.
We randomly sampled 25 MRI images from the fastMRI DICOM dataset (Zbontar et al., 2018; Knoll et al., 2020).
The RSNA clinical trial processor was used to anonymize the whole dataset, and each image has been manually inspected to make sure that there is no protected information.
According to the fastMRI paper (Zbontar et al., 2018), "This dataset represents a larger variety of machines and settings than are present in the raw data."
Also, the image values are represented in uint16 rather than ﬂoat numbers, and their value range can be very large (0 to several thousand).
So we normalized the image values to be
[Table]
between range [ − 1 , 1] using img =
(img - img.min())/(img.max()-img.min())* - 1 .
However, the trained StyleGAN2 outputs are always slightly beyond [ − 1 , 1] , and we clipped the values in both test image generation and inversion.
Therefore, the new examples are out-ofdistribution.
We kept the same experimental setup of 8 × acceleration and 20 dB SNR and tested inversion using the spherical constraint and Gaussianization layers on both only using the style vectors and using the additional noise maps.
We report the inversion results in Table 10 and Fig. 13.
To save experiment time, we only use one 1D Gaussianization layer (either Yeo-Johnson or Lambert W × F X )
[Table] [Table]
First, none of the inversion results shown here are visually comparable to the ground truth, which shows the necessity of having a large training dataset that covers the target distribution.
Also, it is vital to ensure that the pre-processing and value ranges of training and target images are the same.
Second, in terms of metrics, the results are comparable to but slightly worse than those on synthetic test images reported previously.
This is expected because these images are out of distribution.
Besides, in cases where we applied Gaussianization layers, using both noise maps and style vectors gives better results than using style vectors only.
In addition, Gaussianization layers improved results when we also optimized the noise maps in addition to the style vectors.
However, we found that only using the style vector and the spherical constraint gives the best result on this test set, which contradicts our results reported in Table 5.
Our interpretation is that the dimension of style vectors is still relatively lower than noise maps or the latent tensor used in Glow, so our Gaussianization layers may be more effective in dealing with the latter two types of parameters.
D R EPARAMETERIZATION SCHEMES FOR S TYLE GAN2 AND G LOW
Fig.
14 illustrates the reparameterization scheme for StylgeGAN2 using the Gaussianization layers.
The patch size for the style vectors and noise maps are 32 and 1 × 8 × 8 , respectively.
Note that we transpose the latent tensors in Fig.
14 compared to the notations in corresponding equations and algorithms.
Also, as pointed out in Gu et al.
(2020), we also ﬁnd that a single latent code is insufﬁcient for image reconstruction.
Thus, we use multiple style vectors that were fed into different intermediate layers.
[Figure]
[Figure] [Figure]
Parameters
[Figure]
Parameters
[Figure]
Unfolding vvvv
Vectors
Tensor
Tensor (1 channel)
Tensor (4 channels) vvvv
Squeezing
For Glow, we came up with two reparameterization schemes (Fig. 15).
The ﬁrst one (P1) reparameterizes patches from a latent tensor with the same dimension as the output image.
The dimensions of the patches and the image are 3 × 8 × 8 and 3 × 128 × 128 , respectively.
Then a multi-scale squeezing operation maps the latent tensor into a list of tensors corresponding to different scales of
Glow.
Glow uses the list of tensors as the input.
The second scheme (P2) reparameterizes patches extracted directly from tensors in the above list.
We also partition each tensor into 3 × 8 × 8 patches.
The unfolding and squeezing operations are illustrated in Fig.
In the multi-scale architecture of
Glow, the squeezing operation is applied recursively on half of the output tensors cut in the channel direction (Dinh et al., 2016; Kingma & Dhariwal, 2018).
E M ORE DETAILS OF THE G AUSSIANIZATION LAYERS E. ICA LAYER
The overall ICA layer is summarized in algorithm 1.
We set a maximum number of the ﬁxed-point iterations to reduce computational cost and ensure accurate gradient computation that can pass the ﬁnite-difference convergence test.
The gradients are backpropagated through the loops without using the trick introduced in App.
G.
Whitening The FastICA algorithm typically requires that the data are pre-whitened.
Let V be the data matrix whose columns are data vectors.
We ﬁrst subtract the mean from the data vector using V ← V − mean ( V , dim = 1) .
Then we compute the data covariance matrix using
C = (1 − η ) 1 N − 1 V V (cid:62) + η I , where we use a small constant (e.g., η = 0 .
001 ) to blend the empirical covariance matrix and an identity matrix to avoid ill-conditioning.
After these data preparation steps, the ZCA whitening used throughout the study ﬁrst computes the eigenvalues Λ and eigenvectors D , and then output the whitened data using V ← D Λ − 1 / 2 D (cid:62) V .
Alternatively, we can use the following steps to whiten the data, which are also used later in ICA iterations to decorrelate column vectors in the orthogonal matrix (Hyvarinen, 1999):
[Equation]
[Equation] [Equation]
Input :
V
N error
η
− α
8 ; maximum iteration numbers J =
K =
Whitening stage
1 V =
V − mean ( V , dim = 1) , C = (1 − η ) 1 N − 1 V V (cid:62) + η I ; 2 V ← ZCA-whitening ( V ) ; // ICA stage 3 W
W ∗ ← I , j ←
; 4 while j do
W ← 1 N (cid:104)
α V φ (cid:0) W (cid:62) V (cid:1) (cid:62) − W diag (cid:0) φ (cid:48) (cid:0) W (cid:62) V (cid:1) 1 (cid:1)(cid:105) Eq. 8 ;
W = W 0 ← W / (cid:112) (cid:107) W (cid:62) W (cid:107) 2 , k ← 1 ; while k
K do
W ← 32 W − 12 W W (cid:62) W ; if (cid:107) W
(cid:107) then k
< (cid:15) then break;
W ∗ ← W , j ← j + 1 ; 15 return P ← W (cid:62) V .
The modiﬁed FastICA iterations.
As stated in Hyvärinen (1999), the objective function for one neural unit of the weight vector w i and input v is
[Equation]
where Φ is the contrast function ( e.g. , logcosh).
The original derivations convert this constrained optimization to an unconstrained one using Lagrange multipliers.
However, this procedure is unnecessary since the matrix W is orthogonalized after each iteration, and the input vectors have been pre-whitened.
Therefore, we only need to solve the following equation
[Equation]
19) whose Jacobian is
[Equation]
where φ is the derivative of Φ .
The Newton iteration scheme is thus
[Equation]
To improve the convergence, we damp the iterations by a parameter α ∈ (0 , 1) .
Also, using the same technique to convert the Newton iterations to ﬁxed-point iterations in Hyvärinen (1999); Hyvarinen (1999), we arrive at the modiﬁed ﬁxed-point iteration scheme:
[Equation]
) followed by the aforementioned decorrelation procedure after each step.
The convergence of the modiﬁed FastICA iterations can be proved similarly as in Oja & Yuan (2006).
E. P OWER TRANSFORMATION LAYER Algorithm
We propose to use the power transformation or Yeo-Johnson transformation (Yeo & Johnson, 2000) to reduce the skewness of distributions:
[Equation]
where p is an input value, s is an output value, and λ is the parameter to be estimated.
As shown in
Fig. 3(a), the form of the Yeo-Johnson activation function depends on parameter λ .
If λ = 1 , the mapping is an identity mapping.
If λ > 1 , the activation function is convex, compressing the left tail and extending the right tail, reducing the left-skewness.
If λ < 1 , the activation function is concave, which oppositely reduces the right-skewness.
The only parameter λ is determined by solving an optimization problem that minimizes the negentropy:
[Equation]
where p is the input data vector with entries p
i,i =1 , ··· ,n .
We use a custom operator based on SciPy’s implementation using Brent’s algorithm (Brent, 2013) to ﬁnd an approximate minimum of Problem 24.
Continuing from the approximate minimum, we use
Brent’s root ﬁnding algorithm
(Brent, 2013) to ﬁnd the minimum where the gradient vanishes.
Since the parameter λ depends on input data, we need to back-propagate the gradient through the optimization process.
The power transformation layer is summarized in algorithm 2.
E. L AMBERT W × F X LAYER
Due to noise and inaccurate forward models, we observe that the distribution of latent vector values tends to be shaped as a heavy-tailed distribution during the inversion process.
To reduce the heavytailedness, we adopt the Lambert W × F X method detailed in Goerg (2015).
Let X be a random variable whose CDF is F X , with mean µ X and standard deviation σ X .
The following transformation with a heavy-tail parameter δ ≥ 0 :
[Equation]
where U =
( X − µ X ) /σ
X , is a bijection and maps X to another random variable S with heavier tails.
The transformation Eq.
25 is bijective if δ ≥ 0 , and we can use the Lambert W function to ﬁnd its inverse.
The Lambert W function W is deﬁned as the inverse of q = W − 1 ( t ) =
t exp ( t ) , where t and q are scalars.
Given q , Halley’s method can be used to ﬁnd t = W ( q ) (Corless et al., 1996).
Hence, the inverse of Eq.
25 is
[Equation]
where
[Equation]
We use the parameterized Lambert W × F X distribution family to approximate a heavy-tailed input distribution and use Eq.
26 to recover a distribution with lighter tails.
In order to make the recovered distribution close to a Gaussian distribution, we compute the optimal heavy-tail parameter δ by minimizing the difference between the kurtosis of the output distribution and 3 (Kurtosis is a common surrogate measure of negentropy (Hyvärinen & Oja, 2000)):
[Equation]
where s is the data vector, and Kurt is the kurtosis.
We constrain δ > 0 , and solve Eq.
28 using the
L-BFGS-B optimizer (Zhu et al., 1997).
In addition, we estimate the mean µ X and standard deviation σ X along with δ using the Iterative
Generalized Method of Moments (IGMM) (Goerg, 2015), which embeds an optimization problem for δ in an outer loop of iterations to estimate σ X and µ X (see algorithm 3).
If the kurtosis of input data vector is not greater than 3, we skip the whole Lambert W × F X layer by directly outputting the data vector.
Input : Data vector s , error tolerance (cid:15) = 10 − 5 , maximum iteration number K = 100 .
Output : vector x whose empirical distribution is less heavy-tailed than s , and its kurtosis ≈ 3 ) 1 Initialize: ξ (0) ← (cid:16) ˆ µ (0) X , ˆ σ (0) X , ˆ δ (0) (cid:17) , k = 0 ; 2 Compute initial kurtosis β 2 =
Kurt ( s ) ; 3 if β 2 ≤ 3 then return s ; 5 end 6 while k < K and (cid:107) ξ ( k ) − ξ ( k − 1) (cid:107) ≥ (cid:15) do u ( k ) ← (cid:16) s − µ ( k ) X (cid:17) /σ
( k ) X ;
Compute δ ( k +1) using Eq. 28; u ( k +1)
← W δ ( k +1) (cid:0) u ( k ) (cid:1) ; x ( k +1)
← u ( k +1)
σ ( k ) X + µ ( k ) X ;
Update µ ( k +1)
X ← E (cid:2) x
( k +1) (cid:3) , and σ ( k +1)
X ← (cid:113) Var (cid:0) x
( k +1) (cid:1) ; ξ ( k +1) ← (cid:16) ˆ µ ( k +1)
X , ˆ σ ( k +1)
X , ˆ δ ( k +1) (cid:17) ; k ← k + 1 ; 12 end 13 return x = W δ (cid:16) s − µ X σ X (cid:17)
σ X + µ X ,
TRAINING
For MRI and eikonal tomography, we generate synthetic brain images as inversion targets using the pre-trained StyleGAN2 weights from Kelkar & Anastasio (2021), which are trained on data from the databases of fastMRI (Zbontar et al., 2018; Knoll et al., 2020), TCIA-GBM (Scarpace et al., 2016), and OASIS-3 (LaMontagne et al., 2019).
As a result, there is no sensitive personal information in our target images.
In the data generation process, we only used one style vector of a dimension of 512.
We picked 100 images that are visually plausible, whose random seeds can be found in our code.
These random seeds were never used to initialize inversion.
In inversion experiments, we usedsuch style parameters: one for the lowest resolution of 4 × 4 , one for the tRGB layer, and two for each resolution from 8 × 8 to 256 × 256 , to avoid inversion crime.
This choice of expanding latent space dimension is also justiﬁed by our observation that only one style parameter vector is generally insufﬁcient for inversion tasks with data from datasets such as CelebA-HQ (Karras et al., 2018).
We used the CelebA-HQ dataset (Karras et al., 2018)
(under the Creative Commons CC BY-NC 4. license) for the deblurring experiments.
All images were downsampled to the resolution of 128 × 128.
We split the 30000 images from CelebA-HQ into the subsets of training (24183 images), validation (2993 images), and testing (2824 images) following the original splits from CelebA (Liu et al., 2015).
For the inversion tests, we randomly selected 100 images from the test set.
For the hyper-parameters of the Glow networks, we used 4 multi-scale levels and 32 ﬂow-steps, and we only used additive coupling layers.
Fig.
17 reports the training process.
For each epoch, we computed the training negative log-likelihood (NLL) averaged throughout the epoch, and the validation NLL at the end of the epoch.
The validation curves suggest that it is better to usemulti-scale levels.
We chose the network weights from the epoch before the validation NLL stopped to decrease: 850 for the CelebA-HQ dataset.
All training was conducted using 8 × 32 GB Nvidia
V100 GPUs with a batch size of 64.
We used the Adam optimizer (Kingma & Ba, 2015) with a learning rate of 10 − 4 , as well as β 1 = 0 .
9 and β 2 = 0 .
G G RADIENT COMPUTATION OF THE OPTIMIZATION - BASED DIFFERENTIABLE LAYERS
[Equation]
As shown in Fig.
18, in the power transformation and Lambert W × F X layers, there are operators whose outputs are obtained by solving optimization problems formally described as
[Equation]
where l denotes the objective function that deﬁnes the operator (combining the objective function and layer h θ in Fig. 18), symbol θ stands for the optimal output, a scalar in our cases but can also be a vector in general situations.
The optimal condition is
[Equation]
where the subscript denotes partial derivatives.
The optimal condition implicitly deﬁnes a forward operator of the following form:
[Equation]
The backward operator is
[Equation]
where χ is the objective function of an inverse problem.
Differentiating Eq.
30 with respect to x , we have
[Equation]
using the implicit function theorem.
Then, to back-propagate the gradient from ∂χ∂ θ to ∂χ∂ x , we use
[Equation]
[Equation]
where H θ is the Hessian matrix of χ with respect to θ .
In our problems, the output θ is a scalar, so it is easy to use automatic differentiation to compute L θ directly and hence L − 1 θ .
Otherwise, if θ has many parameters, we can ﬁrst solve the following linear system with an auxiliary vector λ :
[Equation]
(35) and then compute the gradient using
[Equation]
(36) a technique also known as the adjoint-state method.
Note that there is no need to compute the Hessian explicitly, but one can use automatic differentiation to compute the vector-Hessian product λ H θ and utilize iterative linear solvers like GMRES (Saad & Schultz, 1986) to solve the linear system.
As a ﬁnal note, we check the accuracy of our gradient computation using the ﬁnite-difference convergence test based on Taylor expansion:
[Equation] [Equation]
where δ x is a random vector with a unit (cid:96) 2 norm, and ∇ f ( x ) is computed using our custom gradient.
We here deﬁne f as a composite function that maps the (vector) output of an forward operator to a scalar, e.g. , f ( x ) = (cid:107) op forward ( x ) (cid:107) 22 .
Once we decrease (cid:15) , we should see that the error term f ( x + (cid:15)δ x ) − f ( x ) − (cid:15) ∇ f ( x ) (cid:62) δ x decreases at a rate of the second order.
All our layers passed this test, as the example shown in Fig.
This test should be conducted in double precision.
H O RTHOGONAL REPARAMETERIZATION
We also propose to reparameterize the latent vector z using an orthogonal matrix R :
[Equation]
where v is ﬁxed during an inversion, and we treat R as the parameter instead.
There are various ways to parameterized an orthogonal matrix.
We choose the Cayley parameterization:
[Equation]
, where W = − W
[Equation]
39) one of the best reported in Liu et al. (2021).
Therefore, the DGM-regularized inversion using orthogonal reparameterization is
[Equation]
The speciﬁc reparameterization schemes for StyleGAN2 and Glow are the same as in App.
D. For style vectors, we use the full dimension of 512 because of the observation described in Table 6.
For
Glow and noise maps in StyleGAN2, we use patch sizes 3 × 8 × 8 and 1 × 8 × 8 , respectively, to save computation time and memory.
I G AUSSIAN TYPICAL SET
The Gaussian annulus theorem states that a high-dimensional standard Gaussian distribution has most of its probability mass concentrated within an annulus area around a high-dimensional sphere:
Theorem 1 (Gaussian Annulus Theorem (Blum et al., 2020))
For an n-dimensional standard
Gaussian, for any β ≤ √ n , all but at most 3 e − cβ 2 of the probability mass lies within the annulus √ n − β ≤ | x | ≤ √ n + β , where c is a ﬁxed positive constant.
This theorem gives a necessary geometric condition of typical samples from a high-dimensional standard Gaussian.
However, the converse is not true: not all vectors whose (cid:96) 2 norm is √ n − β ≤ | x | ≤ √ n + β are typical examples from the standard Gaussian distribution.
As a result, a DGM does not necessarily map a latent vector staying within the Gaussian annulus geometrically to a plausible image, which is demonstrated in Fig.
The formal deﬁnition of a typical set is as follows.
Deﬁnition 1 (Cover & Thomas (2012))
Let p X
( x ) be a distribution whose support is X .
The typical set A ( n ) (cid:15) is deﬁned as the set of sequences ( x 1 , x 2 , · · · , x n ) ∈
X n , x i ∼ p X , that satisfy
[Equation]
where H [ X ] is the entropy of random variable X .
Now a random vector x ∈ R n ∼ N ( 0 , σ 2 I ) can be factorized as i.i.d.
random variables that are distributed as N (0 , σ 2 ) .
Therefore, we can regard x as an i.i.d.
sequence and give the following deﬁnition:
Deﬁnition 2 (Gaussian Typical Set)
A Gaussian typical set is the typical set A ( n ) (cid:15) of x ∈ R n ∼
N ( 0 , σ 2 I ) .
The following theorem guarantees that a typical sample from x ∈ R n ∼ N ( 0 , σ 2 I ) resides in the
Gaussian typical set with very high probability.
Theorem 2 (Cover & Thomas (2012))
For every (cid:15) > 0 , the typical set has probability P (cid:16) A ( n ) (cid:15) (cid:17) > 1 − (cid:15) with a sufﬁciently large dimension n .
This theorem is a direct application of the asymptotic equipartition property (AEP), which is based on the i.i.d.
assumption of sequence entries and the weak law of large numbers.
Similar to the
Gaussian annulus theorem, Theorem 2 depicts a geometric property of the Gaussian typical set: it is concentrated in an annulus near a shell of radius σ √ n , which can be directly veriﬁed by deﬁnition.
Of course, equivalently, if a vector whose (cid:96) 2 norm signiﬁcantly deviates from √ n cannot be a typical sample from a standard Gaussian.
However, one cannot assert that a vector is sampled from the
Gaussian typical set by only checking its norm.
J M ISCELLANEOUS TOPICS
Original
Observed β = 0 .β
=
.β
= 1 .β =
[Figure]
Ours
J. I NVERSION USING G LOW
In the Glow-regularization framework, we solve the inverse problem by ﬁnding the maximum a posteriori (MAP) estimate from
[Equation]
where M denotes the parameter space.
The probability density p M introduces our a priori knowledge and is represented by a normalizing ﬂow (Glow) g θ , which is a differentiable invertible mapping between two distributions, parameterized by neural network parameters θ : m =
g θ ( z ) , where z is the latent vector.
After training, the log probability density of a given model m is
[Equation]
where Z stands for the latent space.
When we use the trained network in inversion, we freeze the network weights; hence, we drop θ in g θ hereafter in our notation.
Therefore, Eq. 3 for Glowregularized inversion is
[Equation]
The new regularization term in Eq.
3 is R (cid:48)
( z ) = − β (cid:0) log p Z ( z ) − log (cid:12)(cid:12) det J g
( z ) (cid:12)(cid:12)(cid:1) , and we have shown the results with different β s in Fig.
20 and Table 7.
J. D UALITY OF KL DIVERGENCE
As also shown in Papamakarios et al.
(2017), the KL-divergence between two distributions does not change under a differentiable invertible transformation, so
[Equation]
where p ∗ M is the target distribution in the physical parameter space, and p ∗ Z is the corresponding latent-space distribution under the normalizing ﬂow.
This means that minimizing the forward KL divergence in the M domain or physical parameter space is equivalent to minimizing the reverse
KL-divergence in the Z domain or the latent space.
This fact and Theorem 2 imply that a well-trained normalizing ﬂow maps samples from the target distribution into the Gaussian typical set with very high probability and vice versa.
J. I NVARIANCE OF KL- DIVERGENCE AND MULTI - INFORMATION
For the completeness of this paper, we brieﬂy prove the key properties of KL-divergence and multi-information used in the Gaussianization framework.
First, we prove that the KL-divergence is invariant under differentiable bijections.
We write the deﬁnition of the KL-divergence between distribution p x and q x of random vector x as
[Equation]
Suppose there is a differentiable and invertible transformation T that maps x to u : u = T ( x ) .
Then, we know the PDF under change of variable is
[Equation]
where we deﬁne the Jacobian matrix J T ( x ) =
∂
u ∂ x ( x ) .
Therefore, we have
[Equation]
Then, we show that the multi-information does not change under component-wise invertible differentiable transformations.
The multi-information is deﬁned by the KL-divergence between the joint distribution p ( x ) and the marginals:
[Equation]
We deﬁne the bijection u
= T ( x ) speciﬁcally with invertible differentiable transformations for each component T i ( x i ) =
u i .
Then, using the KL-divergence invariance we just proved and that (cid:81)
Dj p j x ( x j ) =
(cid:81)
Dj p j u ( u j ) | det J T ( x ) | , we get I ( x ) =
I ( u ) .
J. T HE ROLLING OPERATION
After applying one set of the Gaussianization layers to the latent tensor, it might be desirable to apply the layers to a different set of non-overlapping patches from the latent tensor again.
The rolling operation shifts the latent tensor in the horizontal and vertical directions by half of the patch size w before patch extractions.
The values at the boundaries are wrapped around to the opposite sides.
One can, for example, use the torch.roll command to implement this functionality.
K A DDITIONAL DISCUSSIONS
[Figure]
We point out the limitations and potential improvements of this work.
First, requiring latent vectors to traverse within the Gaussian distribution/typical set means that the statistics of the training dataset will dominate the results.
Fig.
21 shows that our method cannot restore the eyeglasses because it is not a signiﬁcant feature within our training data.
One needs to pay particular attention to training dataset construction to ensure that it represents the typical features to be investigated.
One should also interpret their results with this caveat in mind.
On the other hand, one may combine
Gaussianization layers and latent space manipulation techniques (Shen et al., 2020) for a guided inversion, e.g. , requiring that the inverted images should have eyeglasses.
Second, we only use one set of Gaussianization layers.
One can further improve the performance by using additional sets of layers, with possible different partition schemes, such as using the rolling operation (App. J.4) to shift the patch extraction locations.
Third, since the deep generative models are highly nonlinear, the results may get stuck in local minima.
We mitigate this problem by starting inversion from multiple randomly initialized latent vectors.
However, this technique is computationally expensive.
It is an important research direction to improve the initialization process.
The Gaussianization layers are based on optimization and ﬁxed-point iterations, which can be computationally expensive.
To study possible simpliﬁcations, we estimated the forward and backward time for various combinations of sub-layers using 100 repeated experiments (Table 11).
The computation was performed using one Nvidia V100 GPU and an Intel Xeon Platinum 8276 CPU.
We always turn on the ICA layer (together with whitening) since the ablation studies show that it contributes the most to performance increase (Table 2 and Table 9).
Combining results from those ablation studies, we can see that the option ICA+YJ strikes a good balance between performance and computation time.
It has a shorter runtime than ICA+Lambert and ICA+YJ+Lambert.
At the same time, it gives comparable and sometimes the best results.
On the other hand, the additional computation time from Gaussianization layers is well justiﬁed by the performance increase, especially if the physics simulation is much more time-consuming.
For example, the eikonal solver for the experimental setup in our study has a forward runtime of 0.293 ± 0.013 seconds and a backward runtime of 5.933 ± 0.
seconds, respectively.
In addition, our implementation can also be greatly improved by more efﬁcient utilization of GPUs and rewriting bottleneck components using high-performance codes.
[Table]
We also provide insights into why Gaussianization layers perform better than other methods.
First, the spherical constraint only relies on the geometric property of the Gaussian typical set (App. I).
However, as shown in Fig.
[Table]
it is insufﬁcient to only constrain the norm of latent vectors in inversion.
Gaussianization is necessary to ensure inverted results are in the range of DGMs.
Second, the orthogonal reparameterization generally underperforms the Gaussianization layers.
Our interpretation is as follows: rather than actively destroying latent tensor patterns like the Gaussianization layers do, orthogonal reparameterization cannot prevent permuting the ﬁxed typical Gaussian latent tensors into ones with spatial patterns.
In addition, we have shown in Fig.
10 that the whitening transformation alone is ineffective in keeping the DGM outputs within the range.
We thus interpret that reducing higher-order dependencies is essential.
Therefore, the noise regularization and the whitening layers yield less satisfying results than the Gaussianization layers.
We believe this study can beneﬁt both the broader scientiﬁc community and the general public.
However, we point out that the MRI and eikonal tomography experiments are purely synthetic and numerical.
They do not fully reﬂect realistic medical imaging conﬁgurations.
One should be cautious about applying our techniques to real data and interpreting the results.