CodeFormer——Towards Robust Blind Face Restoration with Codebook Lookup Transformer

链接：CodeFormer

摘要

盲脸恢复是一个高度病态的问题，通常需要辅助引导来1）改善从退化输入到期望输出的映射；2)   补充输入中丢失的高质量细节

在本文中，作者证明了在小代理空间中学习离散码本先验，通过将盲人脸恢复作为潜码预测任务，极大地减少了恢复映射的不确定性和模糊性，同时为生成高质量的人脸提供了丰富的视觉原子。在这种模式下，作者提出了一个基于transformer的预测网络，名为CodeFormer，用于对低质量人脸的全局组成和上下文进行建模，以进行潜码预测，即使在输入严重退化的情况下，也能发现与目标人脸非常接近的自然人脸。为了增强对不同退化的适应性，作者还提出了一个可控的特征转换模块，允许在保真度和质量之间灵活权衡。由于具有表现力的码本先验和全局建模，CodeFormer在质量和保真度方面都优于目前的技术水平，对退化表现出卓越的鲁棒性，在合成数据集和真实数据集上的大量实验结果验证了该方法的有效性。

引言

在野外拍摄的人脸图像经常遭受各种退化，如压缩、模糊和噪声。恢复这样的图像是高度病态的，因为退化引起的信息损失导致在低质量(LQ)输入的情况下产生无限的高质量(HQ)输出。在盲目恢复中，具体的退化是未知的，病态进一步升高。尽管随着深度学习的出现取得了进展，但在巨大的图像空间中，在没有额外指导的情况下学习LQ-HQ映射仍然是棘手的，导致早期方法的恢复质量不是最优。为了提高输出质量，降低LQ-HQ映射的不确定性和补充高质量细节的辅助信息是必不可少的。

各种先验被用来减轻这个问题的不适定性，包括几何先验、参考先验和生成先验。虽然改进的纹理和细节被观察到，这些方法往往遭受对退化的高敏感性或有限的先验表现力问题。这些先验对人脸恢复没有足够的指导，因此它们的网络本质上依赖于LQ输入图像的信息，而这些信息通常是高度损坏的。因此，LQ-HQ映射的不确定性仍然存在，并且由于输入图像的退化导致输出质量下降。最近，基于生成先验，一些方法通过迭代潜在优化或直接潜在编码将退化的人脸投影到连续的无限空间中。尽管输出的真实度很高，但在退化严重的情况下，很难找到准确的潜在向量，导致结果的保真度较低(图1(d))。为了提高保真度，在这种方法中通常需要在编码器和解码器之间进行跳过连接，然而，当输入严重退化时，这种设计同时会在结果中引入伪影。

与上述方法不同的是，该方法将盲人脸恢复作为一种编码预测任务，在学习到的离散码本先验的小的有限代理空间中进行，具有较好的抗退化鲁棒性和丰富的表达性。码本是通过使用矢量量化自编码器自构建HQ人脸来学习的，该自编码器与解码器一起存储丰富的HQ细节用于人脸恢复。与连续生成先验相比，码本项目的组合形成一个只有有限基数的离散先验空间。通过将LQ图像映射到更小的代理空间(例如1024个代码)，LQ- hq映射的不确定性显著减弱，增强了对各种退化的鲁棒性，如图1(d-g)所示。此外，码本空间具有更强的表现力，在感知上近似于图像空间，如图1(h)所示。这种性质允许网络减少对输入的依赖，甚至不需要跳过连接。

尽管基于码本的离散表示已被用于图像生成，但用于图像恢复的准确潜码组成仍然是一个非平凡的挑战。现有的工作是通过最近邻(NN)特征匹配来查找码本，由于LQ输入的固有纹理通常被破坏，这种方法在图像恢复中不太可行。LQ图像的信息丢失和多样性退化不可避免地会扭曲特征分布，影响准确的特征匹配。如图1(b)(右)所示，即使在LQ图像上对编码器进行了微调，LQ特征也不能很好地聚到确切的代码上，而是分散到附近的其他代码簇中，因此在这种情况下，最近邻匹配是不可靠的。

为此，本文提出了一个基于transformer的潜码预测网络，称为CodeFormer，以利用LQ面的全局组合和远程依赖关系来更好地进行代码预测。具体来说，Transformer模块以LQ特征作为输入，预测潜码token序列，该序列被视为码本空间中人脸图像的离散表示。由于采用全局建模来弥补LQ图像中的局部信息丢失，所提出的CodeFormer对严重退化具有鲁棒性，并保持整体一致性。对比图1(f-g)所示的结果，所提出的CodeFormer能够比最近邻匹配(如眼镜)恢复更多细节，从而提高了恢复的质量和保真度。

此外，本文提出了一个具有可调系数的可控特征转换模块，以控制LQ编码器到解码器的信息流。这样的设计允许在恢复质量和保真度之间进行灵活的权衡，以便在它们之间实现连续的图像转换。该模块增强了CodeFormer在不同退化情况下的适应性，例如在严重退化的情况下，可以手动减少携带退化的LQ特征的信息流，从而产生高质量的结果。配备了上述组件，所提出的CodeFormer在现有数据集以及新引入的WIDER-Test数据集(由从WIDER-Face数据集收集的970张严重退化的人脸组成)中表现出卓越的性能。除了面部恢复，该方法也证明了它在其他具有挑战性的任务上的有效性，比如面部修补，这些任务需要来自其他区域的远程线索。系统的研究和实验证明了本方法比以往的工作的优点。

受到VQVAE、VAGAN等工作中码本学习的启发，本文研究了一种用于盲人脸恢复的离散代理空间，通过全局建模预测潜码序列来利用离散码本先验，并通过固定编码器来确保先验的有效性。这样的设计允许该方法充分利用码本，使其不依赖于低质量线索的特征融合，显著提高了人脸恢复的鲁棒性。

CodeFormer的框架图。(a)首先重建任务来学习一个离散的码本还有解码器，用于保存人脸图像的高质量视觉部分。(b)随后固定码本和解码器，引入一个Transformer模块用于潜码序列预测，建模低质量输入的全局人脸组成。另外，引入一个可控的特征转换模块来控制LQ Encoder到解码器的信息流。注意，此处的连接是可选的，当输入严重退化时，可以禁用它以避免不利影响，并且可以调整标量权重w以在质量和保真度之间进行权衡。

3.方法

这项工作的主要重点是利用离散表示空间，减少恢复映射的不确定性，并补充退化输入的高质量细节。由于局部纹理和细节在低质量输入中丢失和损坏，因此使用Transformer模块对自然面的全局组成进行建模，从而弥补了局部信息的丢失，实现了高质量的恢复。总体框架如上图所示，首先结合矢量量化的思想，并通过自重构预训练一个量化的自编码器，以获得一个离散码本和相应的解码器；然后使用码本组合和解码器的先验信息进行人脸恢复。基于先前的代码本，然后使用Transformer来准确预测来自低质量输入的潜码组合。此外，引入了一个可控的特征转换模块，在恢复质量和保真度之间进行灵活的权衡，该方法的训练分为三个阶段。

3.1阶段一 潜码学习

为了减少LQ-HQ映射的不确定性并补充高质量的恢复细节，首先对量化自编码器进行预训练，以学习上下文丰富的码本，这提高了网络的表达性以及抗退化的鲁棒性。高质量图像首先通过编码器Eh嵌入为一个压缩的特征，然后和VQVAE、VQGAN一样，依据最近邻找到可学习码本中与潜码最相近的编码来替代当前潜码获取量化特征以及相应的潜码token序列s，解码器Dh根据Zc来重建高质量人脸。

为了用码本训练量化自编码器，作者采用了三种图像级重建损失:L1损失L1、感知损失L_per和对抗损失L_adv,由于更新码本项目时图像级损失是不受约束的，因此本文也采用中间码级损失L_feat码来减小码本C与输入特征嵌入Zh之间的距离。由于码本量化操作是不可微的，因此采用直通梯度估计器将梯度从解码器复制到编码器。整个训练的损失函数如下：

Lcodebook = L1 + Lper + Lfeat code + λadv · Ladv， 其中λadv被设定为0.8

码本设置

编码器Eh和解码器Dh分别由12个残差块和5个调整大小层组成，分别用于下采样和上采样。因此，获得了r = H/n = W/m = 32的大压缩比，这使得在第二阶段的全局建模具有很强的抗退化鲁棒性和可处理的计算成本。尽管更多的代码项可以简化重构，但多余的元素可能会导致后续代码预测中的不确定性。因此，将码本项N设置为1024，这足以进行准确的人脸重建，码维d设置为256。

3.2码本查找转换学习(第二阶段)

由于LQ人脸中纹理的破坏，最近邻(NN)匹配通常无法找到准确的人脸恢复代码。如图1(b)所示，具有不同退化的LQ特征可能会偏离正确的代码，并被分组到附近的簇中，从而导致不理想的恢复结果，如图1(f)所示。为了缓解这个问题，本文使用Transformer来为全局相互关系建模，以获得更好的代码预测。

如图2(b)所示，在第3.1节介绍的学习自编码器的基础上，作者在编码器后面插入一个Transformer模块，其中包含9个自注意块。接着固定了码本C和解码器DH，并微调编码器得到编码器EL。为了通过EL获得LQ特征Zl，我们首先将特征展开成m·n个向量，然后将其馈给Transformer模块，Transformer的第s个自关注块计算如下： Xs+1 = Softmax(QsKs)Vs + Xs，其中X0=Zv_l，查询Q、键K和值V是通过线性层从x中得到的，在查询Q和键K上添加正弦位置嵌入P，以增加建模顺序表示的表达性。

在自注意力块之后，采用线性层将特征进行投影。总体而言，Transformer以编码特征Zv l作为输入，以n个码项的概率形式预测m·n码序列{s∈{0，···，| n |−1}m·n。然后，预测的代码序列从学习的码本中检索m·n个各自的代码项，形成量化特征，通过固定的解码器DH生成高质量的人脸图像。受益于采用的大压缩比(即32)，使用Transformer在建模低质量人脸图像的全局相关性方面是有效和高效的。

训练目标

对Transformer模块T进行训练，对编码器EL进行微调恢复，而码本C和解码器DH保持固定。该阶段不使用图像级的重建损失和对抗损失，只需要码级损失:1)交叉熵损失Ltoken代码用于潜码token预测监督，2)L2损失Lfeat ‘代码用于强制LQ特征Zl接近码本中的量化特征Zc，从而降低了token预测学习的难度。

其中，从阶段1(第3.1节)预训练的自编码器中获得潜在代码s的真值，因此根据s从码本中检索量化特征Zc。请注意，经过这一阶段后，该网络在人脸恢复方面已经具备了优越的鲁棒性和有效性，特别是对于严重退化的人脸。

3.3可控特征转化(第三阶段)

尽管第二阶段已经获得了一个很好的面部恢复模型，但作者也研究了面部恢复质量和保真度之间的灵活权衡。因此，作者提出了可控特征变换(CFT)模块来控制从LQ编码器EL到解码器DH的信息流。具体而言，如图2所示，利用LQ特征Fe与α和β的仿射参数进行空间特征变换，对解码器特征Fd进行微调。然后使用可调节系数w∈[0,1]来控制输入的相对重要性。

其中Pθ表示从c(Fe, Fd)的连接特征中预测α和β的卷积堆栈。作者在编码器和解码器之间采用多尺度的CFT模块s∈{32,64,128,256}。这样的设计允许网络在轻度退化时保持高保真度，在严重退化时保持高质量。具体来说，可以使用较小的w来减少对严重退化的输入LQ图像的依赖，从而产生高质量的输出。较大的w将从LQ图像中引入更多信息，以在轻微退化的情况下增强保真度。

训练目标

为了在这一阶段训练可控模块和微调编码器EL，作者在第二阶段保留了Ltf的码级损失，并增加了L1、Lper和Ladv的图像级损失，这些损失与第一阶段相同，只是Irec被恢复输出Ires取代。这一阶段的完全损失是上述损失与其原始权重因子加权的总和。在这一阶段的训练中，我们将w设为1，这样在推理过程中，网络就可以通过调整[0,1]中的w来实现结果的连续过渡。对于推理，除非另有说明，默认设置w = 0.5，以便在输出的质量和保真度之间取得良好的平衡。

实验

-训练数据集。作者在FFHQ数据集上训练模型，该数据集包含70,000张高质量(HQ)图像，所有图像都被调整为512×512进行训练。为了形成训练对，作者使用一系列退化方法从HQ对应图像中合成LQ图像：

其中HQ图像首先与高斯核kσ进行卷积，然后进行尺度r的下采样。之后，在图像中加入加性高斯噪声nδ，然后应用质量因子q的JPEG压缩。最后，将LQ图像的大小调整回512×512，这一过程中分别从[1,15]，[1,30]，[0,20]和[30,90]中随机抽取σ， r， δ和q。

-测试数据集。作者在合成数据集CelebA-Test和三个真实数据集(LFW-Test、WebPhoto-Test和新提出的WIDER-Test)上评估了该方法。CelebA-Test包含从CelebA-HQ数据集[20]中选择的3,000张图像，其中LQ图像是在与本文的训练设置相同的退化范围下合成的。这三个真实数据集分别包含三种不同程度的退化，即轻度(LFW-Test)、中度(WebPhoto-Test)和重度(WIDER-Test)。LFW- test由LFW数据集[17]中每个人的第一张图像组成，包含1711张图像。WebPhoto-Test由407张从互联网上收集的低质量人脸组成。本文提出的WIDER- test包含来自WIDER face数据集的970张严重退化的人脸图像，为评估盲人脸恢复方法的泛化性和鲁棒性提供了更具挑战性的数据集。

设置。

本文将512 × 512 × 3的人脸图像表示为16 × 16的编码序列，对于所有阶段的训练，使用Adam优化器，批大小为16。将第一阶段和第二阶段的学习率设置为8×10−5，第三阶段的学习率设置为2×10−5。这三个阶段分别用1.5M、200K和20K次迭代进行训练。本文的方法是用PyTorch框架实现的，并使用四个NVIDIA Tesla V100 gpu进行训练。

指标。

对于具有ground truth的CelebA-Test的评价，本文采用PSNR、SSIM和LPIPS作为指标，还使用ArcFace网络中特征的余弦相似度来评估身份保持，表示为IDS。对于没有真实值的真实世界数据集的评估，使用了广泛使用的非参考感知度量:FID和MUSIQ (KonIQ)。

拓展

面部颜色增强

作者使用与GFP-GAN (v1)相同的颜色增强(随机颜色抖动和灰度转换)来微调得到面部颜色增强模型。将本文的方法与GFP-GAN (v1)[37]在具有颜色损失的真实世界老照片(来自CelebChild-Test数据集)上进行比较。所提出的CodeFormer生成高质量的面部图像，具有更自然的颜色和忠实的细节。

人脸补全

所提出的编解码器可以很容易地扩展到人脸补全任务，并且即使在大掩码比下也表现出良好的性能。为了构建训练对，作者使用一个公开可用的脚本来随机绘制不规则的多段线蒙版以生成遮罩面。作者将本文的方法与两种最先进的人脸绘制方法CTSDG和GPEN以及用于码本查找的最近邻匹配进行了比较。如图9所示，在口罩较大的情况下，CTSDG和GPEN会出现冲突。在本文的框架中使用最近邻匹配大致重建了面部结构，但它也无法恢复完整的视觉部分，如眼镜和眼睛。相比之下，CodeFormer能生成没有笔画和伪影的高质量自然面孔。

局限性

本文的方法是建立在一个预训练的自动编码器与一个码本。因此，自编码器的性能和表达能力会影响该方法的性能。1)虽然Transformer的全局建模大大缓解了身份不一致问题，但在配件等一些罕见的视觉部件中仍然存在不一致，目前的码本空间无法无缝地表示图像空间。在码本空间中使用多个尺度来探索更细粒度的视觉量化可能是一种解决方案。2)虽然在大多数情况下，CodeFormer具有很强的鲁棒性，但在侧面方面，CodeFormer相对于其他方法的优势有限，也不能产生很好的结果，如图10所示的失败案例。这是意料之中的，因为FFHQ训练数据集中只有很少的侧面，因此，码本无法为这种情况学习足够的代码，导致重建和恢复的效率都较低。

结论

本文旨在解决盲脸修复的基本问题。利用学习到的离散但表达性强的小码本空间，将人脸恢复转化为潜码token预测，显著降低了恢复映射的不确定性，简化了恢复网络的学习。为了弥补局部损失，本文通过一个富有表现力的Transformer模块来探索退化面的全局组合和依赖关系，以获得更好的代码预测。得益于这些设计，本文的方法具有良好的表达能力和较强的抗严重退化的鲁棒性。为了提高该方法对不同退化的适应性，本文还提出了一个可控的特征转换模块，允许在保真度和质量之间灵活权衡。实验结果表明了该方法的优越性和有效性。