Diverse Image Inpainting with Bidirectional and Autoregressive Transformers

Traditional methods address image inpainting challenge through either image diffusion [5, 1] or using image patches [3, 15, 9]. However, diffusion-based methods often introduce diffusion-related blurs, which tends to fail while the missing or corrupted image regions are large [6, 2, 4]. Patch-based methods can work well for the inpainting of stationary background with repeating patterns. However, they struggle in completing large missing regions of complex scenes as the patch-based approach relies heavily on patch-wise matching of low-level features.

Generative adversarial networks (GANs) [14] have been investigated extensively in various image synthesise tasks such as image translation [31, 36, 19, 53, 51, 56, 54], image editing [48, 44, 43], image composition [24, 59, 58, 52, 57, 55], etc. Specifically for image inpainting, Pathak et al. [32] first apply adversarial learning to the image inpainting task. To further improve the adversarial learning within local regions, Iizuka et al. [18] introduce an extra local discriminator to enforce the local consistency. As the local discriminator uses fully-connected layers and can only deal with missing regions of fixed shapes, Yu et al. [49] inherit the discriminator from PatchGAN [19] due to its great success in image translation. Yan et al. [46] propose patch-swap to make use of distant feature patches for the better inpainting quality. Liu et al. [25] design partial convolutions to alleviate the negative influence of the masked regions. Yu et al. [50] present a novel free-form image inpainting system based on an end-to-end generative network with gated convolutions. To generate reasonable structures and realistic textures, Nazeri et al. [30] and Xu et al. [45] utilize edge maps as structural guidance for image inpainting, and Ren et al. [35] instead propose to use edge-preserved smooth images as structural guidance. Liu et al. [27] propose feature equalizations to improve the consistency between structures and textures. As aforementioned methods focus on reconstructing the ground truth instead of generating pluralistic inpainting, they are constraint to generate a deterministic inpainting image for each incomplete image.

To achieve pluralistic image inpainting with plausible filling contents, Zheng et al. [64] propose a VAE-based network with a dual pipeline, which trades off between reconstructing ground truth and maintaining the diversity of the inpainting results. Similarly, Zhao et al. [61] propose a VAE-based model and leverage a reference image to improve the diversity. Although the above methods achieve certain diversities to some extent, completion quality of VAE-based methods is limited due to variational training. Recently, Zhao et al. [62] propose a co-modulated GAN to incorporate the image condition and the stochastic generation of unconditional generative model for diverse inpainting. Peng et al. [33] introduce a hierarchical vector quantized variational auto-encoder (VQ-VAE) to quantize the context representation and archive diverse structure generation in an autoregressive way. Sharing a similar framework with us, Wan et al. [41] propose to apply transformer for diverse structure generation using the objective of BERT [11]. In contrast, we propose a novel Bidirectional and Autoregressive Transformer (BAT) which inherits the advantages of autoregressive models and bidirectional models and achieves superior image inpainting performance.

To relieve the pressure of quadratic complexity incurred in transformer, we adapt the low-resolution image with the size of $32\times 32\times 3$ to represent the coarse structure. As the autoregressive generation requires discrete distribution, the pixel value should be treated as classes to the model, which leads to the dimensionatliy of $256^{3}$ for each pixel of the 8-bit RGB images. Following Chen et al. [8], a color palette is applied to further reduce the dimensionality to $512$ while faithfully preserving the main structure of original images, which is generated by $k$-means clustering of RGB pixel values with $k$=512 from ImageNet [10] dataset.

Autoregressive (AR) modeling and masked language modeling (MLM) in BERT [11] are two representative objectives for exploiting large language corpora in language processing tasks. Given a discrete sequence $X=\left\{x_{1},x_{2},\dots,x_{L}\right\}$ where $L$ is the length of $X$, AR model is optimized by maximizing the unidirectional likelihood:

$$\log P(x;\theta)=\mathbb{E}_{X}\sum\limits_{t=1}^{L}\log P(x_{t}|X_{<t};\theta),$$

(1)

where $\theta$ is the parameters of the model. In contrast, MLM aims to reconstruct corrupted data with the masked positions $M=\left\{m_{1},m_{2},\dots,m_{K}\right\}$, where $K$ is the number of masked tokens. Each masked position of the corrupted data is indicated by a special token $[M]$ following BERT [11]. Denoting the masked tokens as $X_{M}$ and unmasked tokens as $X_{\setminus M}$, the objective of MLM can be formulated by:

$$\log P(X_{M}|X_{\setminus M};\theta)=\mathbb{E}_{X}\sum\limits_{m_{k}\in M}\log P(x_{m_{k}}|X_{\setminus M};\theta).$$

(2)

AR and MLM differ from two aspects as defined in Eqs. 1 and 2. The first aspect lies with output dependency, where MLM predicts the masked tokens separately and independently which may oversimplify the complex context dependency in the data [37]. As a comparison, AR factorizes the predicted tokens with the product rule, which establishes the output dependency and produces better predictions. The second aspect lies with context dependency, where AR is only conditioned on the tokens up to the current position (in a fixed order), while MLM has access to bidirectional contextual information. Therefore, MLM is more suitable for image inpainting as the missing or corrupted image regions often have arbitrary shapes with rich variation in the neighboring background.

We propose a novel Bidirectional and Autoregressive Transformer (BAT) that inherits the advantages of AR and MLM to achieve bidirectional context modeling and output dependency simultaneously. The training objective of the BAT is formulated by:

$$\mathcal{L}_{BAT}=\mathbb{E}_{X}\sum\limits_{m_{k}\in M}\log P(x_{m_{k}}|X_{\setminus M},M,X_{<m_{k}};\theta).$$

(3)

We first project all the tokens into a $d$-dimensional token embedding and add a learnable position embedding over the token embedding to preserve the positional information. Unlike XLNet [47] which randomly permutes the input sequence to capture the bidirectional context, we permute all unmasked tokens $X_{\setminus\Pi}$ in the front while maintaining the original order of the masked tokens for better predicting their positions. Moreover, the positional information of all masked tokens will be conditioned for better modeling of the full input sequence (e.g. the counts and positions of masked tokens in the sequence). The proposed BAT model is then adopted to predict the masked tokens as illustrated in Fig. 3.

As shown in Fig. 3, there is a masked sequence $X=\{x_{1},[M],[M]$ $,x_{4},[M]\}$ with length $L$ = 5 and positions $2,3,5$ being masked. After permutation and inserting the full mask tokens, we have the non-predicted tokens $(X_{\setminus M},M)=(x_{1},x_{4},[M],[M],[M])$ which provides the bidirectional context. For the predicted part, we have the input tokens $([M],x_{2},x_{3})$ to predict their corresponding next tokens i.e. $(x_{2},x_{3},x_{5})$. Here we use the mask token instead of $x_{1}$ to predict $x_{2}$ to encourage the leverage of positional information. We apply bidirectional modeling [11] to non-predicted tokens and autoregressive modeling to the predicted tokens to avoid future information leakage. For example, while predicting $x_{3}$, the model could attend to $x_{4}$ in non-predicted tokens and meanwhile the previously ‘predicted’ token $x_{2}$. Hence, we could capture bidirectional context and establish output dependency simultaneously with the proposed BAT.

In this work, we adapt GPT [34] as our network architecture. The network is a decoder-only transformer that consists of $\mathcal{N}$ stacked decoder blocks. Given an intermediate embedding $H^{n}$ at the $n$-th layer, the decoder block can be formulated by:

	$\displaystyle H^{n}$	$\displaystyle=H^{n}+\text{MA}(LN(H^{n}))$		(4)
	$\displaystyle H^{n+1}$	$\displaystyle=H^{n}+\text{MLP}(LN(H^{n})),$		(5)

where $MA$, $LN$ and $MLP$ stand for multihead self-attention, layer normalization, and fully-connected layers, respectively. For self-attention, we apply a customized mask to the $L\times L$ matrix of attention logits as illustrated in Fig. 3. At the final layer of the transformer, a learnable linear projection is employed to map $H^{\mathcal{N}}$ to logits, which parameterizes the conditional distribution for each pixel.

During inference, we follow the raster-scan order to predict each masked token bidirectionally and autoregressively. We adopt a top-$\mathcal{K}$ sampling strategy to randomly sample from the $\mathcal{K}$ most likely next words. The predicted token is then concatenated with the input sequence as conditions for the generation of next masked token. This process repeats iteratively until all the masked tokens are sampled. Finally, the generated discrete sequence can be converted back to the RGB values with the aforementioned color palette.

As the inpainting diversity can be achieved by sampling the reconstructed structures $S$, we take the advantages of efficiency and texture representation capacity of CNNs to learn a deterministic mapping between low-resolution structures $S$ and high-resolution completed image $I_{out}$. The texture generator thus utilizes CNN layers and adversarial training to up-sample the reconstructed structures and replenish high-fidelity texture details by leveraging the styles of the valid pixels of input image $I_{m}$. In particular, we employ two encoders to encode the low-resolution structures and input images into two high-level CNN representations of the same dimension. We then concatenate them together as the input of a few consecutive residual blocks with different dilation rates. Finally, a SPADE [31] generator is employed to incorporate the modulated style of input images and gradually up-sample the texture features to the target resolution. Meanwhile, all vanilla convolutions are replaced by gated convolution [50].

The training of the texture generator is driven by the combination of several losses including a reconstruction loss, an adversarial loss, and a perceptual loss. For clarity, we denote the texture generator as $G_{t}$, the ground truth as $I_{gt}$, and the completed image as $I_{out}$. Firstly, a reconstruction loss $\mathcal{L}_{rec}$ between $I_{out}$ and $I_{gt}$ can be measured as follows:

$\displaystyle\mathcal{L}_{rec}=||I_{out}-I_{gt}||_{1},$

Besides, a CNN-based discriminator $D$ together with an adversarial loss is employed to synthesize fine texture details. Specifically, the texture generator $G_{t}$ and discriminator $D$ are jointly trained with hinge loss [19], where the adversarial losses for the discriminator and generator are defined by:

	$\displaystyle\mathcal{L}_{adv}^{D}=\mathbb{E}_{I_{gt}}[\textit{ReLU}(1-D(I_{gt})]+\mathbb{E}_{I_{out}}[\textit{ReLU}(1+D(I_{out})]$
	$\displaystyle\mathcal{L}_{adv}^{G_{t}}=-\mathbb{E}_{I_{out}}[D(I_{out})],$

Next, we penalize the perceptual and semantic discrepancy via the perceptual loss [20] with a pretrained VGG-19 network:

	$\displaystyle\mathcal{L}_{perc}=\sum_{i}{\lambda_{i}||\Phi_{i}(I_{out})-\Phi_{i}(I_{gt})||_{1}}$
	$\displaystyle+\lambda_{l}||\Phi_{l}(I_{out})-\Phi_{l}(I_{gt})||_{2},$

where $\lambda_{i}$ are balancing weights, $\Phi_{i}$ is the activation of $i$-th layer of the VGG-19 model (including relu1_2, relu2_2, relu3_2, relu4_2 and relu5_2), $\Phi_{l}$ represents the activation maps of relu4_2 layer which mainly extracts semantic feature. The texture generator is trained by optimizing the combination of aforementioned losses:

$\displaystyle\mathcal{L}_{G_{t}}=\underset{G_{t}}{min}\underset{D}{max}(\lambda_{rec}\mathcal{L}_{rec}+\lambda_{adv}\mathcal{L}_{adv}^{G_{t}}+\lambda_{perc}\mathcal{L}_{perc}),$

where $\lambda_{rec}$, $\lambda_{adv}$, and $\lambda_{perc}$ are empirically set at 1.0, 1.0 and 0.2, respectively, in our implementation.

We conduct experiments over three public datasets that have different characteristics as listed:

• –

  CelebA-HQ [21]: It is a high-quality version of the human face dataset CelebA [28] with 30,000 aligned face images. We follow the split in [50] that produces 28,000 training images and 2,000 validation images, where 1,000 validation images are randomly sampled in evaluations.

• –

  Places2 [65]: It consists of more than 1.8M natural images of 365 different scenes. We adopt the same 800 images from validation set with [41] in evaluations.

• –

  Paris StreetView [32]: It is a collection of street view images in Paris, which contains 14,900 training images and 100 validation images.

We compare our method with a number of state-of-the-art methods as listed:

• –

  GC [50]: It is also known as DeepFill v2, a two-stage method that leverages gated convolutions.

• –

  EC [30]: It is a two-stage method that first predicts salient edges to guide the generation.

• –

  MEDFE [26]: It is a mutual encoder-decoder that treats features from deep and shallow layers as structures and textures of an input image.

• –

  PIC [64]: It is a probabilistically principled framework that leverages VAE to generate diverse image inpainting.

• –

  ICT [41]: It is a diverse inpainting framework that combine the merits of transformers and CNNs for high-fidelity image inpainting.

We perform evaluations by using five widely adopted evaluation metrics: 1) Fréchet Inception Score (FID) [17] that evaluates the perceptual quality by measuring the distribution distance between the synthesized images and real images; 2) mean $\ell_{1}$ error; 3) peak signal-to-noise ratio (PSNR); 4) structural similarity index (SSIM) [42] with a window size of 51; 5) Learned Perceptual Image Patch Similarity (LPIPS) [60] that evaluates the diversity of generated images. The average scores of LPIPS are calculated between random pairs of sampled inpainting results.

The proposed method is implemented in PyTorch. The network is trained using $256\times 256$ images with random irregular masks [25]. The diverse-structure generator and texture generator are trained using $256\times 256$ images with random irregular masks [25]. We train the diverse-structure generator with AdamW [29] with $\beta_{1}=0.9$, $\beta_{2}=0.95$ and learning rate of 3e-4 following [8]. For the texture generator, we use Adam optimizer [22] with $\beta_{1}=0$ and $\beta_{2}=0.9$, and set the learning rate at 1e-4 and 4e-4 for the generator and discriminators, respectively. Learning rate decay is applied for the training of both networks, and the experiments are conducted on 4 NVIDIA(R) Tesla(R) V100 GPU.