[Paper Review] Reduce Information Loss in Transformers for Pluralistic Image Inpainting

Qiankun Liu, Zhentao Tan+7 more|arXiv (Cornell University)|May 10, 2022

Generative Adversarial Networks and Image Synthesis8 citations

TL;DR

This paper proposes PUT, a novel transformer-based framework for pluralistic image inpainting that minimizes information loss by avoiding input downsampling and quantization. It uses a patch-based VQ-VAE (P-VQVAE) for non-overlapping feature encoding and an Un-Quantized Transformer (UQ-Transformer) that processes unquantized features while predicting from quantized tokens, achieving state-of-the-art fidelity, especially on large and complex masked regions.

ABSTRACT

Transformers have achieved great success in pluralistic image inpainting recently. However, we find existing transformer based solutions regard each pixel as a token, thus suffer from information loss issue from two aspects: 1) They downsample the input image into much lower resolutions for efficiency consideration, incurring information loss and extra misalignment for the boundaries of masked regions. 2) They quantize $256^3$ RGB pixels to a small number (such as 512) of quantized pixels. The indices of quantized pixels are used as tokens for the inputs and prediction targets of transformer. Although an extra CNN network is used to upsample and refine the low-resolution results, it is difficult to retrieve the lost information back.To keep input information as much as possible, we propose a new transformer based framework "PUT". Specifically, to avoid input downsampling while maintaining the computation efficiency, we design a patch-based auto-encoder P-VQVAE, where the encoder converts the masked image into non-overlapped patch tokens and the decoder recovers the masked regions from inpainted tokens while keeping the unmasked regions unchanged. To eliminate the information loss caused by quantization, an Un-Quantized Transformer (UQ-Transformer) is applied, which directly takes the features from P-VQVAE encoder as input without quantization and regards the quantized tokens only as prediction targets. Extensive experiments show that PUT greatly outperforms state-of-the-art methods on image fidelity, especially for large masked regions and complex large-scale datasets. Code is available at https://github.com/liuqk3/PUT

Motivation & Objective

• To address information loss in transformer-based pluralistic image inpainting caused by input downsampling and pixel quantization.
• To preserve high-resolution image details and spatial alignment, especially at masked region boundaries.
• To improve inpainting fidelity and diversity on large and complex regions, particularly in large-scale datasets like ImageNet.
• To design a training strategy for un-quantized transformers that still enables discrete token prediction via random quantization during training.
• To eliminate misalignment artifacts between high- and low-resolution representations common in prior methods.

Proposed method

• Introduces P-VQVAE, a patch-based, non-overlapping auto-encoder that encodes image patches into latent features without downsampling.
• Employs dual codebooks in P-VQVAE to separately represent masked and unmasked patches for improved feature disentanglement.
• Uses a decoder in P-VQVAE to reconstruct masked regions while preserving unmasked regions unchanged.
• Proposes UQ-Transformer that takes unquantized feature vectors as input but predicts quantized tokens as targets, avoiding quantization-induced information loss.
• Trains UQ-Transformer with randomly quantized input features during training to simulate inference conditions, enabling accurate generation without full quantization at inference.
• Maintains high-resolution input throughout the network, eliminating resolution misalignment between input and prediction.

Experimental results

Research questions

• RQ1 Can avoiding input downsampling and quantization improve image inpainting fidelity in transformer-based models?
• RQ2 How does processing unquantized features in a transformer affect the quality and diversity of generated image regions?
• RQ3 To what extent does the use of a non-overlapping patch design in the auto-encoder reduce interference between masked and unmasked regions?
• RQ4 Does random quantization during training of the UQ-Transformer effectively enable inference with discrete tokens while preserving feature fidelity?
• RQ5 How does PUT compare to existing methods in handling irregularly shaped and large masked regions?

Key findings

• PUT achieves state-of-the-art performance on FID, LPIPS, and FVD metrics across FFHQ, Places2, and ImageNet datasets, especially outperforming prior methods on large masked regions.
• The ablation study shows that PUT with dual codebooks (PUT) performs better than single codebook (PUT one) and non-quantized inference (PUT no_ref), confirming the benefit of structured codebook design.
• PUT outperforms PUT conv, which uses a CNN-based encoder with sliding window processing, demonstrating the superiority of non-overlapping patch encoding in avoiding masked-unmasked region interference.
• The model with random quantization during training (PUT) significantly outperforms the variant without it (PUT qua0), validating the effectiveness of the training strategy for UQ-Transformer.
• PUT eliminates misalignment artifacts at masked region borders—evident in ICT’s results—producing smooth, natural outputs even with irregular masks.
• On complex datasets like ImageNet, PUT shows a substantial improvement in image fidelity over existing transformer-based methods, highlighting its robustness to scene complexity.