聚类对比学习：SwAV & PCL模型浅析

2022 年 9 月 6 日 PaperWeekly

©作者 | 段易通

单位 | 清华大学博士生

研究方向 | 机器学习

作为一种自监督学习方法，对比学习（Contrastive Learning）由于在表征学习上有效性和在不同领域中的普适性，已经成为了近年来的研究热点之一， 这篇文章将介绍对比学习的一个分支——聚类对比学习 ，涉及到两篇论文：SwAV 和 PCL。关于对比学习的基本概念这里就不过多介绍了，还不熟悉的同学可以先看一些相关资料。

我们知道，对比学习是通过 比较正负样本 对来进行学习的，对于一般的对比学习方法，正样本往往是从样本自身得到的，比如对该样本做增强（augmentation），而负样本则是从 batch 中随机挑选的样本。然而这种构造方法可能会面临两个问题：

1. 为了加强模型的分辨能力，我们往往需要在一个 batch 中加入足够多的负样本，许多实验也表明大的 batch size 可以提高模型性能。然而由于一般的对比学习方法需要对 batch 中的样本进行两两比较，计算复杂度为，这就导致 batch size 会受到显存大小的约束，给对比学习的应用带来了障碍。

2. 随机挑选负样本的方式可能会将一些实际上很相似的样本作为负样本（例如若锚定样本是一只狗的图片，而随机挑出来的负样本恰好也是一只狗的图片，那么即使两个样本实际上很相似，模型也会将其作为负样本），这样可能会影响模型的性能。

聚类对比学习就是想要解决上述的问题，顾名思义，该方法 不直接做两两样本的对比，而是先对样本进行聚类，然后在类之间进行对比学习 ，由于样本的“类别”是通过无监督的聚类方法得到的，因此整个学习过程中并不需要样本标签，仍然还是在做自监督学习。通过聚类后再对比的操作，我们就可以大大减小对比的数量，降低计算复杂度，而且同一类下的不同样本也互为正样本，不会将相似的样本当做负样本。

下面将具体介绍 SwAV 和 PCL，两者都将聚类的思想引入到了对比学习当中。

SwAV

要想在对比学习中引入聚类，首先需要解决的就是聚类稳定性的问题，即要避免所有样本都被分到同一类这种极端情况。因此这里 SwAV 就设计了一种聚类的方法，使得模型可以直接对 batch 中的样本进行稳定的聚类，然后按照聚类结果进行对比学习，帮助模型获得样本表征。

先简单说一下 SwAV 的思路，SwAV 与传统对比学习的区别如下所示。为了帮助聚类，模型中引入了 K 个原型（prototype）变量，每个 prototype 可以理解为是一类样本的通用特征，一共有 K 类；模型会根据原型变量，online 计算出 batch 中的样本属于某类的概率，然后再通过不同 view 的对比，令正样本对的属于同一类的概率更大（即相似的样本属于同一类的概率更大），以此来实现对于样本表征的学习。

具体来说，模型首先提取样本表征 z，然后通过 softmax 函数计算该样本属于 K 个类的概率：

SwAV 要让两个 view 下的样本尽量属于同一类，但是需要注意的是，我们不能直接最大化两个 view 的分类概率（比如直接最小化交叉熵），这样可能会导致平凡解（即模型将所有样本都分到一个类别里面）。

为了避免这样的情况，SwAV 这里计算了一个最优分配概率 q，即图中的 code，它是按照最优运输（optimal transport）的思想，在满足类间平衡约束的情况下，将 batch 中样本按照 q 分配到 k 个类中，使得与原型相似（內积大）的样本尽量被分到一起，即：

其中正则项

用来控制分配概率的离散程度，越大则分配的概率越均匀。这里我们假设在实际采样的时候batch中不同类的样本数量是相等的，因此约束条件要求分配概率在不同类上是均匀的。

SwAV 采用了 Sinkhorn-Knopp 算法来迭代求解 q，然后计算交叉熵

及其交换项来作为两个 view（）之间的对比学习 loss，其最终形式为：

PCL

论文标题：

Prototypical Contrastive Learning of Unsupervised Representations

论文链接：

https://arxiv.org/abs/2005.04966

代码链接：

https://github.com/salesforce/PCL

PCL 采用了期望-最大化框架，即 EM 算法的思路，流程如下图所示，简单来说可以分为 E 步和 M 步：

E 步：momentum encoder 从所有样本中提取表征，并对表征做 k-means 聚类，得到每个样本所属类别、以及聚类中心。

M 步：从数据集中采样若干个 batch，然后分别通过 encoder 和 momentum encoder 获得表征，计算聚类对比 loss，梯度更新 encoder。

具体来说，PCL 的 EM 算法的推导过程如下：

这里采用 EM 算法来最大化（6），其中 E 步用来估计分类概率，这里采用 k-means 算法，通过聚类来估计得到 one-hot 概率；对于 M 步，我们假设 k 个聚类的先验分布是均匀的，则有：

接着，假设给定聚类中心的条件分布是高斯分布，则有：

其中，这里认为 v 和 c 是归一化后的向量（），将其代入式（6）中，则原先的对数似然最大化可以写做：

其中代表一个聚类中点的集中程度，这里通过动量 encoder 得到的特征来自适应地计算，具体如下所示：

其中表示该聚类下的样本数量，是平滑参数，模型可以通过该参数控制聚类的稳定（具体论述在原文 3.3 节，但我没看懂...）；最后 PCL 的 loss 可以写做：

该归一化 loss 称作 ProtoNCE，主要分为两部分，第一项是原始的 InfoNCE loss（不了解的同学可以看我之前介绍 InfoNCE 的文章 [3]），后一项是在不同聚类数量下（即种聚类结果）的对比 loss。

总结

这篇文章介绍了两种聚类对比学习的方法，其核心问题在于如何保证对比学习过程中聚类的稳定性，SwAV 采用了给定约束条件下的最优运输理论来实现均匀的聚类；而 PCL 则通过引入自适应的聚类集中度来调节聚类。完成对样本聚类之后，SwAV 得到的是“软”的分类概率，然后通过计算两个 view 之间的分类交叉熵进行对比学习；而 PCL 得到的是 one-hot 聚类结果，然后把每个聚类看做是一个高斯分布，通过最大化对数似然的方式对比进行学习。