[ArXiv] A Simple Baseline for StyleGAN Inversion

[pdf][page][code]

Question

改进目前的反演方法

Contribution

1. a shallower backbone with multiple efficient heads across scales
2. multi-layer identity loss and multi-layer face parsing loss
   to the loss function
3. multi-stage refinement

以上为文中给出的Contribution，以下为我自己的总结

提出了一个简单但是出奇的有效的baseline，以pSp为基础，对网络各个方面的结构进行了探究，并进行一定程度的化简和优化，使得总体性能并没有下降，但是效率更高

介绍了网络结构的两点核心，对应Contribution 1

1. 从编码器的不同空间层次中提取的latent code，而且网络不是越深越好，较浅的网络就足以进行latent code的预测
2. 相比于pSp的18个map2style预测头，该方法精简到了3个，每个size上只有一个预测头，包含一个不同大小的全剧平均池化层和一个全连接层，更加轻量高效

以及两个新的loss，对应Contribution 2

1. multi-layer identity loss，提供了更强的语义对齐监督，从而大大提高了重建图像与输入的真实图像之间的身份一致性
2. multi-layer face parsing loss，这有助于捕捉局部的面部细节（例如，眼睛，鼠标），以进行更细粒度的重建。

多阶段细化学习，对应Contribution 3

与ReStyle类似，forward过程经过多次迭代，上一次迭代产生的结果会用于下一次迭代中，预测残差

实验表明，该方法优于现有的基于学习的方法，甚至达到了基于优化方法的性能，并且多个应用程序证明该方法在真实图像编辑任务中具有较好的泛化能力

Method

architecture

IDGI与pSp都是之前基于学习的Inversion方法，其中pSp是目前最先进的模型，它设计了两个金字塔状部分，对两部分特征图进行融合，在用18个(1024分辨率)ConvHead从中提取latent code。虽然表现很好，但或许网络中有很多冗余，降低了网络的效率，于是提出了以下三个问题

1. Is there any guideline for the efficient backbone design?
2. How to split the latent code into different feature levels?
3. Can we use cheaper regression heads?

question 1

对于第一个问题，是否有一个设计backbone的准则，使得backbone更加有效，避免网络出现无效的计算

为了找到一个合适的backbone深度，大致将相同分辨率特征视做相同的语义级别，使用了一个类似于IDGI的全连接头预测latent code，发现当从$\frac{1}{8}$分辨率级别上提取的特征足以匹配StyleGAN的最大语义级别，甚至可以得到比$\frac{1}{16}$更好的结果，这意味着更深的网络并不代表有更好的反演结果

该部分详细结果将在实验中列出

question 2

受pSp的启发，我们可以从不同语义级别预测latent code的不同部分，按上一节所述，我们从$\frac{1}{2}$、$\frac{1}{4}$、$\frac{1}{8}$三个级别进行预测，所以latent code也要划分成三个部分分别预测，那么问题是，我们该如何划分才能得到最好的预测结果

假设latent code三个部分的大小分别是$n_1,n_2,n_3$，则问题可以被描述为

$$n_{1}^{*}, n_{2}^{*}, n_{3}^{*}=\arg \max_{n_1,n_2,n_3} Q(n_1,n_2,n_3) + \lambda P(n_1,n_2,n_3)$$

$$s.t.n_1+n_2+n_3=18$$

以上$Q,P$分别是度量反演质量和模型大小的函数，由于$Q$依赖于训练过程，很难得到确切的函数，所以使用二分法求最优解。我们首先认为$n2,n3$是一个整体，只更改$n_1$的值，找到一个最优解，然后固定$n_1$，只修改$n_2$的值，再找到一个最优解。最终，$n_1,n_2,n_3$的值分别为9,5,4

question 3

如上所述，pSp中有18个ConvHead，每个都将特征压缩为$1\times 1$分辨率，但是我们发现这么多ConvHead并不是必要的，反而会带来多余的开销和信息损失。于是设计了一个简单但是高效的头，只有一个平均池化层和一个全连接层，且总共只有三个，分别对应上部分的三个级别，并负责生成对应级别的latent code

Multi-Stage Refinement

采用了和ReStyle相似的多层迭代机制，区别是输入时不用cat均值图像，其他基本一致，公式表示如下

$$W_t=\begin{cases} E_t(x)\qquad\qquad\qquad\qquad\, t=1\\ E_t(x,G(W_{t-1}))+W_{t-1}\,\,\, t>1 \end{cases}$$

loss

loss有三个部分，普遍使用的loss和两个新设计的loss

Common Losses

包括常用的L2 loss和LPIPS loss，如下

$$\mathcal L_2=||x-G(E(x))||_2$$

$$\mathcal L_{LPIPS}=||F(x)-F(G(E(x)))||_2$$

Multi-Layer Identity Loss

身份作为人脸图像的关键属性，保持身份信息不变是反演的重要任务之一，于是使用了一个预训练的人脸识别网络（ArcFace）的多层特征，分别计算余弦相似度

$$\mathcal L_{m\_id}=\sum_{i=1}^{5}(1-cos(R_i(x),R_i(G(E(x)))))$$

其中$cos$是余弦相似度，$R_i$为输入为$x$时，人脸识别网络$R$的第$i$个语义级别对应的特征

Multi-Layer Face Parsing Loss

由于人脸识别网络侧重于人脸的全局身份特征，因此仅靠身份损失无法进行对细节的精确重建，所以为了更好的重建局部信息，引入了一个预训练的面部解析网络$P$获得多层特征，以提供更局部的信息，该loss如下

$$\mathcal L_{m\_par}=\sum_{i=1}^{5}(1-cos(P_i(x),P_i(G(E(x)))))$$

overall loss

全部loss定义如下

$$\mathcal L=\lambda_1\mathcal L_2+\lambda_2\mathcal L_{LPIPS}+\lambda_3\mathcal L_{m\_id}+\lambda_4\mathcal L_{m\_par}$$

其中$\lambda_1,\lambda_2,\lambda_3,\lambda_4$分别为$1,0.8,0.5,1$

Results

backbone为SE-ResNet50，只保留尺寸为$\frac{1}{16}$之前的部分，StyleGAN为在FFHQ上训练的StyleGAN2

Quantitative and Qualitative Comparison

Table 1列举了各个方法的各项指标，*是没有进行残差学习的网路，我总结了以下几点

1. Ours的1-3Stage之间对比，迭代确实增强了反演的效果，各项指标都有增加
2. Ours-2Stage∗与其他Stage对比，分数与1Stage一致，说明当没有残差学习时，迭代没能加强反演的效果，所以迭代可能是一种对残差网络结构的利用，我会在Think部分谈下我的理解
3. Ours-1Stage与pSp和IDGI-Encoder相比，效果更好了，说明该方法即使不经过迭代，各项指标也要比目前最先进的基于学习的方法好，而且更轻量
4. Ours与基于优化的I2S和混合方法IDGI对比，在经过迭代后SSIM能达到优化方法的水平，PSNR比混合方法好但还没达到优化的水平，IDS即使Ours-1Stage都比之前的所有方法好了，推测可能是使用了更好的loss（后来发现不止因为loss，后面会有补充）

以下是对比图

Ablation Study

对Size、Head、Loss三个上面讨论的部分进行了消融实验，证明了设计的有效性，和上面的猜测有出入，IDS的提升也有网络结构的功劳。对于pSp中的金字塔结构也进行了消融实验

(d)+FPN是Table 2中(d)情况加上与pSp相同的金字塔结构，发现并不能导致性能的提高

对于多层身份和解析loss的作用，有以下对比图

w/o为没有loss的情况，我们发现Multi-ID确实比Single-ID的效果要好，而Multi-Parsing loss的引入使得网络对眼镜和舌头等细节的关注度提高，以更好的进行重建

Applications

应用部分只选取图像编辑部分的对比图放在这里

并不好评价，感觉有些面部细节的缺失，而且没有和专为编辑设计的e4e进行比较

Think

关于目前的基于学习的反演方法，一个主要的趋势就是添加迭代，但是每次迭代都会使总参数量增加一个网络的大小，而在后面的几次迭代中，可能并不需要整个网络的参数，所以后面几次迭代在计算量或者参数量上可能都是有冗余的。那么我们需要解决的问题就是，如何将后面几次迭代的网络化简，将迭代的网络化为一整个更轻量有效的网络，对于这方面的想法我会在下面详细列举

Why Work

我们可以换个角度，将迭代视作好几个一样的网络连在一起，只进行一次输入，如下图

其中，三个网络是一模一样的，但是显而易见的是后面两个并不需要这么多的参数，所以在一定程度上，经过残差学习，后面的网络被化简了，在计算中并没有用到全部的部分，这也是为什么失去残差学习后，迭代失去了效果，残差结构的缺失使得单独一个网络难以拟合多个网络的效果（猜测增加参数量或者将三个网络分开训练会解决这个问题）

如何改进

对上面的结论，我有一点猜测，如果Ours-1Stage的分数比Ours-3Stage只迭代一次的分数高的话，就说明该网络还有改进空间，我们单纯的将三个网络分开训练就可以获得更好的分数

此外，我们还可以可以化简后面两个网络，并对其进行针对性得改进，比如让其更专注于细节等等，最终有可能得到的结果是，增加了一些网络的参数量，但是获得了更快的速度和更好的效果

突然有了想法写在这里，如第一个网络用来预测初始latent code，后面两个网络预测残差，模拟一个先生成再优化的过程，效果会不会更好