图像风格转换 之 CUT 论文解读与感想

Contrastive Learning for Unpaired Image-to-Image Translation（CUT）是发布于2020年ECCV的一篇关于图像风格转换的文章。本文最大的特点就是将目前大热的对比学习应用到了图像风格转换问题上来，实现了一种轻量级的图像风格转换模型，并且本文可以进一步实现单张图片之间的风格转换。

文章切入点

无配对图像风格转换领域近些年来最成功和最经典的模型无疑就是CycleGAN了，CycleGAN通过两组风格转换网络$G_{A2B}$和$G_{B2A}$实现了两种风格之间的相互转换，使得网络形成了一种闭环式的训练模式，通过引入独特的cycle consistency loss和identity loss对模型进行监督。然而我们的目的是将风格A转换到风格B，这种训练模式在面对单方向图像风格转换时无疑引入了过度的限制，浪费大量计算资源在辅助网络上。
通过常识我们可以知道，对于图像A和其对应的风格转换图像$A^B$，对应空间位置上的图像应该具有相似的信息（例如相似的外部轮廓）。作者通过这个常识来入手，引入了对比学习的思想，提出了一种全新的损失函数——InfoNCE loss。互信息是对比学习中一个最基本和最重要的概念，它指的是两个特征向量所隐含的相同的信息。通过将$A$和$A^B$对应位置互信息最大化，不仅可以监督生成网络发现两个对应位置的共性部分（风格转换时应该被保留的部分）还能监督生成网络去忽略掉不同位置的特性部分（风格转换时应该被去掉的部分）。

具体实现和模型设计

首先提一下，本文的结构是一个基本的GAN-based图像风格转换网络，可以表达为
$$A^B=G(A)$$，
所以自然的要引入对应的对抗损失：

我们知道对比学习最重要的部分就是如何构造正负样本来求互信息，所以第一个非常关键点是寻找正负样本。本文构造正负样本的方法非常巧妙，作者通过将图像中的patch看做一个对象，将$A^B$的一个随机位置$a$处的patch看做是锚点，然后$A$的$a$位置patch看做是一个正样本，将$A$中剩余的部分的patch看做是负样本，这其实就回应了前面的直观理解。让我们给出损失函数的数学表达：

其中$v$是$A^B$的$a$处patch图像的特征向量，$v^{+}$和$v^{-}$分别对应$A$中$a$处和不在$a$处的图像的patch的特征向量，$\tau$是一个超参数。通过最小化上式，便可以最大化$v$和#v^+#的互信息，而最小化$v$和$v^{-}$的互信息。
接下来是第二个非常关键的点，我们在确定了正负样本后，如何获得对应的特征向量。这里作者巧妙的将前面的生成网络G拆成了两部分：$G=G_{dec}(G_{enc})$
这样随着我们就可以通过$G_{enc}$来获得高维特征向量而不用去引入其他辅助网络，这里这个技巧也是CUT可以实现轻量化的一个关键所在。这里还要说一点，我们知道随着$G_{enc}$越深，其产生的特征信息（$w\times h\times c$）也越来越细长，对应的每一个特征向量（$w\times h\times 1$）所表征的patch面积也越大，因此这里作者并不是用的某一层$G_{enc}$的输出，而是用了多层的输出，每一层都求前面对比损失。此外，本文也引入了SimCLR中映射头的概念，使得patch的特征向量具有更好的可分性。给每一层$G_{enc}$添加一个映射头，应该算是本文添加的唯一的辅助网络了，然而2层全连接映射头的计算量相对于生成和鉴别网络来说是可以忽略不计的。
最终通过对$G_{enc}$的每一层输出，以及输出的每个patch的特征向量求互信息，上述损失函数就可以写为：

其中$L$是总的层数，$S_l$是对应层总的patch。

讨论和个人理解

先说一下本文做的一些有趣的消融实验：
一、
我们说前面我们的PatchNCE是将$A$输入到生成器$G$中得到$A^B$，然后将不同风格的二者求对比损失，在cycleGAN中有一个著名的损失叫做identity loss，它是将$B$输入到生成器中，力图得到的图像没有变化，即$B^B=G(B)\approx B$。identity loss的作用是为了使得生成器$G$忽略掉输入图像的风格，而只专注于输入图像的内容，这样可以避免生成图像产生额外不必要的变化。在这篇文章中，作者同样模仿着提出了一个identity loss：
$$L_{PatchNCE}(G,H,Y)$$
以此将模型分为了完整版（CUT），其优化目标是：

轻量化版（FastCUT），其优化目标是：

并以此做了很多对比试验。从试验结果来看，添加了identity loss虽然会提高一些计算量（memory：CUT=3.33，FastCUT=2.25，CycleGAN=4.81），但是产生的图像质量更好（不同的数据集下质量提升大约17%到38%），并且稳定性大大提高，个人认为这部分额外计算量是可以接受的。
二、
我们开始的时候说，构造的负样本是从目标图像其他位置选的patch（本文称之为内部patch），其实我们完全可以从其他图像选一些patch来当做负样本（本文称之为外部patch），本文就做了关于这个问题的对比试验：

可以看出来，在仅有外部patch和内部patch+外部patch的情况下，对比只有内部patch，性能都有不同程度的下降。这里文中给出了一点直观的解释：

通过使用内部统计数据，生成器不需要对大的类内变量(如白马vs.棕马)建模，这对生成输出斑马是不必要的

这里我也给出个人的理解：对比损失实际上是促使生成器发现正样本和目标patch之间的相同点，发现负样本和目标patch之间的不同点。当正样本和负样本都聚焦于同一张图像时，由于正样本和目标patch的相同点是图像内容，那么$A$的图像内容便会被$G_{enc}$所注意，而由于负样本们的风格和正样本的风格是相同的，那么如果$G_{enc}$将注意力放在了风格上面，无疑会使得负样本和目标patch的互信息增大，所以为了负样本互信息最小化，$G_{enc}$会去忽略掉$A$的风格信息。以此，便达到了我们理想的$G_{enc}$，即最大化保留输入图像的内容信息，而最小化输入图像的风格信息（这实际上和identity loss的作用殊途同归）。然而当使用外部patch时，$G_{enc}$完全可以为了欺骗对比损失，将部分注意力集中于$A$和$B$的风格上的相似性上，这样便会冲淡对于$A$结构上的关注度，从后面的Qualitative ablation results 实验也可以看出，当使用外部patch时，生成的图像会对原输入图像的内容部分作出一些意外的变动。这正是这篇文章给我的关于对比学习的玄学之处，如无必要，切勿引入额外信息分散模型注意！