[论文解读] 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 Synthesis被引用 8

一句话总结

本文提出 PUT，一种基于 Transformer 的新型框架，用于复数图像修复，通过避免输入下采样和量化来最小化信息损失。它使用基于图像块的 VQ-VAE（P-VQVAE）对非重叠特征进行编码，并采用无量化 Transformer（UQ-Transformer）处理未量化的特征，同时以量化标记作为预测目标，从而在大而复杂的掩码区域上实现了最先进的保真度。

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

研究动机与目标

• 解决由于输入下采样和像素量化导致的基于 Transformer 的复数图像修复中的信息损失问题。
• 保留高分辨率图像细节和空间对齐性，尤其是在掩码区域边界处。
• 在大型和复杂区域，特别是 ImageNet 等大规模数据集上，提升修复的保真度和多样性。
• 设计一种适用于无量化 Transformer 的训练策略，即使在训练期间通过随机量化实现离散标记预测。
• 消除先前方法中常见的高分辨率与低分辨率表示之间的错位伪影。

提出的方法

• 提出 P-VQVAE，一种基于图像块、非重叠的自编码器，可在不进行下采样的情况下将图像块编码为潜在特征。
• 在 P-VQVAE 中使用双码书，分别表示掩码和未掩码的图像块，以改善特征解耦。
• 在 P-VQVAE 中使用解码器重建掩码区域，同时保持未掩码区域不变。
• 提出 UQ-Transformer，其输入为未量化的特征向量，但以量化标记作为目标进行预测，从而避免因量化导致的信息损失。
• 在训练过程中对 UQ-Transformer 使用随机量化的输入特征，以模拟推理条件，从而实现在推理阶段无需完全量化即可实现精确生成。
• 在整个网络中保持高分辨率输入，消除了输入与预测之间分辨率的错位。

实验结果

研究问题

• RQ1 避免输入下采样和量化是否能提升基于 Transformer 模型的图像修复保真度？
• RQ2 在 Transformer 中处理未量化特征如何影响生成图像区域的质量和多样性？
• RQ3 自编码器中采用非重叠图像块设计在多大程度上减少了掩码区域与未掩码区域之间的干扰？
• RQ4 在 UQ-Transformer 训练期间使用随机量化是否能有效实现在推理阶段使用离散标记，同时保持特征保真度？
• RQ5 PUT 在处理不规则形状和大范围掩码区域方面与现有方法相比表现如何？

主要发现

• PUT 在 FFHQ、Places2 和 ImageNet 数据集上的 FID、LPIPS 和 FVD 指标上均达到最先进性能，尤其在大范围掩码区域上显著优于先前方法。
• 消融实验表明，采用双码书的 PUT（PUT）优于单码书（PUT one）和无参考量化推理（PUT no_ref），证实了结构化码书设计的优势。
• PUT 显著优于使用 CNN 编码器和滑动窗口处理的 PUT conv，证明了非重叠图像块编码在减少掩码-未掩码区域干扰方面的优越性。
• 在训练中使用随机量化的模型（PUT）显著优于未使用该策略的变体（PUT qua0），验证了 UQ-Transformer 训练策略的有效性。
• PUT 消除了掩码区域边界处的错位伪影——这在 ICT 的结果中明显可见——即使在不规则掩码下也能生成平滑、自然的输出。
• 在 ImageNet 等复杂数据集上，PUT 相较于现有基于 Transformer 的方法在图像保真度方面表现出显著提升，凸显了其对场景复杂性的鲁棒性。