由颜水成老师带领的 Sea AI Lab 提出了一种多粒度自监督学习框架 Mugs[1],用以学习不同粒度的非监督特征,从而满足不同下游任务对不同粒度甚至多粒度特征的需求。在相同的实验设置下(相同数据集和模型等),该方法大幅超越了目前最好的自监督学习方法。在没有使用额外数据的情况下,该方法在 ImageNet 数据集上取得了目前最高的线性评估准确率(linear probing accuracy)82.1% 以及最高的 KNN 分类准确率 80.3%。
详细结果请参看 paperswithcode 的Leaderboard:https://paperswithcode.com/sota/self-supervised-image-classification-on?p=mugs-a-multi-granular-self-supervised
一:背景和动机
通过精心设计上游任务(pretext task),自监督学习从这些上游任务上直接拿到无监督数据的伪标签,并且利用这些伪标签训练深度学习模型。实验结果表明通过自监督学习到的特征在很多下游任务上展现出比常规监督学习更强的迁移性能和泛化性能。目前大量的无标签数据可以在互联网上获得,自监督学习已经展示出比常规监督学习更大的潜能。
在自监督学习中,许多下游任务其实需要不同粒度的特征,例如粗粒的特征或者细粒度的特征。举例来说,常规的分类问题一般会用更多的粗粒度特征,而细粒度的分类问题则需要更多的细粒度特征。更重要的是,在实际应用中,许多下游任务需要多粒度的特征来处理数据中的复杂结构。一个典型的例子是 ImageNet 上的分类问题。首先,注意到 ImageNet 数据非常贴近实际数据,拥有类别上的层级结构,也就是很多小类拥有相似的属性从属于一个大类(例如犬科类里面包含很多不同品种的狗和狼)。这样一来,为了区分不同的大类别,例如犬科类和鸟类,我们就需要粗粒度的特征。为了进一步区分更细的类别,譬如狗和狼,我们则需要稍微细粒度一点的特征。最后为了细分不同品种的狗,我们需要更加细粒度的特征。许多其他应用,例如物体检测与识别,也面临着类似的多粒度特征需求。
然而,现在的自监督学习方法一般都只侧重于单粒度特征的学习。譬如,最经典的对比学习 MoCo[2] 构建实例识别任务来区分单个实例。这样一来,MoCo 侧重学习更多实例级的细粒度特征,但是它不考虑数据中粗粒度的类别结构(cluster structure)。另外一类具有代表性的自监督学习是基于聚类的自监督学习,包含 DINO[3],DeepCluster[4]等。这类方法将类似的实例聚到同一个虚拟类别中,从而学习聚类级别的粗粒度特性。然而,它不能很好地处理对细粒度特性有所需求的下游任务。因此,在下游任务的特征偏好未知的情况下,我们应该构建一个学习多粒度特征的自监督学习框架,以便尽可能多地处理不同下游任务。
为了解决上述问题,该研究提出了一种多粒度自监督学习 (Mugs) 框架来显示地学习视觉数据的多粒度特征。具体来说,Mugs 主要学习三种不同粒度特征:
- 第一是 instance-level(实例级别的)细粒度特征,用于将不同的实例区分开;
- 第二是稍微粗粒度的 local-group 数据特征。这类特征主要是为了保证非常相似的样本能够聚集成小簇,从而形成一个 local-group 上的整体语义(例如对应到非常相似的狗)。这样一来,local-group 特征则会将第一类过于分散的实例特征进行高阶语义的重新聚集,也就让相似的样本享有相似的特征;
- 第三类是数据中更加粗粒度的 group 特征(例如犬科特征)。这类特征是保证相似的 local-groups 能够靠的更近,从而保证数据整体语义的一致性。通过这样的层级数据语义学习,该方法能够学到更贴合实际应用的特征,从而更加满足于不同下游任务对不同粒度特征的需求。
二:方法
如前面所述,该方法主要通过三种不同粒度的监督来学习三种不同粒度的特征, 包含 1) instance-level 细粒度特征,2)稍微粗粒度的 local-group 特征, 3)粗粒度的 group 特征。下面将介绍这三种特征的具体学习方法。
图一:Mugs 自监督学习整体框架。(a)展示了总体框架。对于每幅图像,Mugs 将两次随机增强的样本输入学生和教师骨干网络。然后, Mugs 采用三种不同粒度的监督 / 损失: 1)instance discrimination loss, 2)local-group discrimination loss,3) group discrimination loss。(b)展示了学生和老师中都有的 local-group discrimination 模块。这个模块首先平均所有 patch tokens 得到 average token,并从一个 buffer 里面找到当前 average token 的 top-k 邻居。然后,它使用一个小的 transformer 来聚合 average token 和它的 k 个邻居,从而获得一个 local group feature(即新的 class token)。
A). instance-level 细粒度特征学习
和对比学习 MoCo 一样,Mugs 通过 instance discrimination supervision(实例判别监督)来分散不同的实例特征,从而监督实例级别的细粒度特征学习。
具体来说,给定样本, Mugs 进行两次数据增强得到, 然后分别送进老师和学生 backbone 得到。接着 Mugs 把里面的 class token 分别送入对应的投影网络。最后,Mugs 将学生网络对应的输出进一步输入到一个预测网络用来预测老师的输出。为了达到让老师和学生输出一致的目的,Mugs 使用 infoNCE 损失函数来实现:
其中,分别对应老师和学生分支最后的输出。代表的是一个缓存 buffer,用于存贮过往的。这样一来,正如图一中的效果图所示,Mugs 把同一个图片对应的特征拉进,把不同图片的特征推远,从而分散开 instance-level 的特征。
B). 稍微粗粒度的 local-group 特征学习
如前所述,细粒度特征通常不足以满足不同的下游任务(例如分类),因为缺乏足够的高级数据语义。为了学习更高级别的细粒度(即稍微粗粒度)特征,这里也称为 local-group feature,Mugs 提出了一种新颖而有效的 local-group supervision。直观地讲,如图一中第三个球面所示,这种监督鼓励数据特征具有小但分散的 local-group 结构。目的是保证高度相似的样本能够在特征空间上聚集,避免前面所介绍的 instance discrimination supervision 将相似的样本推得过远,从而避免破坏一定的语义。这样一来,local-group supervision 有助于捕获数据中的更高级别语义。接下来将首先介绍其细节,然后更加详细地讨论它的好处。
给定两个数据增广后的样本, 假设老师和学生骨干网络得到。如图一中的子图(b)所示,Mugs 把里面的 patch tokens 分别进行平均池化得到。在另一方面,Mugs 使用一个 buffer 以先进先出的方式来保存过往的。接着,Mugs 分别为当前的里面找到它们对应的 top-k 邻居。然后,Mugs 借助一个小的 transformer 将融合成一个 local-group feature:
其中,代表的是一个两层的 transformer:它的 input class token 是,input patch tokens 是,输出 class token。有着同样的功能。因为和它的 top-k 邻居很相似,它们一起组成一个 local-group。因此,它们的融合也被称为的 local-group feature。最后 Mugs 拉进同一张图片的两个数据增强的 local-group feature ,并且推远不同图片的 local-group feature:
其中,分别对应老师和学生的投影网络。代表学生网络的预测网络,代表一个缓存 buffer,用于存贮过往的。这样一来,正如图一中的第三个球面特征效果图所示,Mugs 将分散的特征聚集为小的聚类中心,从而学习到稍微粗粒度的 local-group 特征。
这种 local-group discrimination supervision 给 Mugs 带来两方面的好处。
- 它是上述 instance discrimination supervision 的有效补充。针对同一张图片的两个数据增强后的样本,这种监督要求这两个样本的 local-group 相同。这意味着除了样本本身对齐之外,它们对应的邻居也应该很好地对齐。这种邻居对齐一个更具挑战性的对齐问题,也增强了更高级别的 local-group 语义对齐。
- 这种监督鼓励高度相似的样本形成 local-group,并将不同的 local-groups 分散开来,从而提高了特征的语义判别能力。这是因为 a)这种监督所寻求的邻居数目 k 非常小(大约 10),这样在同一 local-group 的样本高度相似,有助于形成非常紧凑的 local-group; b)这种监督进一步将不同的 local-groups 分开。通过这两个方面,这种 local-group 通过考虑数据中的 local-group 结构,促进了更高级别的细粒度特征学习,有效地补充了 instance discrimination supervision。