Intermediate Layer Optimization for Inverse Problems using Deep Generative Models
Abstract
We propose Intermediate Layer Optimization, a novel optimization algorithm for solving inverse problems with deep generative models.
Instead of optimizing only over the initial latent code, we progressively change the input layer we optimize over, obtaining successively more expressive generators.
We also experiment with different loss functions and utilize a perceptual loss combined with standard mean squared error.
We empirically show that our approach outperforms the state-of-the-art inversion methods introduced in StyleGAN2 and PULSE.
Introduction
We study how deep generators can be used as priors to solve inverse problems like inpainting, superresolution and denoising.
We focus on unsupervised image reconstruction techniques that rely on a pre-trained generator, building on the general framework introduced by Bora et al.
[1] (see also [2] for an overview).
The central optimization problem that appears in unsupervised image reconstruction is the inversion of a deep generative model, i.e. ﬁnding a latent code that explains the measurements.
This can be performed for different generators, e.g. DCGAN or more recently the powerful StyleGAN2 [3, 4] as shown in the excellent results obtained by PULSE [5].
Unfortunately, inverting a generator with even 4 layers is NP-hard [6].
Prior work [1] used gradient descent to minimize the MSE and showed good empirical performance for numerous inverse problems including inpainting and compressed sensing with random Gaussian measurements using DCGAN.
This does not work as well for deeper generators e.g. BigGAN as discussed in [7].
PULSE [5] showed signiﬁcantly better results speciﬁcally for super-resolution by improving the latent space optimization and using the powerful generator of StyleGAN2.
We propose a novel optimization method for solving general inverse problems using a technique we call Intermediate Layer Optimization (ILO).
Our method adaptively changes which layer is optimized, moving from the initial latent code to intermediate layers closer to the pixels.
By optimizing intermediate layers we expand the range of the generator to better satisfy the measurements.
This has to be done very carefully since intermediate layers can produce non-realistic images, which has to be carefully regularized.
We experiment with different loss functions and utilize a perceptual loss combined with standard mean squared error.
One beneﬁt of our approach is that training biases seem to be mitigated, as shown in our results.
Background: Inversion is the problem of projecting an image on the range of a generator .
Given a real image x ∈ R n and a generator G ( z ) , the goal of inversion is to ﬁnd the latent code z ∗ ∈ R k , so that G ( z ∗ ) ∈ R n approximates the given image x as well as possible.
The canonical way to perform inversion is to use gradient descent to solve the optimization problem:
[Equation]
where z is the latent code, x is the given image, dist ( · , · ) is a distance function that measures how close the generated image is to the given image and R ( z ) is a regularization term that forces the latent code to stay close to the original distribution.
Several papers have independently introduced this formulation for the inversion [8, 1, 9, 3].
Recent work has generalized this framework to solve more general inverse problems [1, 9, 10, 5, 2].
The general optimization problem is:
[Equation]
where f is a forward operator that produces measurements from the real data.
For example, for inpainting f is a masking operator that hides certain pixels of a given image.
MSE has been the most common loss used for solving inverse problems but we show how more careful combinations with perceptual loss functions improves performance.
Using Perceptual Loss: The LPIPS [11] distance metric has been shown to be effective at measuring perceptual differences between two images.
This makes it a natural candidate for inversion since the goal is to match the generated image to the observed image.
We show in the experiments that incorporating the LPIPS loss function yields superior image reconstructions.
The observed image in inpainting is a partially observed image - an image distortion that LPIPS was not trained for.
To address this, we minimize the perceptual distance between the generated image and the superimposed reconstruction (see equation 3).
In comparison to MSE, we show that a weighted combination of
LPIPS and MSE is more effective at image reconstructions under the downsampling and additive noise degradation operator.
Intermediate Layer Optimization: We show how to adapt image inversion to perform image inpainting by searching for latent vectors that match the observed pixels.
Our loss function for a given image I is deﬁned as:
[Equation]
where M is a binary mask, (cid:12) denotes element-wise multiplication, and M C is the complementary mask, i.e. M C ( i, j )
= 1 ( M ( i, j ) = 0) .
Intermediate Layer Optimization obtains ﬁne-grained image reconstructions.
ILO is motivated by the problem of inversion, for which we argue that the optimization scheme of Equation 3 can be further improved.
For the problem of inpainting, f is a masking operator that we use to force the model to match only the observable pixels of the given image I .
Consider now the case in which the model is observing only the top half of a human face and it is asked to complete it.
Solving the optimization problem of Equation 3 usually leads to realistic faces that can not fully match the observed pixels.
In other words, super-imposing the original image and the completion leads to observable visual discontinuities.
The cause of the problem is that the generator may be unable to produce an image that adequately matches the observations.
We therefore need a reconstruction that respects the boundary conditions and matches the observed pixels.
ILO is designed to solve exactly this problem by expanding the range of the generator.
We gradually optimize consecutive layers of the generator to introduce more ﬂexibility to match observed pixels.
Given a deep generator
G = ( g n ◦ g n − 1 ... ◦ g 2 ◦ g 1 )( z 1 ) , the algorithm runs in rounds, where in each round the initial layer is discarded.
At the ﬁrst round, the algorithm optimizes over the input latent codes.
At the second round, we remove the ﬁrst layer from the generator and we optimize over ( g n ◦ g n − 1 ... ◦ g 2 )( z 2 ) , initializing z 2 = g 1 ( z 1 ) to stay near the manifold of realistic images before running gradient descent.
We repeat this process of removing the initial layer and initializing values from the previous round until the MSE loss becomes very small.
There is a trade-off between the ability to match observed pixels and generation of realistic images.
Speciﬁcally, if we optimize only over the ﬁrst latent code layer (as previously done e.g. in [1] and
PULSE [5]), we obtain realistic faces that fail to perfectly match the boundary conditions .
On the contrary, if we start optimizing only the ﬁnal layers near the pixels, the generated images do not resemble human faces.
However, if we progressively move from the early layers to later layers of the network, we add more ﬂexibility to match the boundary conditions and we are also constrained to move in the manifold of realistic reconstructions.
We switch to the next layer when the loss function ﬂattens, which can be done in an unsupervised manner.
As we also show in the Experiments, this simple heuristic works surprisingly well.
Stochastic Noise Addition for Denoising:
For denoising, simply inverting a noisy high-resolution image creates grainy reconstructions due to the expressive power of StyleGAN2.
We address this issue by minimizing the loss function for a given noisy image I :
[Equation]
where N (0 , σ 2 ) is a randomly generated noise added to the generated image at each optimization step.
Our method requires knowing the variance of the Gaussian noise added to the original image.
Experimental Evaluation
Our experiments are performed using the state of the art StyleGAN2 [4], a ramped-down learning rate scheduler with initial learning rate of 0 .
1 and the regularization term R ( z ) deﬁned as the geodesic loss between the latent vectors.
We optimize over the 18 different latent vectors and the ﬁrst 5 noises.
We experiment under three different settings: inpainting, denoising, super-resolution.
For each task, we use the ideas introduced in the previous section and we compare against the best previously proposed techniques that are based on a formulation of the inversion problem for the image reconstruction task.
Namely, for inpainting we use ILO and the formulation of Equation 3, for denoising we use Stochastic
Noise Addition and the formulation of Equation 4.
For all experiments we use a combination of
LPIPS [11] and MSE.
We attempt reconstruction of high-resolution (1024 × 1024) real and generated images.
[Figure]
Inpainting.
We ﬁrst show results for the problem of inpainting and compare our framework consisting of ILO and the formulation of Equation 3 against the natural extension of PULSE [5] for inpainting.
The results are shown in Figure 1.
The ﬁrst four images have masks that were chosen to remove important features and structural integrity of the facial structure.
As shown, the combination of ILO and our new loss function consistently gives better completions compared to the baselines.
Even when the observations are extremely sparse (e.g. in the last row we observe only 5% of the pixels), our solution produces accurate reconstructions of the original image.
To further validate the ILO innovation, we measure the MSE loss between the real and the generated images when we use: (i) just PULSE, (ii) ILO applied for different number of rounds.
The results are shown in Figure 2.
As shown in the plot, we switch to the next layer in ILO when the MSE error ﬂattens, achieving lower
MSE.
Denoising.
For denoising, we use the Stochastic Addition of Noise framework we introduced earlier together with the ILO method.
Previously proposed techniques perform poorly for denoising and hence we also compare with BM3D [12] a standard denoising method.
Results are shown in
Figure 3 for zero-mean additive Gaussian noise with σ = 25 .
[Figure]
Super-resolution.
Finally, we compare our ILO method for super-resolution with PULSE [5].
Results are shown in Figure 3 (right-part).
As shown in the Figure, ILO leads to successful reconstructions.
Further, ILO seems to lead to less biased reconstructions since it expands the expressive range of the generator by optimizing intermediate layers.