图像修补论文阅读：MAT算法笔记

标题：MAT: Mask-Aware Transformer for Large Hole Image Inpainting
会议：CVPR2022
论文地址：https://ieeexplore.ieee.org/document/9879508/
官方代码：https://github.com/fenglinglwb/MAT
作者单位：香港中文大学、Adobe公司、香港中文大学（深圳）、上海AI实验室

Abstract

最近的研究表明了在修复问题中建模长距离相互作用的重要性。为了实现这一目标，现有的方法要么使用独立的注意力技术，要么使用transformers，但考虑到计算成本通常是在较低的分辨率下使用。在本文中，我们提出了一种新的基于transformer的大空洞修复（large hole inpainting）模型，它结合了transformers和卷积的优点来高效地处理高分辨率图像。我们精心设计了框架的每个组成部分，以保证恢复图像的高保真度和多样性。具体来说，我们定制了一个面向修复的transformer block，其中的注意力模块仅从部分有效的token中聚合非局部信息，用动态掩码表示。大量实验证明了该新模型在多个基准数据集上的SOTA性能。

1. Introduction

图像补全（又称图像修补）是计算机视觉中的一个基本问题，其目的是用似乎是真的内容填补缺失的区域。它的应用包括图像编辑、图像重定位、照片修复和目标去除等。
在修补中，对上下文信息进行建模是至关重要的，特别是对于较大的掩码。为缺失区域创建合适的结构和纹理需要上下文理解，根据图像中的非局部先验使用远处信息。之前的工作使用堆叠的卷积来达到大的感受野和建模长距离关系，这在对齐的（例如，人脸、身体）和纹理密集的（例如，森林、水）数据上效果很好。在处理结构复杂的图像时（即图1中第2行的第1个例子），全卷积神经网络（CNNs）难以表征长距离区域之间的语义对应。这主要是由于CNNs的固有性质、有效感受野的增长缓慢以及附近像素不可避免的主导性。为了显式建模修补中的长距离依赖关系，之前的方法提出在基于CNN的生成器中使用注意力模块。然而，受限于平方计算复杂度，注意力模块仅仅应用于相对较小尺度的特征图，且次数较少，其中的长距离建模没有得到充分利用。

与将注意力模块应用于CNNs不同，transformer是一种天然的处理非局部建模的结构，其中的注意力是每个block中的基本组件。近期的进展采用transformer结构来解决修补问题。然而，受复杂度问题的影响，现有工作仅利用transformers推理低分辨率的预测（例如32×32）用于后续处理，因此生成的图像结构较为粗糙，影响了最终的图像质量，尤其是在大尺度的掩码上。
在本文中，我们开发了一种新的修补transformer，能够为大掩码的修补生成高分辨率的完整结果。由于一些区域缺少有用信息（这在给定的掩码排除了大多数像素时是常见的），我们发现常用的transformer block（LN→MSA→LN→FFN）在对抗训练中表现出较差的性能。为此，我们改制了原始的Transformer block，通过移除传统的层标准化、用特征拼接的融合学习代替残差学习，来增加优化稳定性，也提高了性能。我们分析了为什么这些修改对于学习是至关重要的，并从实证上证明了它们是非平凡的。此外，为了处理从高分辨率输入中提取到的所有token之间可能存在的大量相互作用，我们提出了一种新的多头自注意力（multi-head self-attention，MSA）变体，命名为多头上下文注意力（multi-head contextual attention，MCA）。它只使用部分有效的token计算非局部关系。采用的令牌的选择由一个动态掩码表示，其由输入掩码初始化，并结合空间约束和长距离相互作用进行更新，在不牺牲有效性的情况下提高效率。此外，我们在提出的框架中合并了一个新颖的风格操纵模块，内在地支持多元化生成。如图1所示，我们的方法视觉上逼真且内容极其多样化地成功填补了大空洞。我们的贡献总结如下：

• 我们开发了一个新颖的修补框架MAT。它是第一个能够直接处理高分辨率图像的基于transformer的修补系统。
• 我们精心设计了MAT的组件。所提出的多头上下文注意力通过利用有效的token来高效地进行长距离依赖建模，用动态掩码表示。我们还提出了一种改进的transformer block，使训练大掩码更加稳定。此外，我们设计了一个新颖的风格操纵模块来提高多样性。
• MAT在Places和CelebA-HQ等多个基准数据集上达到了SOTA。它也使多元化补全成为可能。

2. Related Work

图像修补是计算机视觉领域长期存在的问题。早期基于diffusion的方法将相邻的未损坏的信息传播到空洞中。在内部或外部搜索空间中，基于patch或exemplar的方法根据人类定义的距离度量，借助外观相似的patch来填补缺失区域。PatchMatch提出了一种多尺度patch搜索策略来加速修补过程。此外，偏微分方程和全局或局部的图像统计量在之前的文献中有大量研究。尽管传统方法往往可以获得视觉上逼真的结果，但缺乏高级理解阻碍了它们生成语义合理的内容。
在过去的几年中，深度学习在图像修补上取得了巨大的成功。Pathak等人将对抗训练引入到修补中，并利用基于编码器-解码器的结构来填补空洞。后来，U-Net结构的许多变体被开发出来用于图像修补。此外，还提出了更复杂的网络或学习策略来生成高质量的图像，包括全局和局部判别、上下文注意力、部分和门控卷积等。多阶段生成也得到了大量关注，其中目标边缘、前景轮廓、结构和语义分割图等中间线索被广泛利用。为了允许高分辨率的图像修补，已经有很多方法做了一些尝试来开发渐进生成系统。
最近，研究人员将研究重点转向更具挑战性的场景，其中最具代表性的问题是多元化生成和大空洞填补。对于前者，Zheng等人提出了一个具有两条并行路径的概率原则化框架，能够产生多个似然解。UCTGAN通过优化KL散度将实例图像空间和掩码图像空间投射到一个公共的低维度的流形空间，从而允许缺失内容的多元化生成。随后，Wan等人和Yu等人利用双向注意力或自回归transformers来完成这一目标。尽管这些方法提高了多样性，但由于变分训练和基于光栅扫描顺序的生成，它们的修补和推理性能有限。另一方面，一些工作被提出用于解决大空洞的修补问题。例如，CoModGAN利用调制技术处理大尺度的缺失区域。在本工作中，我们开发了一个新颖的框架来同时实现高质量的多元化生成和大空洞填补，将长距离上下文相互作用和无限制生成的最佳效果带到了图像修补任务中。

3. Method

给定一幅掩码过的图像，记为${\mathbf{I}}_{\mathrm{M}}=\mathbf{I}\odot \mathbf{M}$，图像修补的目的是对缺失区域填补视觉上具有吸引力且语义合适的内容。在这项工作中，我们提出了一个用于大掩码修补的掩码感知transformer（mask-aware transformer，MAT），支持条件的长距离相互作用。此外，鉴于图像修补问题的不适定性（不确定性），即可能存在许多可行的解决方案来填补大空洞，我们的方法旨在支持多元化生成。

3.1. Overall Architecture

如图2所示，我们提出的MAT结构由一个卷积头部、一个transformer主体、一个卷积尾部和一个风格操纵模块组成，充分发挥了transformers和卷积的优点。具体来说，卷积头部被用来提取token，然后，具有5个stage、每个stage中的transformer block有不同分辨率（有不同的token数量）的主体通过提出的多头上下文注意力（MCA）建模长距离相互作用。对于来自主体的输出token，采用基于卷积的重建模块将空间分辨率上采样到输入尺寸。此外，我们采用另一个Conv-U-Net来细化高频细节，依靠CNNs的局部纹理细化能力和效率。最后，我们设计了一个风格操纵模块，使模型能够通过调制卷积的权重来提供各式各样的预测。

3.2. Convolutional Head

卷积头部接收不完整的图像${\mathbf{I}}_{\mathrm{M}}$和给定的掩码$\mathbf{M}$作为输入，并生成$\frac{1}{8}$大小的特征图用于token。它包含4个卷积层，一个用于改变输入维度，其余的用于对分辨率降采样。
我们使用卷积头部主要出于两个原因。首先，在早期视觉处理中融入局部归纳先验对于更好的表征和可优化性仍然至关重要。另一方面，它被设计用于快速降采样以降低计算复杂度和内存开销。此外，我们通过实验发现该设计要优于ViT中使用的线性投射头，正如补充材料中所验证的。

3.3. Transformer Body

transformer主体通过建立长距离关联来处理token。它包含具有所提出的调整过的transformer block的5个stage，具有由额外的掩码引导的高效注意力机制。

3.3.1 Adjusted Transformer Block

我们提出了一种新的transformer block变体来处理大空洞掩码的优化。具体来说，我们去除了层标准化（LN），使用融合学习（使用特征拼接）代替残差学习。如图3所示，我们将注意力的输入和输出拼接起来并使用一个全连接（FC）层：
$$\begin{array}{rl}{\mathbf{X}}_{k,\ell }^{\prime }& =\mathrm{FC}([\mathrm{MCA}({\mathbf{X}}_{k,\ell -1}),{\mathbf{X}}_{k,\ell -1}])\\ {\mathbf{X}}_{k,\ell }& =\mathrm{MLP}({\mathbf{X}}_{k,\ell }^{\prime })\end{array}$$其中，${\mathbf{X}}_{k,\ell }$是MLP模块第$k$个stage中第$l$个block的输出。经过几个transformer block后，如图3所示，我们采用一个卷积层和一个整体的残差连接。注意，由于之前的文献已经表明3×3卷积足以为transformers提供位置信息，所以我们舍弃了在transformer block中的位置嵌入。因此，这里的传播只取决于特征相似性，促进了长距离相互作用。
分析。 transformer的常规结构包含两个子模块，多头自注意力（MSA）模块和MLP模块。在每个模块之前应用一次层标准化，在每个模块之后应用一次残差连接。然而，在处理大尺度的掩码时，我们观察到使用常规的block优化不稳定，有时会发生梯度爆炸。我们将这个训练问题归因于大比例的无效token（它们的值几乎为零）。在这种情况下，层标准化可能会极大地放大无用的token，导致训练不稳定。此外，残差学习通常鼓励模型学习高频内容。然而，考虑到大多数token开始时是无效的，在GAN的训练中没有适当的低频基础很难直接学习高频细节，这使得优化变得更加困难。将这种残差学习替换为拼接会得到明显更好的结果。

3.3.2 Multi-Head Contextual Attention

为了处理大量的token（对于512×512的图像最高可达4096个token）和给定token中的低保真度（至多90%的token是无用的），我们的注意力模块利用了移动窗口和动态掩码，能够使用少量可行的token进行非局部相互作用。计算有效token的加权和作为输出，如图3所示，其公式为：
$$\mathrm{Att}(\mathbf{Q},\mathbf{K},\mathbf{V})=\mathrm{Softmax}(\frac{\mathbf{Q}{\mathbf{K}}^T+{\mathbf{M}}^{\prime }}{\sqrt{d_k}})\mathbf{V}\text{\hspace{1em}(3)}$$其中，$\mathbf{Q}$、$\mathbf{K}$、$\mathbf{V}$分别是query、key、value矩阵，$\frac{1}{\sqrt{d_k}}$是缩放系数。掩码${\mathbf{M}}^{\prime }$表示为：
$${\mathbf{M}}_{i,j}^{\prime }=\{\begin{array}{c}0,\mathrm{if}\mathrm{token}j\mathrm{is}\mathrm{valid}\\ -\tau ,\mathrm{if}\mathrm{token}j\mathrm{is}\mathrm{invalid}\end{array}\text{\hspace{1em}(4)}$$其中，$\tau$是一个大的正整数（在我们的实验中为100）。在这种情况下，无效token的聚合权重几乎为0。在每次注意力之后，我们通过把$w\times w$尺寸的窗口移动$(\lfloor \frac{w}{2}\rfloor ,\lfloor \frac{w}{2}\rfloor )$像素，实现跨窗口联系。
掩码更新策略。 掩码（${\mathbf{M}}^{\prime }$）指出了token是否有效，其由输入掩码初始化并在传播过程中自动更新。更新遵循一个规则，即只要之前至少有一个有效token，那么窗口中的所有token在经过注意力后都被更新为有效。如果窗口中的所有token都是无效的，那么它们在经过注意力之后仍然是无效的。如图4所示，经过从a到b的注意力，左上窗口中的所有token都变为有效，而其它窗口中的token仍然为无效。经过几次窗口移动和注意力后，掩码将更新变得全部有效。

分析。 对于被缺失区域主导的图像，默认的注意力策略不仅不能借助可见的信息来修补空洞，而且会破坏有效的像素。为了减少色差或模糊，我们提出只使用有效的token（由一个动态掩码选择）来计算关系。

3.4. Style Manipulation Module

受之前文献的启发，我们设计了一个风格操纵模块来赋予我们的框架多元化生成。它通过附加一个噪声输入，改变重建过程中卷积层的权重标准化来操纵输出。为了增强噪声输入的表征能力，我们强制条件图像（image-conditional）风格${\mathbf{s}}_c$从图像特征$\mathbf{X}$和无条件噪声（noise-unconditional）风格${\mathbf{s}}_u$中学习，公式为：
$$\begin{array}{rl}& {\mathbf{s}}_u=\mathcal{E}(\mathbf{n})\\ & {\mathbf{X}}^{\prime }=\mathbf{B}\odot \mathbf{X}+(1-\mathbf{B})\odot \mathrm{Resize}({\mathbf{s}}_u)\\ & {\mathbf{s}}_c=\mathcal{F}({\mathbf{X}}^{\prime })\end{array}\text{\hspace{1em}(5)(6)(7)}$$其中，$\mathbf{B}$是一个随机的二值掩码，上面的值以概率$p$设置为1，以概率$1-p$设置为0，$\mathcal{E}$和$\mathcal{F}$是映射函数。如图2所示，最终的风格表征通过融合这两种风格表征得到：
$$\mathbf{s}=\mathcal{A}({\mathbf{s}}_u,{\mathbf{s}}_c)\text{\hspace{1em}(8)}$$其中，$\mathcal{A}$是一个映射函数。然后，将卷积的权重$\mathbf{W}$表示为：
$$\begin{array}{rl}& {\mathbf{W}}_{ijk}^{\prime }={\mathbf{W}}_{ijk}\cdot {\mathbf{s}}_i\\ & {\mathbf{W}}_{ijk}^{\prime \prime }={\mathbf{W}}_{ijk}^{\prime }/\sqrt{\sum _{i,k}{{\mathbf{W}}_{ijk}^{\prime }}^2+ϵ}\end{array}\text{\hspace{1em}(9)(10)}$$其中，$i$、$j$、$k$分别表示卷积的输入通道、输出通道和空间位置，$ϵ$是一个小的常数。不同风格表征的调制导致了多元化的输出。同时，我们将噪声注入包含到我们的框架中，以进一步增强生成的多样性。

3.5. Loss Functions

为了提高生成的质量和多样性，我们在两个阶段都采用非饱和对抗损失（non-saturating adversarial loss）来优化我们的框架，而不考虑像素级的MAE或MSE损失，它们通常会导致均值模糊（averaged blurry）结果。我们还使用了$R_1$正则化，记为$R_1=E_x||\Delta D(x)||$。此外，我们采用了感知损失（perceptual loss），经验上系数较低，因为我们注意到它可以更容易地进行优化。
对抗损失。 我们将对抗损失计算为：
$$\begin{array}{rl}& {\mathcal{L}}_{\mathrm{G}}=-{\mathbb{E}}_{\hat{x}}[\log (D(\hat{x}))]\\ & {\mathcal{L}}_{\mathrm{D}}=-{\mathbb{E}}_x[\log (D(x))]-{\mathbb{E}}_{\hat{x}}[\log (1-D(\hat{x}))]\end{array}\text{\hspace{1em}(11)(12)}$$其中，$x$和$\hat{x}$分别表示真实图像和生成图像。我们对图2中的两阶段生成都应用对抗损失。
感知损失。 感知损失计算为：
$${\mathcal{L}}_{\mathrm{P}}=\sum _i{\eta }_i||{\varphi }_i(\hat{x})-{\varphi }_i(x)|{|}_1\text{\hspace{1em}(13)}$$其中，${\varphi }_i(\cdot )$表示预训练的VGG-19网络的层激活。我们只考虑$con{v}_{4\_4}$和$con{v}_{5\_4}$的高层特征，允许缩放系数${\eta }_i$为$\frac{1}{4}$和$\frac{1}{2}$的修补结果的变化。
总体损失。 生成器的总体损失为：
$$\mathcal{L}={\mathcal{L}}_{\mathrm{G}}+\gamma R_1+\lambda {\mathcal{L}}_{\mathrm{P}}\text{\hspace{1em}(14)}$$其中，$\gamma =10$和$\lambda =0.1$。

4. Experiments

4.1. Datasets and Metrics

我们在512×512分辨率的Places365-Standard和CelebA-HQ数据集上进行实验。具体来说，在Places数据集上，我们分别使用来自训练集和验证集的180万张和3.65万张图像来训练和评估我们的模型。对于CelebA-HQ，训练和验证拆分由24183张和2993张图像组成。尽管是在512×512的图像上训练，我们在补充材料中表明我们的模型可以很好地推广到更大的分辨率。
就大空洞场景而言，根据Zhao等人，我们选择感知指标，包括FID、P-IDS和U-IDS进行评估。我们认为使用像素级的L1距离、PSNR和SSIM是不合适的，因为初步的工作已经表明这些指标与人类对图像质量的感知相关性很弱，特别是对于不适定的大尺度图像修补问题。尽管LPIPS是在深度特征空间中计算的，但像素级的评估仍然极大地惩罚了针对大空洞的各式各样的修补系统。因此，我们仅在补充材料中使用它作为参考。

4.2. Implementation Details

在我们的框架中，我们将头部、主体和重建模块的卷积通道数和FC维度数设置为180。5层transformer组的block数和窗口尺寸分别是 { 2 , 3 , 4 , 3 , 2 } \{2,3,4,3,2\} {2,3,4,3,2}和 { 8 , 16 , 16 , 16 , 8 } \{8,16,16,16,8\} {8,16,16,16,8}。最后一个Conv-U-Net首先将分辨率降采样到$\frac{1}{32}$，然后再上采样到原始尺寸，其中，在不同尺度的卷积层数和通道数借鉴Karras等人。映射网络由8个FC层组成，风格操纵模块由卷积后接1个AvgPool层实现。与之前的方法不同，我们的transformer结构不使用预训练。
所有实验均在8个NVidia V100 GPU上进行。根据Zhao等人，我们在Places365-Standard上训练了50M张图像，在CelebA-HQ上训练了25M张图像。batch size为32。我们采用${\beta }_1=0$和${\beta }_2=0.99$的Adam优化器，将learning rate设置为$1\times 10^{-3}$。自由形状掩码（free-form mask）在补充文件中进行了描述。

4.3. Ablation Study

在本节中，我们梳理了我们框架中哪些组件对最终性能的贡献最大。为了实现快速探索，我们仅在256×256分辨率的Places中使用100K训练图像（≈5.6%），用5M个样本（完全设置的10%）训练模型。我们采用前10K个验证图像进行评估。定量比较见表1。

Conv-Transformer结构。 我们探究由transformers建模的长距离上下文关系是否对修补大空洞有用。将transformer block替换为卷积block（表1中模型“B”），我们发现在所有指标上都有明显的性能下降，尤其是在P-IDS和U-IDS上，表明修补图像失去了一定的保真度。此外，我们在图5中给出了一些可视化的例子。与全卷积网络相比，我们的MAT利用长距离上下文很好地重建了球网的结构和恐龙骨架的纹理，显示出长距离相互作用的有效性。

调整过的Transformer Block。 在我们的框架中，由于传统的设计容易导致不稳定的优化，因此我们开发了一个新的transformer block，在这种情况下我们需要降低transformer主体的学习率。如表1所示，我们的设计（模型“A”）获得了优于使用原始transformer block的模型“C”的性能，FID提高了0.39。如图5所示，我们注意到我们的设计产生了更有视觉吸引力的结果，支持高质量的图像补全。特别是对于第一个例子，即使缺失区域非常大，我们的方法仍然可以恢复出语义一致且视觉逼真的室内场景。
多头上下文注意力。 为了使用逼真的内容快速地填补缺失区域，我们提出了多头上下文注意力（MCA）。为了更深入地理解，我们从有效的token中构建了一个没有部分聚合（partial aggregation）的模型。定量结果如表1中模型“D”所示。值得注意的是，FID下降了0.1，而其它指标变化不大。我们认为，本文提出的上下文注意力有助于保持颜色一致性和减少模糊。如图5所示，没有MCA的模型在第一个例子中生成了颜色不正确的内容，在第二个例子中生成了模糊的伪影。定量和定性结果都验证了我们MCA的能力。
风格操纵模块。 为了处理大的掩码，除了有条件的长距离相互作用，我们还引入了无条件的生成。为了量化我们框架的无条件生成能力，我们除去了噪声风格操纵。从表1模型“E”的结果来看，我们发现P-IDS和U-IDS有很大的差距，表明随机噪声风格的调制进一步提高了补全图像的自然度。
重建中的高分辨率。 由于二次方增加的计算复杂度，现有的工作采用transformers合成低分辨率的结果，通常为32×32，用于后续处理。相比之下，我们的MAT结构利用其计算效率，在重建阶段能够有高分辨率的输出。如表1所示，我们的完整模型“A”比模型“F”取得了显著的改进，证明了高分辨率预测的重要性。

4.4. Comparison with State of the Arts

我们将提出的MAT与一些SOTA方法进行了比较。为了公平比较，我们使用公开可用的模型在相同的掩码上进行测试。如表2所示，MAT在CelebA-HQ和Places上都取得了优于或可比的结果。值得注意的是，我们只使用了一个子集Places365-Standard（1.8M图像）来训练我们的模型，远远少于CoModGAN（8M图像）和BigLaMa（4.5M图像）。此外，我们的方法比次优的CoModGAN和基于transformer的ICT更加轻量化（parameter-efficient）。如图8所示，与其它方法相比，所提出的MAT以更少的伪影恢复出更加照片真实的（photo-realistic）图像。例如，我们的方法成功地恢复出视觉上令人满意的鲜花和规则的建筑结构。

4.5. Pluralistic Generation

我们框架的内在多样性主要源于风格操纵。如图6所示，风格变体导致了不同的补全。从图6中的第一个例子中，我们观察到了一个从噘嘴微笑到露牙大笑的变化。第二个例子显示了不同的人脸轮廓和外观。至于最后一个，我们发现了不同的窗户和屋顶结构。

4.6. Limitations and Failure Cases

在没有语义标注的情况下，MAT经常在处理具有各种形状的目标时挣扎，比如跑动中的动物。如图7所示，由于缺乏语义上下文理解，我们的方法未能恢复出猫和车。此外，受限于注意力中的降采样和预定义的窗口尺寸，我们需要对图像进行pad或resize，使其尺寸为512的倍数。

5. Conclusion

我们提出了一种用于多元化大空洞图像修补的掩码感知transformer（MAT）。利用所提出的调整过的transformer结构和部分注意力机制，所提出的MAT在多个基准测试集上取得了SOTA的性能。此外，我们设计了一个风格调制模块来提高生成的多样性。大量的定性比较证明了我们的框架在图像质量和多样性方面的优越性。