Maximum a Posteriori on a Submanifold: a General Image Restoration Method with GAN

1 INTRODUCTION . Image restoration has been researched for many years , but in a case-by-case way ( Park et al. , 2003 ; Mairal et al. , 2008 ; Guillemot & Le Meur , 2014 ) . Almost all image restoration algorithms are only designed for certain type of images or degradation . This research paradigm has some obvious disadvantages . It is exhausting to invent new algorithms or train new models for slightly different situations . Even if we can , those specialized solutions are not so elegant , because they are very unlike one another even though the problems they focus on are fundamentally so similar . It is worth noting that any image degradation process can be abstracted as a many-to-one function . More specifically , for any given degradation process , one degraded image could be degraded from many possible original images . From that point of view , we propose a general method for various image restoration problems , such as denoising , deblurring , super-resolution and inpainting . Our algorithm chooses the most probable original image from all those possible original images , and uses it as the restoration of the given degraded image . To be more precise , the general image restoration is formulated as a constrained optimization problem . Its objective is to maximize a posteriori probability of latent variables , and its constraint is that the image generated by these latent variables must be the same as the degraded image . Recent progress of density estimation techniques makes our algorithm possible . In the field of image generation , Generative Adversarial Networks ( GANs ) make a huge success in recent years ( Goodfellow et al. , 2014 ; Radford et al. , 2015 ) . As research continues , images generated by GANs become more and more realistic and clear , and training procedure of GANs become more and more stable ( Salimans et al. , 2016 ; Arjovsky et al. , 2017 ) . Besides being an image generation technique , GANs can also be used for density estimation . The generator part of a GAN is an implicit probability distribution model , and it will converge to a good estimator of the data distribution after training . In this work , we solve the inference problem with the probability density estimated by a GAN . Figure 1 provides an illustration of how our image restoration method works . There are four dashed boxes from left to right in Figure 1 , corresponding to four different phases of image capture and restoration process . Images in the first dashed box are original images , which are clear and undegraded . These images undergo a series of degradation in the second dashed box , and then are captured by our camera . In the image restoration process , we hope to estimate the original images with the degraded images we captured . As we pointed out before , every degraded image could be degraded from many possible original images . To be more precise , there is a particular subset of the original image manifold for any degraded image , and all image samples on the submanifold could be degraded to the given degraded image . Images in the third dashed box are those samples on the submanifold , and they are arranged in ascending order of log-likelihood from left to right . Images marked by yellow boxes are samples with the highest log-likelihood in their group , and they are placed in the last dashed box as restoration outputs . Overall , the contributions of this work are mainly in two aspects : 1 . We propose a general method for various image restoration problems . In the method , we explicitly use density information estimated by a GAN , an implicit model ; and we directly solve the image restoration problem , an inference problem , with a GAN , a generative model . To the best of our knowledge , our work is the first to do those two things . 2 . We propose a new algorithm to solve the optimization problem in our method . The new optimization algorithm is a first-order iterative algorithm for constrained problems , and it works well even for problems with highly nonlinear objective and constraints . These features make it especially suited to neural network related constrained optimization problems . 2 RELATED WORK . The most similar work to ours is proposed in Yeh et al . ( 2017 ) . They propose a semantic image inpainting method , which can generate missing content with a trained GAN . They search in the latent space of the trained GAN for the image which is closest to the corrupted image , and use the discriminator loss of the trained GAN as an indicator of how realistic their restoration is . Their motivation is similar to our work , but unfortunately , there is a major theoretical flaw in their method . Goodfellow et al . ( 2014 ) already prove that the discriminator is unable to identify how realistic an input is after several steps of training , if the GAN has enough capacity . During the training , the information of the data distribution gradually transfer from the discriminator to the generator . Ideally , the generator will have all the information of the data distribution while the discriminator will have none . That is why we use the generator of a trained GAN as an implicit probability density model in our method . Another difference between their work and ours is that they only focus on image inpainting problem , while our method applies to various image restoration problems . The maximum a posteriori ( MAP ) has existed for a long time as a classic estimation method ( Campisi & Egiazarian , 2016 ) . But before GANs , people do not have a probability density model which is good enough to describe the distribution of images . After GANs make a huge success in image generation , researchers start to use them in image restoration tasks to get more realistic results ( Isola et al. , 2017 ; Bousmalis et al. , 2017 ) . Ledig et al . ( 2017 ) and Sønderby et al . ( 2016 ) try to use the MAP estimation on GANs to solve image super-resolution problem . However , they only use the MAP estimation implicitly and indirectly , while our method use it explicitly and directly . We suspect that all methods do implicit MAP estimation on GANs would require redesigning or retraining when the image restoration task changes , and this makes implicit methods not as general as our explicit method . Ulyanov et al . ( 2017 ) is another work which is seemingly similar to ours , but they are actually quite different . Their work uses a randomly-initialized neural network as an image prior to solve various image restoration problems . The prior in their method is elaborate , neural network related but still handcrafted , while in our method the prior is learned from data . So our data-driven prior has better adaptability to specific image distribution . 3 MAXIMUM A POSTERIORI ON A SUBMANIFOLD . 3.1 FORMULATION . Consider a general image degradation model as follows , x̃ = F ( x , Ω ) ( 1 ) where x , x̃ , and Ω represent the original image , the degraded image , and the parameters of the degradation model , respectively . The image degradation function F is a deterministic function . That means , given an original image x and a particular set of parameters Ω , the image degradation model will always produce the same degraded image x̃ . Our goal is to get a reasonable estimate of x with given x̃ and F . In this paper , we use the maximum a posteriori probability ( MAP ) estimate of x as the restoration of x̃ . Compared to MSE-based method , MAP estimate of x is perceptually more convincing . We can perform inference by maximizing the posterior p ( x , Ω|x̃ ) : { x̂ , Ω̂ } =argmax x , Ω p ( x , Ω|x̃ ) = argmax x , Ω p ( x̃|x , Ω ) p ( x|Ω ) p ( Ω ) p ( x̃ ) ( 2 ) where x̂ and Ω̂ represent MAP estimate of x andΩ . Note that p ( x̃ ) is always positive and does not depend on x andΩ , and typically we assume that x andΩ are independent . Therefore , { x̂ , Ω̂ } =argmax x , Ω p ( x̃|x , Ω ) p ( x ) p ( Ω ) ( 3 ) Note that x̃ = F ( x , Ω ) is a deterministic function , i.e. , p ( x̃|x , Ω ) = δ ( x̃− F ( x , Ω ) ) . Therefore , the estimation is equivalent to { x̂ , Ω̂ } = argmax x , Ω p ( x ) p ( Ω ) s.t . ‖x̃− F ( x , Ω ) ‖ = 0 ( 4 ) Here we write p ( x ) more specifically as pr ( x ) , which stand for the probability density of real data distribution . We can estimate pr ( x ) with the generator part of a trained GAN , which is an implicit probability distribution model with distribution pG ( x ) . The trained generator G represents a mapping from latent space of z to data distribution of original image x , i.e. , pr ( x ) = pG ( x ) , and pG ( x ) is a probability density function implicitly defined by x = G ( z ) , where z is typically sampled from some simple distribution , such as the uniform distribution or the normal distribution . Assuming G : Rn → Rm is an injective function , the estimation is equivalent to { ẑ , Ω̂ } = argmax z , Ω pG ( G ( z ) ) p ( Ω ) s.t . ‖x̃− F ( G ( z ) , Ω ) ‖ = 0 ( 5 ) and x̂ = G ( ẑ ) ( 6 ) Generally the dimension of vector space of z is far lower than the dimension of vector space of x . Note that pG ( x ) is nonnegative if and only if x is on the low dimensional manifoldM defined by x = G ( z ) , we can replace the probability density on the original space pG ( G ( z ) ) in Eq . ( 5 ) by the probability density on the manifold pM ( z ) , and end up with the same estimation result ẑ . According to Pennec ( 2004 ) , the probability density on the manifold can be calculated by pM ( z ) = p ( z ) √ detGram ( ∂G∂z1 , . . . , ∂G ∂zn ) ( 7 ) where Gram represents the Gram matrix , and √ detGram ( ∂G∂z1 , . . . , ∂G ∂zn ) is the volume of the parallelotope spanned by the vectors ( ∂G∂z1 , . . . , ∂G ∂zn ) , so the square root of the Gram determinant can serve as a local scale factor . It has an effect similar to the Jacobian determinant , but we can only use the Gram determinant here because G is a function from Rn to Rm , and generally n is much less than m. The Gram matrix can be simply calculated byGram ( ∂G∂z1 , . . . , ∂G ∂zn ) = V TV , where V is anm×n matrix , whose entries are given by V ij = ∂xi∂zj . Therefore , Eq . ( 5 ) is equivalent to { ẑ , Ω̂ } = argmax z , Ω p ( z ) p ( Ω ) √ detV TV s.t . ‖x̃− F ( G ( z ) , Ω ) ‖ = 0 ( 8 ) To solve the estimation problem efficiently , we represent probabilities in Eq . ( 8 ) in logarithmic space , i.e. , { ẑ , Ω̂ } = argmax z , Ω − 1 2 log detV TV + log p ( z ) + log p ( Ω ) s.t . ‖x̃− F ( G ( z ) , Ω ) ‖ = 0 ( 9 ) Matrix V TV is a positive-definite matrix , so we can use Cholesky decomposition to calculate log detV TV efficiently , i.e. , log detV TV = 2 tr ( log ( chol ( V TV ) ) ) ( 10 ) Finally we deduce a set of expressions which can be calculated directly , and their final outcome x̂ is the restored image we want , i.e. , { ẑ , Ω̂ } = argmax z , Ω − tr ( log ( chol ( V TV ) ) ) + log p ( z ) + log p ( Ω ) s.t . ‖x̃− F ( G ( z ) , Ω ) ‖ = 0 ( 11 ) and x̂ = G ( ẑ ) ( 12 ) Note that ( G ( z ) , Ω ) form a low dimensional manifold which is embedded in the space of ( x , Ω ) , and the feasible solutions of Eq . ( 11 ) is on a subset of the manifold , which is defined by ‖x̃ − F ( G ( z ) , Ω ) ‖ = 0 . So our method basically makes a MAP estimate on a submanifold . Figure 2 is a toy example to show how our method works in a very visible way . Suppose there is a grayscale original image x , which has only three pixels . Then it is downsampled to only one pixel during the image capture process , and our task is to estimate xwith the one pixel image we captured . Suppose we have trained a GAN as an implicit model of data distribution of x . More specifically , the generator of the trained GAN represents a mapping from its input noise z to data distribution of x . The left part of Figure 2 describes the two dimensional latent space of z . We use the saturation of orange color to represent probability density level , i.e. , a thicker orange color means a higher probability density . So the uniform orange color in the latent space means that the input noise z is sampled from a uniform distribution . Then the two dimensional vector z is mapped to three dimensional space of image x by the generator of the trained GAN , and the big orange square in the latent space of z is transformed into a twisted torus in the three dimensional data space of x , which is described in the right part of Figure 2 . Some areas in space of z expand during the transformation , while other areas shrink . We can find this out by comparing the red and blue quadrilateral between the latent and data space . Therefore , the probability density on the torus is no longer uniform . The orange colors of the expanded areas become lighter , and the colors of the shrunken areas become thicker . Quantitatively speaking , the square root of the Gram determinant in Eq . ( 7 ) is the local area scale factor of the mapping , and its inverse , of course , is the local density scale factor . The pale yellow plane in the data space represents the constraint in the toy example . All points on the plane would exactly be downsampled to the one pixel image we captured . So the intersection curve of the plane and the torus is the submanifold we are looking for , and that white curve is the feasible set of the toy problem . In this problem , p ( z ) is a constant in the domain , and degradation parameters Ω does not exist at all . According to Eq . ( 8 ) , what we need to do is to maximize the inverse of the square root of the Gram determinant on the submanifold . In other words , the point with the thickest orange color on the intersection curve is the restored image x̂ , the MAP estimate on the submanifold . We can find out that the method is both intuitive and rational for this toy example .

Summary: The authors propose a method for image restoration, where the restored image is the MAP estimate. A pretrained GAN is utilized to approximate the prior distribution of the noise-free images. Then, the likelihood induces a constraint which is based on the degradation function. In particular, the method tries to find the latent point for which the GAN generates the image, which if gets degraded will match the given degraded image. Also, an optimization algorithm is presented that solves the proposed constrained optimization problem.