一篇解决表示学习坍塌问题的工作报告

1   表示学习中的坍塌问题

在表示学习中，一种很常见的做法是利用孪生网络的结构，让同一个样本的不同数据增强后的表示相似。比如对于一张图像，使用翻转、裁剪等方法生成另一个增强图像，两个图像分别经过两个共享参数的编码器，得到表示，模型的优化目标是让这两个表示的距离近。

这种方法一个比较大的挑战在于，模型在训练过程中容易出现坍塌问题。模型可以把所有样本的表示都学成完全相同的常数向量（比如每个样本模型的输出都是全0向量），这样就能满足上述两个表示距离近的要求了。

2   解决坍塌问题的方法

业内有很多解决坍塌问题的方法，主要可以分为4种类型：

对比学习方法：在训练一对正样本对时，同时采样大量的负样本，让正样本之间离得近，负样本之间离得远，避免模型偷懒把所有样本的表示都学成一样的。

聚类方法：在训练过程中增加一个聚类过程，将样本分配给不同的类簇，然后在类簇级别进行对比学习。

基于蒸馏的方法：通过模型结构的角度避免坍塌问题，学习一个student network来预测teacher network的表示，teacher network是student network参数的滑动平均，teacher network不通过反向传播更新参数。

信息最大化方法：让生成的embedding中每一维的向量相互正交，使其信息量最大化，这样可以避免各个维度的值信息过于冗余，防止坍塌问题。

目前常用的负样本采样方法，一个比较大的问题是计算开销大，取得好的效果往往需要大量负样本，因此有了MoCo等对比学习框架。本文提出的方法基于信息最大化的思路，能够只使用正样本对实现表示学习的同时，防止坍现象的发生。关于对比学习的常用经典方法，可以参考这篇文章：对比学习中的4种经典训练模式。

3   信息最大化方法历史工作

基于信息最大化的方法典型的工作有两篇ICLM 2021的文章，分别是Whitening for Self-Supervised Representation Learning（ICML 2021，W-MSE）和Barlow twins: Self-supervised learning via redundancy reduction（ICML 2021）。下面介绍一下这两篇文章的整体思路，Facebook的这篇论文也是基于这个思路设计的。

W-MSE的整体网络结构如上图，模型中输入一对互为正样本的样本对（例如一个图像的不同增强形式），使用共享参数Encoder分别编码后，增加一个whitening模块，对每个batch内的所有embedding进行白化，让embedding的各个维度变量线性无关，后面再接norm处理。下图形式化表明了W-MSE的用途，通过白化+norm让样本形成一个球形分布，正样本之间距离近，每个样本需要调整自己在圆周上的位置拉进正样本之间的距离，形成最终的簇。

Barlow twins方法也是同样的思路，只是做法不同。在得到一对正样本的两个embedding后，计算两个向量各个变量之间的相关性矩阵，优化这个矩阵接近对角矩阵，以此实现embedding中各个变量之间线性无关。

4   VICREG

Facebook团队提出的VICREG是上述基于信息最大化表示学习方法的一个扩展，损失函数主要包括variance、invariance、covariance三个部分。模型的主体结构如下，仍然是一个孪生网络结构，输入一个样本的不同view互为正样本。

variance部分主要是让batch内每个样本embedding向量的每一维变量之间的方差不小于一个值，这样可以有效防止每个样本的向量都对应同一个点的情况，防止坍塌的发生。

invariance部分即让正样本对的表示embedding距离尽可能小，是正常的表示学习loss。

covariance借鉴了Barlow twins中的思路，让batch内的embedding非对角线元素尽可能为0，即让表示向量各个维度变量线性无关。

最终的表示学习loss是由上述3个loss的加和：

5   实验结果

作者进行了一些实验验证VICREG表示学习方法的效果。在ImageNet数据集上，使用不同方法预训练的表示接一个线性层进行分类，或者基于预训练的表示使用部分数据进行finetune，对比图像分类效果，实验结果如下：

作者也对比了在不同方法上是否采用variance或covariance等loss的效果变化：

6   总结

本文介绍了Facebook AI团队在2022年ICLR的一篇解决表示学习坍塌问题的工作，顺带介绍了基于信息最大化的防坍塌工作，有助于帮助我们更深层次理解表示学习的原理，以及如何解决实际应用表示学习、对比学习时遇到的坍塌问题。

审核编辑：刘清