原文链接 https://doi.org/10.48550/arxiv.2308.07314

前言：

在学习本论文前，你需要先学习VQGAN模型和CILP模型。

整体架构图：

第一阶段：离散码本学习阶段

与 VQGAN类似，本阶段通过自编码器与可学习码本分别对高质量域（HQ）与低质量域（LQ）进行特征编码，在训练中捕获各自领域的专属特征。HQ 路径与 LQ 路径采用完全相同的结构与训练设置，以保证特征表示的一致性。以下以 HQ 重建路径为例进行说明，LQ 路径训练流程与其完全一致。

如图所示，编码器 $E_H$ 将高质量人脸图像 $I_h\in {\mathbb{R}}^{H\times W\times 3}$ 编码为压缩特征 $Z_h\in {\mathbb{R}}^{m\times n\times d}$。随后进行矢量量化：将 $Z_h$ 中的每个特征向量，替换为可学习码本 C h = { c k ∈ R d } k = 0 N − 1 C_h = \{c_k \in \mathbb{R}^d\}_{k=0}^{N-1} Ch​={ck​∈Rd}k=0N−1​ 中距离最近的码元，得到量化特征 $Z_h^c\in {\mathbb{R}}^{m\times n\times d}$，计算方式为：
$$Z_h^{c(i,j)}=\arg \underset{c_k\in C_h}{\min }{\Vert Z_h^{(i,j)}-c_k\Vert }_2^2$$
最后，解码器 $D_H$ 利用量化特征 $Z_h^c$ 重建出高质量人脸图像 $I_h^{rec}$。

训练目标

与现有方法设置一致，本阶段采用四种损失进行训练，包括三种图像级重建损失与一种码本级损失：
$$\begin{array}{rl}{L}_1& =\Vert I_h-I_h^{rec}{\Vert }_1,\\ L_{per}& =\Vert \Phi (I_h)-\Phi (I_h^{rec}){\Vert }_2^2,\\ L_{adv}& =\log D(I_h)+\log (1-D(I_h^{rec})),\\ L_{code}& =\Vert \text{sg}(Z_h)-Z_h^c{\Vert }_2^2+\beta \Vert Z_h-\text{sg}(Z_h^c){\Vert }_2^2.\end{array}$$
其中，$\Phi$ 为 VGG19 特征提取器，$D$ 为基于补丁的判别器，$\text{sg}(\cdot )$ 为停止梯度算子，$\beta =0.25$。

最终的码本学习总损失为：
$$L_{\text{codebook}}=L_1+{\lambda }_{per}{L}_{per}+{\lambda }_{adv}{L}_{adv}+L_{code}$$
其中 ${\lambda }_{per}=1.0$，${\lambda }_{adv}=0.8$。

第二阶段：关联阶段

如图所示，在上一阶段得到高质量域（HQ）与低质量域（LQ）的编码器（$E_H$ 和 $E_L$）后，我们将输出特征（$Z_h$ 和 $Z_l\in {\mathbb{R}}^{m\times n\times d}$，如$Z_h$ 和 $Z_l\in {\mathbb{R}}^{16\times 16\times 256}$）展平为对应的特征块（$P_i^H$ 和 $P_i^L$，$i\in ([1,...,m\times n])$。这种展平操作使我们能够构建相似度矩阵（$M_{\text{assoc}}\in {\mathbb{R}}^{N\times N}$，其中 $N=m\times n$），用于量化不同特征块之间的相似度。

具体而言，我们计算每个特征块的余弦相似度，并约束矩阵对角线上的相似度最大化，(对角线为高质量图像和低质量图像空间位置对应的图像块)，通过施加该约束，两个编码器会在空间位置与特征层级上都将相近的特征块关联起来，在整个关联过程中保持空间对应关系。融合来自两个编码器的特征块后，我们得到两个关联后的编码器，记为 $E_H^A$ 和 $E_L^A$，它们融合了各自域的专属信息，将用于后续阶段。

训练目标

我们对高质量（HQ）与低质量（LQ）重建路径进行联合训练，并将特征关联模块纳入其中。为实现特征关联，我们采用交叉熵损失（$L_{CE}^H$ 和 $L_{CE}^L$）对相似度矩阵 $M_{assoc}$ 进行有效约束：
$$L_{CE}^H=-\frac{1}{N}\sum _{i=1}^N\sum _{j=1}^C{y}_{i,j}\log (p_{i,j}^h),L_{CE}^L=-\frac{1}{N}\sum _{i=1}^N\sum _{j=1}^C{y}_{i,j}\log (p_{i,j}^l),$$
其中，$N$ 表示特征块数量，$C$ 表示类别数，在本文中 $C=N$。$y_{i,j}$ 为真实标签，$p_{i,j}^h$ 为相似度矩阵 $M_{assoc}$ 在高质量域的余弦相似度得分，$p_{i,j}^l$ 为低质量域的得分。

补充真实标签 $y_{i,j}$ ：

我们把特征图切成了 N 个块（N = 16*16 = 256）：

• HQ 第 1 块、HQ 第 2 块 … HQ 第 N 块
• LQ 第 1 块、LQ 第 2 块 … LQ 第 N 块

我们希望：HQ第i块 只和 LQ第i块 对应，HQ第i块 不应该和 LQ其他块 对应

所以真实标签 $y_{i,j}$ 就是：

✔ 当 i = j（同一个空间位置），$y_{i,j}=1$ → 是“正确匹配”， 当 i ≠ j（不同空间位置），$y_{i,j}=0$ → 是“错误匹配”

---

特征关联阶段的最终损失为：
$$L_{assoc}=L_1+{\lambda }_{per}\cdot L_{per}+{\lambda }_{adv}\cdot L_{adv}+L_{code}+(L_{CE}^H+L_{CE}^L)/2,$$
我们沿用码本学习阶段的损失与权重，以保持特征表示的稳定性。（前四个损失和码本学习阶段一致，只有最后一个损失为该阶段特有。）

---

第三阶段 特征融合与码本预测阶段

在特征关联阶段得到两个关联后的编码器 $E_H^A$ 与 $E_L^A$ 后，如图所示，我们用两个编码器同时对低质量图像 $I_l$ 进行编码。具体来说，分别用 $E_H^A$ 和 $E_L^A$ 从 $I_l$ 中提取特征 $Z_h^A\in {\mathbb{R}}^{m\times n\times d}$ 与 $Z_l^A\in {\mathbb{R}}^{m\times n\times d}$。

我们使用多头交叉注意力MHCA模块融合两路特征，生成融合特征 $Z_f^A\in {\mathbb{R}}^{m\times n\times d}$：
$$Z_f^A=\text{MHCA}(Z_h^A,Z_l^A)$$MHCA 使模型能关注特征空间的不同维度，更好地捕捉两路编码器的有效信息。

得到融合特征 $Z_f^A$ 后，我们采用基于 Transformer 的分类方法预测对应的码本索引 $s$。首先将 $Z_f^A\in {\mathbb{R}}^{m\times n\times d}$ 展平为 ${\hat{Z}}_f^A\in {\mathbb{R}}^{(m\cdot n)\times d}$，输入 Transformer 得到预测索引 s ∈ { 0 , … , N − 1 } m ⋅ n s\in\{0,\dots,N-1\}^{m\cdot n} s∈{0,…,N−1}m⋅n。

此过程中，码本学习阶段得到的高质量码本 $C_h$ 与高质量解码器 $D_H$ 被冻结。利用预测索引 $s$ 在 $C_h$ 中检索对应特征 $Z_f^c$，输入解码器 $D_H$ 生成最终高质量修复图像 $I_{res}$。该步骤通过融合双域信息，有效提升人脸修复效果。

训练目标

我们使用两种损失函数来有效训练多头交叉注意力（MHCA）模块和Transformer模块，以确保特征融合与码本索引预测的学习效果。

第一种是 L2 损失 $L_{code}^{feat}$，它用于约束融合特征 $Z_f^A$ 尽可能接近来自高质量码本 $C_h$ 的量化特征 $Z_h^c$。该损失保证特征被合理融合，并同时保留来自高质量域（HQ）与低质量域（LQ）的有效信息。

第二种是用于码本索引预测的交叉熵损失 $L_{code}^{index}$，使模型能够精准预测高质量码本 $C_h$ 中对应的码本索引 $s$。

$$L_{code}^{feat}=\Vert Z_f^A-\text{sg}(Z_h^c){\Vert }_2^2,L_{code}^{index}=\sum _{i=0}^{mn-1}-{\hat{s}}_i\log (s_i)$$

其中，真实量化特征 $Z_h^c$ 和真实码本索引 $\hat{s}$ 均来自码本学习阶段；$Z_h^c$ 是通过真实索引 $\hat{s}$ 从高质量码本 $C_h$ 中检索得到的。

特征融合与码本预测的最终损失为：
$$L_{predict}={\lambda }_{feat}\cdot L_{code}^{feat}+L_{code}^{index}$$

在实验中，我们设置 L2 损失的权重 ${\lambda }_{feat}=10$。

总结

本文提出一种三阶段递进式低质量人脸修复方案：首先借鉴VQGAN为高低质量人脸域分别训练带可学习码本的自编码器，挖掘双域独立特征；再通过特征块余弦相似度矩阵与位置约束交叉熵损失，完成双域特征空间对齐与关联；最后利用多头交叉注意力融合双域特征，结合Transformer预测高质量码本索引，依托冻结的高质量码本和解码器重建清晰人脸，通过多阶段专属损失协同优化，实现跨域人脸高质量修复。