自监督对比学习的泛化性理论

自监督学习范式已经在各种 CV 和 NLP 等任务上超越了直接用监督学习的性能。尽管近来自监督学习的算法层出不穷，但是对它泛化性好的理论理解仍然十分缺乏，只能盲目尝试，因此限制了自监督学习进一步的发展。本工作尝试对自监督对比学习的泛化性进行一些理论上的理解。

懒人版：下面视频的前 18 分钟

【2022.03 北大报告】华为诺亚黄维然：自监督学习理论

https://www.zhihu.com/zvideo/1483773806724997120

自监督学习 Self-Supervised Learning 是目前人工智能领域大热的一个方向，图灵奖得主 Yann Lecun 和 Yoshua Bengio 都认为自监督学习是机器智能达到人类水平的关键[1]。这种学习范式已经在各种 CV 和 NLP 等任务上超越了直接用监督学习的性能。尽管近来自监督学习的算法层出不穷，但是对它泛化性好的理论理解仍然十分缺乏，只能盲目尝试，因此限制了自监督学习进一步的发展。我们最近的工作尝试对自监督对比学习的泛化性进行一些理论上的理解。论文地址如下：

Towards the Generalization of Contrastive Self-Supervised Learning

https://www.weiranhuang.com/publications/contrastive2021.html

我们对自监督对比学习的泛化性进行分析，提出了一个统一的理论框架来解释对比自监督学习泛化性好的原因。我们对唯一的自监督信息——数据增强，进行数学上的定量刻画，并基于此推导出自监督对比学习在下游任务的泛化误差上界。根据该误差上界，得出泛化性好的三个条件：正样本对齐 Alignment of positive samples、类中心错开 Divergence of class centers、增广数据的聚焦性 Concentration of augmented data。其中，前两个条件是可以通过对比学习算法进行优化的——因此我们证明了两个常见算法 SimCLR 和 Barlow Twins 都在隐式地满足前两个条件；而第三个条件只依赖输入的训练数据——因此我们在真实数据集上测试，发现用我们理论预测出不同训练数据的泛化性变化趋势能跟真实性能变化趋势一致。

自监督学习是通过大量无标签数据来预训练一个模型，学到数据表示 representation，然后根据不同的下游任务，把学到的表示接在不同的分类器 / 回归器上，并用下游任务的有标签数据对分类器 / 回归器进行训练，从而适用于下游任务。自监督对比学习首先把 n 个训练数据进行两次数据增广得到 n 对“正样本”，然后把“正样本对”在特征空间中拉近从而学到特征提取器。由于只拉近正样本会导致学到常数平凡解（所有样本被映射到特征空间的同一个点），为了避免塌缩，不同的损失函数（比如 SimCLR/MoCo 等算法的 InfoNCE 引入负样本并将其远离正样本；Barlow Twins 等算法的互相关损失把特征的各个分量解耦）和训练方法（比如 BYOL/SimSiam 引入 stop-gradient）被提出。

尽管自监督对比学习已经被实验广泛验证了它在下游任务的泛化性，但是它的理论分析还十分有限。有一些工作从互信息最大化的角度理解 InfoNCE，指出 InfoNCE 的相反数是正样本互信息的下界，因此最小化 InfoNCE 其实是在最大化互信息的下界。然而，最新的结果[2]指出，如果去优化更紧的互信息下界，并不能得到更好的表示。这说明通过最大化互信息的角度也许并不一定能够完全解释 InfoNCE 的成功。

Arora 等人[3]提出一个理论框架来直接分析 InfoNCE 在下游任务的泛化性，但是他们基于一个假设：正样本是从同一个潜在类别 latent class 中选取的，比如两张不同的狗的图片作为一对正样本。而在实际中，一对正样本是同一个图片的两个不同的增广，不同的图片（即使属于同一个潜在类别）是作为负样本的，因此这个假设和实际情况出入较大。Wang 等人[4]提出 alignment 和 uniformity 作为预测下游性能的经验性指标，并没有严格的泛化性理论保证。

同时，上述两个工作都避开了去分析数据增广在自监督对比学习中的角色，然而它却是自监督对比学习有效的关键——因为仅有的人类先验知识是通过数据增广构造正样本来注入学习过程的（即人类可以识别出增广后的样本来自同一个图片，说明人类认知这两个增广后的图片所用的特征在这个两个图片中都具备，而自监督对比学习正是通过把上述人类经验教给机器从而学到好的特征）。

另外，Haochen 等人[5]从矩阵分解的角度来分析自监督对比学习，但是他们的理论分析仅仅适用于他们的 spectral contrastive loss，无法分析广泛引用的 SimCLR 或者 Barlow Twins 的损失函数。

除了上述现有自监督对比学习泛化性理论的局限性之外，还有一些有趣的实验现象无法完全解释。

• 更加丰富的数据增广能够让学到的特征空间具有更好的聚类性质。上图从左到右分别是 (a) 训练前的特征空间、(b) 只用一种类型的数据增广学到的特征空间、(c) 使用多种类型的数据增广学到的特征空间。可以看到 (a) 几乎看不出聚类，(b) 可以看出红色和蓝色有一定的聚类，但是还有很大的重叠区域，而 (c) 聚类性质则十分明显。而由于更好的聚类性质往往代表更好的下游分类精度（因为下游任务只训一个线性分类器），因此在下游任务上的表现 (c) 是优于 (b) 的。SimCLR 论文中也观测到类似的现象——更加丰富的数据增广具有更有的下游任务表现。直观上这个现象似乎是理所当然的，这能否从理论上给出定量刻画？此外，SimCLR 论文也做了数据增广两两组合的实验，发现 crop&color 效果最好。这能否有理论解释？
• 样本层面的拉近为什么会得到类别层面的聚类？比如 SimCLR 在优化的时候，仅仅让同一样本的两个数据增广靠近，而把同一潜在类别中其它样本视作该样本的负样本从而被推开，这等价于是让 2n 个增广后的样本做 n 分类，n 是训练数据量。但是在 SimCLR 学到的特征空间中，样本是按类别显著聚类在一起的（比如在 cifar-10 上的结果就是 10 类而不是 n 类）。这一点是如何保证的？
• 前面提到如果特征空间按类别聚类越好，下游任务的性能通常也越好。但是 Barlow Twins 并没有直接优化特征空间的几何性质让它聚类，而是优化特征的相关性性质（将特征的各个分量解耦），这为什么能保证特征空间在几何上有聚类的效果？

本文我们主要聚焦在理论分析自监督对比学习的泛化性 (Motivation 1)，有了我们的理论结果，就可以轻松地解释上述提到的实验现象 (Motivation 2)。首先，最重要的一点就是理解并用数学刻画出数据增广的作用。

直观来看，同一个类别的样本很可能具备相似的增广图片。比如上图中的两个狗的图片（绿框和蓝框），它们在像素层面是完全不同的（不同的姿势，不同的背景等等）。但是如果我们对它们做数据增广比如随机裁剪 random crop，那么它们有可能得到在像素层面十分接近的两个狗头（灰框）。两个相似的图片，它们的特征也应当接近 (Lipschitz 连续)，对应特征空间中的两个灰点。自监督对比学习的目标是把正样本在特征空间中拉近，因此以灰点为锚点，绿点和蓝点都往灰点靠拢，最终这两个狗的图片（绿框和蓝框）在特征空间聚到一起。

根据上述想法，我们定义一个新的两个图片之间距离叫做 augmented distance，它是两个图片在各自的增广中各选取一个增广，使得它们的像素距离最小。从上图的例子可以看出，augmented distance 从某种程度上刻画了两个图片的语义距离。紧接着, 我们可以定义 -增广, 即超过  部分的样本, 它们之间两两 augmented distance 都小于  。当  越大或者  越小的时候, 代表更加 sharp 的 concentration, 见下图。

有了上述-增广来定量刻画数据增广，我们通过理论推导出自监督对比学习泛化性好的三个条件：

1. Alignment of positive samples：这是不同自监督对比学习算法的共同目标，即正样本在特征空间中要被拉得越近越好；
2. Divergence of class center：不同类中心在特征空间里离得足够远，这保证了学到的表示不会塌缩到平凡解。这跟 Wang 等人[4]提出的 uniformity 并不相同，后面 3.1 节会具体讨论。
3. Concentration of augmented data：数据增广的聚焦性要好 sharp concentration，即越大的 \sigma\sigma 和越小的 \delta\delta 。

注意到，其中 1 和 2 两个条件是可以通过设计好的对比学习算法来满足的，但是条件 3 不依赖于算法。为了验证上述 3 个条件，我们分析了 SimCLR 和 Barlow Twins 这两个算法的 loss，并从理论上证明他们均可以满足 1 和 2 两个条件；为了验证条件 3，我们通过实验选取不同的数据增广，来验证越 sharp 的 concentration 确实能带来更好的下游任务效果。

我们用粗体的  来表示一个（随机）样本, 由它进行数据增广后的图片集合记作 , 那么 InfoNCE loss 可以写成：

而 Barlow Twins 中的 cross-correlation loss 可以写成：

其中  且  对每一维  是归一化的。

标准的自监督学习算法的评估方法是在学到的表示后面接一个线性分类器，并在下游数据上对线性分类器进行训练，由此得到的下游任务分类器，并以它的性能来作为评估学到的表示的好坏。为了简化数学分析，我们考虑不需要训练的最近邻分类器作为下游任务的分类器，即

其中  是第  类  的类中心。事实上，最近邻分类器可以看作是线性分类器的特例（见原文），因此实际中用线性分类器得到的下游任务性能会高于最近邻分类器。我们用误差率来测量最近邻分类器的性能，即

有了上述定义，下面我们研究为何自监督对比学习可以得到一个小的误差率。

我们首先给出上面 1.3 Main Idea 中 augmented distance 的数学定义：

以及 -增广的严谨定义：

当我们有一个更丰富的数据增广  时, 根据 augmented distance 定义, 我们有 , 因此当  不变时，  会变小。也就是说，更丰富的数据增广会带来更加 sharp 的 concentration。后面我们会看到更加 sharp 的 concentration 意味着更好的下游性能, 由此我们可以得出, 更丰富的数据增广带来更好的下游性能。

为了刻画正样本拉近的程度, 我们定义集合 里面的元素的两个增广都足够接近。进一步定义  越小表示 alignment 越好。

注：我们可以证明  可以用 alignment 来 upper bounded 如下, 具体细节见原文。

有了数据增广的定义，下面我们不加证明地直接给出本文的第一个定理：

为了更直观的理解上述定理, 我们首先来看一个简单的例子：任何两个来自同一个潜在类别的样本 都存在一个相同的增广图片 , 而且通过自监督对比学习正样本都 perfect aligned , 见下图：

在这种情况下, 来自同一个潜在类别的样本都被映射到了同一个点上, 因此我们只需要任意小的角度就能把不同的类别都准确区分开, 也就是  。事实上, 经过简单计算我们有 , 根据定理 1 得到, 这和我们直观分析的结论完全一致。

从定理 1 可以看出自监督性能好  即  小）的三个条件:

1. Alignment of positive samples：这也是不同自监督对比学习算法优化的共同目标。由定理 2 可 知 alignment 越小  越小, 再由定理 1 可知  也越小;
2. Divergence of class centers：不同的类中心的距离要足够大, 也就是足够小的  。从定理 1 可以看到, divergence 和 alignment  、concentration  都有关。越好的 alignment 和越 sharp 的 concentration（从而更小的  ) 会放松 divergence 的条件。这也可以从上图中的简单例子看出：perfect alignment 和 prefect concentration 下, 只需任意小的 divergence 即可。
3. Concentration of augmented data：当  固定的时候, 越 sharp 的 concentration 意味着越大的 , 从而直接降低了  的上界。

定理 1 提供了一个理论框架在分析不同的自监督学习算法的泛化性。上述条件 1 和 2 可以通过算法来优化（后面第 4 节我们将通过两个具体的例子 SimCLR 和 Barlow Twins 来揭示这一点），而条件 3 不依赖于算法，是由预先给定的增广数据来决定（后面第 5 节我们将通过实验来验证这一点）。

两个工作中都涉及 alignment，它们具有相同的含义，因为这是不同自监督对比学习算法优化的共同目标；但是我们提到的 divergence 和他们提出的 uniformity 在概念上完全不同。

• Uniformity 需要所有的数据点均匀的分布在特征球面上；我们的 divergence 刻画的是类中心之间的 cosine distance。我们不要求所有的数据点尽可能散开分布，相反，我们希望同类别的数据点尽量接近。同时，好的 alignment 和 concentration 可以放松对 divergence 的要求。因此，我们的 divergence 是一种更加准确和精细的刻画。
• Alignment 和 uniformity 是一种用来预测下游性能的经验指标，而我们的 alignment 和 divergence 是由定理 1 导出的，具有泛化性的理论保证。
• 另外，当 concentration 较差时（比如用恒等变换作为数据增广的方法），即使有 perfect alignment 和 perfect uniformity，下游性能仍然很差，这再次说明 alignment 和 uniformity 指标的局限性。

自监督对比损失函数一般可以分解成两部分，一部分是为了拉近正样本，另一部分是正则项为了保证不会学到塌缩解，即

其中, 第一项也就是 alignment, 即  。不同的损失函数的差别主要在于第二项, 一个有效的正则项应当让 divergence 满足定理 1 的条件。原文中的本节针对不同 loss 的数学分析较多, 在这里我们直接给出最后的结论:

首先, 我们证明了两种常见算法 SimCLR 和 Barlow Twins 的损失函数确实可以被分解成上面相关的两部分  和 ; 其次, 我们可以证明, 对于  有 , 对于 Barlow Twins 有 ; 最后, 我们给出下游性能的理论保证, 对于 SimCLR 有 , 对于 Barlow Twins 有

对于 SimCLR 的 InfoNCE loss，把它拆成 pos 和 reg 两项后得到

其中第二项正则项是 LogExp 形式，这对于避免塌缩很关键。如果用线性函数来替换第二项（有些文献成为 simple contrastive loss），

由于  和  的独立性, 第二项可以化简成 , 因此最小化第二项只能要求学到的表示均值为零。但是均值为零加上  无法避免表示发生维度塌缩 dimensional collapse（维度塌缩的概念最早由 @清华MARS Lab [6] 提出）, 比如所有的数据点被映射到特征球面上相对的两个点上, 如下图：

这也和 Wang 等人[7]发现的现象吻合：当 InfoNCE 的温度增加时，uniformity 变差（正则项变差）。因为当温度趋于正无穷时，InfoNCE 退化到上述 simple contrastive loss。

Barlow Twins 的 loss 被设计用来做特征分量解耦，那它为什么能满足 alignment 和 divergence 呢？

首先，把它拆成两项

那么

经过简单的计算，不难得到

也就是说，互相关矩阵的对角元其实决定了 alignment。

当 alignment 较好时, 可以预期 , 于是第二项中, 而对于每个类别中的不同样本基本都聚集在所属类中心的周围, 因此可以用  来近似代替 , 于是上述过程可以总结成

可以看到  距离我们想要的 divergence  只有一步之遥了。事实上, 利用矩阵 trace 的轮换性, 我们就能得到

这样就能保证 divergence 得到满足。

我们主要验证：更 sharp 的 concentration 预示着更好的下游性能。因此我们设计三组实验：

1）在第 3 节的定义 1 下面我们论证了更丰富的数据增广能推导出更 sharp 的 concentration，因此第一组实验我们观测越来越多的数据增广种类对下游性能的影响。我们考虑 5 种增广类型：(a) random cropping (b) random Gaussian blur (c) color dropping (d) color distortion (e) random horizontal flipping。

可以看到，随着数据增广种类越来越丰富，下游性能越来越好，且不同算法下趋势一致。有趣的点是加入 (c) 增广之后，下游性能有明显的提升。我们回过头去看 (c) 是什么操作，是把彩色图片变成灰度图片。由于这个操作让图片在像素层面有剧烈变动，因此由 augmented distance 的定义可以预期该距离会变小很多，从而导致 concentration 变得 sharp。一个直观的例子，在 Figure 2 中，如果把右边那个狗的图片替换成哈士奇（黑狗），那么仅仅做裁剪操作并不能让两个狗获得像素层面相似的增广图片，所以 augmented distance 较大，concentration 较差。而当我们引入 color dropping 操作之后，我们可以获得两个黑灰色的狗头，这样我们又可以得到像素层面相似的增广图片，从而能把它们的原始图片拉近。这时对应 augmented distance 变小，concentration 变得 sharp。

2）当我们固定数据增广的类型，而改变增广的强度，那么更强的数据增广会导致更小的 augmented distance，从而获得更 sharp 的 concentration 和 更好的下游性能。因此，第二组实验，我们固定 (a) 和 (d) 两种增广的种类，改变 (d) 的强度，观测对下游性能的影响。

可以看到，随着数据增广强度的增加，下游性能越来越好，且不同算法下趋势一致。

3）上面两组实验虽然一定程度上能够验证我们的结论，但是有人可能觉得实验结果是理所当然的。第三组实验，我们固定数据增广种类的个数为二，尝试所有可能的两两组合，并直接计算出不同组合的 ，从而观测它与下游任务的相关性。

可以看到, 下游性能  和 concentration (  当固定  时) 具有高度相关性：越 小的  代表越 sharp 的 concentration, 从而有更小的下游误差率。如果我们固定一个增广为 (a), 观测它与其他增广的组合, 可以看到不论是下游性能 (绿线) 还是 concentration (紫线、橙 线) 都具有  的排序。另外, 在全部的组合中, concentration 最 sharp 的 组合  成功预测出它也是下游性能最好的组合, 这也和 SimCLR 论文中的发现 crop&color 是 最有效的组合相吻合。

参考

1. ^https://cacm.acm.org/news/244720-yann-lecun-yoshua-bengio-self-supervised-learning-is-key-to-human-level-intelligence/fulltext
2. ^Michael Tschannen, Josip Djolonga, Paul K Rubenstein, Sylvain Gelly, and Mario Lucic. On mutual information maximization for representation learning. arXiv preprint arXiv:1907.13625, 2019.
3. ^Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, and Nikunj Saunshi. A theoretical analysis of contrastive unsupervised representation learning. arXiv preprint arXiv:1902.09229, 2019.
4. ^abcTongzhou Wang and Phillip Isola. Understanding contrastive representation learning through alignment and uniformity on the hypersphere. In International Conference on Machine Learning, pages 9929–9939. PMLR, 2020.
5. ^Jeff Z HaoChen, Colin Wei, Adrien Gaidon, and Tengyu Ma. Provable guarantees for self-supervised deep learning with spectral contrastive loss. Advances in Neural Information Processing Systems, 34, 2021.
6. ^Tianyu Hua, Wenxiao Wang, Zihui Xue, Sucheng Ren, Yue Wang, and Hang Zhao. On feature decorrelation in self-supervised learning. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9598–9608, 2021.
7. ^Feng Wang and Huaping Liu. Understanding the behaviour of contrastive loss. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 2495–2504, 2021.