Rectified Wasserstein Generative Adversarial Networks for Perceptual Image Restoration

Wasserstein generative adversarial network (WGAN) has attracted great attention due to its solid mathematical background, i.e., to minimize the Wasserstein distance between the generated distribution and the distribution of interest. In WGAN, the Wasserstein distance is quantitatively evaluated by the discriminator, also known as the critic. The vanilla WGAN trained the critic with the simple Lipschitz condition, which was later shown less effective for modeling complex distributions, like the distribution of natural images. We try to improve the WGAN training by introducing pairwise constraint on the critic, oriented to image restoration tasks. In principle, pairwise constraint is to suggest the critic assign a higher rating to the original (real) image than to the restored (generated) image, as long as such a pair of images are available. We show that such pairwise constraint may be implemented by rectifying the gradients in WGAN training, which leads to the proposed rectified Wasserstein generative adversarial network (ReWaGAN). In addition, we build interesting connections between ReWaGAN and the perception-distortion tradeoff. We verify ReWaGAN on two representative image restoration tasks: single image super-resolution (4× and 8×) and compression artifact reduction, where our ReWaGAN not only beats the vanilla WGAN consistently, but also outperforms the state-of-the-art perceptual quality-oriented methods significantly. Our code and models are publicly available at https://github.com/mahaichuan/ReWaGAN.

1 Introduction

Modeling complex distributions, such as the distribution of natural images, has been a longstanding difficult problem and recently been tackled by using deep networks. Such deep networks include variational auto-encoder [1], normalizing flow [2], [3], and auto-regressive model [4]. Different from these networks that focus on maximizing the likelihood of training samples, generative adversarial network (GAN) [5] provides indirect modeling of the distribution of interest by training a generator. Indeed, GAN formulates the training of generator as an adversarial game by introducing an auxiliary network known as discriminator. The discriminator is trained to identify generated samples from real samples, while the generator is trained to fool the discriminator so that it cannot correctly distinguish. In theory, the objective of GAN is to minimize the difference between the distribution $P_\theta$ of generated samples and the distribution $P_{r}$ of real samples, by adjusting the parameters $\theta$ of the generator $g_\theta$. The difference between the two distributions $P_\theta$ and $P_{r}$ may be evaluated by the Jensen-Shannon divergence like in the first GAN (known as JSGAN for clarity) [5], the general $f$-divergence like in the $f$-GAN [6], and the Wasserstein-1 distance like in the Wasserstein GAN (WGAN) [7].