MSRA提出学习实例和分布式视觉表示的极端掩蔽模型ExtreMA

本篇分享论文『Extreme Masking for Learning Instance and Distributed Visual Representations』，MSRA&CMU提出学习实例和分布式视觉表示的极端掩蔽模型ExtreMA，图像重建效果甚至由于MAE！

详细信息如下：

• 论文地址：https://arxiv.org/abs/2206.04667

• 代码地址：尚未开源

      01      

摘要

本文提出了一种可扩展的方法，用于同时学习单个token上的分布式表示和整体实例表示。作者使用 self-attention 块来表示分布式token，然后使用 cross-attention 块来聚合整体实例。该方法的核心是使用极大的token masking（75%-90%）作为监督的数据增强。

本文的模型名为 ExtreMA，遵循简单的 BYOL 方法，其中来自未mask子集的实例表示被训练以从完整输入中预测。学习需要模型捕捉实例中的信息变化，而不是鼓励不变性。

本文做出了三个贡献：

1. 随机mask是一种强大且计算效率高的数据增强，用于学习可概括的注意力表示。
2. 每个实例有多个采样，极端masking极大地加速了学习和对更多数据的需求。
3. 分布式表示可以仅从实例监督中学习，这与掩蔽建模中的每个token监督不同。

      02      

Motivation

掩蔽建模（Masked modeling）已成为视觉表示学习的一种可行方法。在通用Transformer架构上，它基于掩蔽信号预测来优化学习目标，而不依赖于经过大量工程设计的图像增强。这种方法已经证明了卓越的微调性能。

然而，掩蔽建模和基于孪生网络的对比预训练方法还是存在显著的差距。用对比目标训练的孪生网络非常适合学习现成的表示。这种根本区别在于它们表示数据的方式。孪生网络为图像提取实例表示，而掩蔽建模在组成图像的各个token上获取分布式表示。在掩蔽建模方法中，没有实例表示被显式建模或提供监督。

在本文中，作者旨在研究实例与分布式表示之间的联系，并探索自监督方法以学习这些表示。作者首先观察到随机mask可以被视为一种新的数据增强方案，以前没有在孪生网络中使用过。对于mask区域，其潜在的自由度随着其大小的组合而增大，允许比裁剪和缩放等传统增强功能更丰富的自监督，后者严重偏向图像中心周围的区域。

更重要的是，传统增强的自监督导致了一个包含实例的多个增强的通用表示向量，这种不变性降低了对空间局部性表示的灵敏度。相反，随机mask保留了未mask区域的原始内容，并且不会改变数据的几何结构。

本文提出了一个简单的模型，称为 ExtreMA，其中来自mask输入的实例表示被训练为从完整视图中预测（以BYOL方式）。mask造成的信息差距鼓励学生网络编码尽可能多的信息，从而引导教师网络变得更强大。作者采用 ViT 在patch上嵌入分布式表示，然后是交叉注意块将分布式表示聚合到实例表示中。实例级学习目标为两种表示提供监督。

在本文的模型中，分布式表示仅是隐式学习的，没有掩蔽建模中使用的相应token级目标。然而，通过研究和可视化注意力图，作者发现本文模型的输出表示与输入token保持准确的对应关系，并且语义集群倾向于从学习的分布式表示中出现。

该模型的一个显着区别是它的有效性具有极高的掩蔽率（75%-90%），而典型的掩蔽建模方法在 50% 到 75% 的范围内效果最佳。这除了带来的计算效率之外，极端掩蔽（extreme masking）的一个关键方面是多个样本之间出现的互补性。

由于视觉数据的高冗余，不同掩蔽样本的可见内容只有在掩蔽率变大时才会变得独立。此外，多个mask显着加快了学习和收敛速度，使系统成为渴望更多数据的快速学习模型。在实践中，多重采样在计算上也很有吸引力，因为用于处理全部内容的教师网络只需要forwarded一次。

ExtreMA 享有孪生表示学习的有利特性。模型中的实例表示可以现成地用于测量语义相似性。除了为具有不同目的的应用程序mask之外，该框架可以进行其他数据增强。然而，与典型的对比学习不同，ExtreMA 不依赖数据增强诱导的不变性来实现泛化。相反，ExtreMA 保留了mask视图中所有可能的有用信息，以便恢复完整信息。

在实验中，作者系统地研究了模型在不同掩蔽率下的行为，在更大的数据集上使用多个mask的收敛特性，以及与各种其他数据增强的集成。基于研究观察，作者还提出了一种新的增强方案，该方案使用共享的图像裁剪，但两个输入视图使用不同的颜色。

本文方法在 ImageNet1k 上的主要结果在微调和linear probing指标方面都优于先前的掩蔽建模方法。此外，作者还评估了半监督学习和语义分割的迁移性能。对于这两种应用，与现有技术相比，ExtreMA 产生了更好的结果。

      03      

方法

3.1. ExtreMA

这项工作探索了使用随机mask作为实例级监督的孪生表示学习的数据增强。该模型使用两个并行网络的架构，其中动量编码器处理完整图像裁剪，基本编码器处理mask图像裁剪。信息间隙，即掩蔽图像区域，是训练基础编码器的监督基础。本文的方法的概述如上图所示。

Extreme Masking

给定输入图像，首先将其划分为不重叠的patch，然后输入到视觉Transformer中。一个固定的正弦位置编码被添加到每个嵌入的patch中，并且根据掩蔽率随机采样一些嵌入的patch。

本文方法的一个关键方面是它以 75%-90% 的极大掩蔽率实现了最佳性能，使基本编码器只处理一小部分 (10%-25%) 的patch。这与掩蔽图像建模相反，掩蔽图像建模在掩蔽率超过 75%时性能会下降 。极高的掩蔽率为网络设置了一个非常困难的借口任务（pretext task）。

ExtreMA 能够成功进行极端掩蔽（extreme masking）的部分原因是动量编码器处理了整个图像，而掩蔽图像建模中的编码器从未接收到完整视图的输入。这会为掩蔽建模产生训练测试差异，因为注意力块需要从处理一小部分token到完整集进行泛化。

极端掩蔽提供了另外两个关键好处。首先，非常大的掩蔽率会导致来自同一图像的独立和互补视图。这显着加快学习和收敛速度。其次，极端掩蔽大大降低了基本编码器的计算成本，尤其是对于具有平方复杂度的Transformer。当使用多个mask时，它们可以共享来自动量编码器的相同学习目标。这使得极端掩蔽成为一个高效的学习器。

极端掩蔽也带来了新的挑战。当掩蔽率增加时，掩蔽增强的空间组合变得更小。在实践中，作者观察到极端掩蔽往往会过度拟合训练数据，尤其是在启用多掩蔽（multi-masking）时。这与“cheating”不同，因为该特征有效地学习了训练数据的语义。过拟合现象只是表明极端掩蔽需要更多的训练数据。

Distributed and Instance Representations

作者采用视觉Transformer是因为它在处理可变大小的输入内容方面的效率和灵活性。视觉Transformer通过自注意机制将输入的视觉token嵌入到分布式表示中。为了允许从实例级别进行监督，需要实例表示。

为了实现这一点，作者使用交叉注意力块将分布式patch级表示聚合成一个带有附加CLS token的单一表示。分布式表示在没有更新的情况下被冻结，这使得交叉注意力块既高效又轻量，投影头和预测器遵循实例表示。

作者已经研究了另外两种表示实例的替代方案，但都不是很好。如果像在 ViT 中一样将实例token作为输入提供给Transformer，优化变得不稳定，这可能是由于输入的大掩蔽率。通过将token表示作为实例表示进行平均池化，该模型将找到学习平均token特征的捷径，而无需跨patch学习注意力。

Learning Objective

来自掩蔽输入的实例表示被训练以通过简单地最小化两个表示之间的余弦距离来预测来自未掩蔽输入的实例表示。在本文中，为了简单起见，作者主要遵循 BYOL，但在本文的实验中，本文的方法也适用于对比损失。本文的学习目标在以下方面与传统的 BYOL 不同。

传统的 BYOL 采用对称损失，其中两个视图被学习以相互预测。这将驱动表示找到两个视图共享的公共子空间，并且对其他信息保持不变。在本文的方法中，由于两个视图之间的共享信息是显而易见的，因此找到视图之间的共性并不能实现有意义的自监督。因此，作者采用不对称损失。掩蔽造成的信息差距鼓励网络推断被掩蔽的区域，而不是寻找共享的子空间。

与掩蔽建模方法相比，本文的学习目标不同。实例级监督没有明确地对每个单独的token强制执行空间推理。尽管如此，作者发现强有力的证据表明本文的模型学习了token上的分布式表示。作者将上图中Transformer块最后一层中四个查询块的注意力图可视化。显示的可视化是 12 个注意力头的平均值。作者观察到patch token倾向于分组为有意义的语义集群。

BYOL Details

作者对投影和预测头的设计选择比原来的 BYOL还要简单。作者用 LayerNorm 替换 BatchNorm，用 GeLU 激活替换 ReLU 激活，使整个框架没有 BatchNorm 并且与其余的Transformer块保持一致。

本文的工作还顺便证明了 BYOL 不依赖 BatchNorm 来防止崩溃。投影头和预测头分别有3和2个隐藏层，隐藏维度为4096，输出维度为256，遵循原设计。

3.2. Representation Properties

为了理解分布式表示和用极端掩蔽增强训练的实例表示，作者将分布式表示反转到像素空间并检查实例表示对局部性的敏感性。这些结果进一步证明该模型在没有类似 BERT 的目标的情况下学习了有意义的分布式表示，并且实例表示保留了详细的视觉信息。

Generative Properties

给定预训练模型和掩蔽图像，可以使用分布式表示对可见patch进行编码，并将这些部分表示反转回像素空间以重建掩蔽patch。具体来说，遵循之前的工作，作者最小化掩蔽图像和重建图像之间可见表示的 L2 距离。这种重建技术能够检查编码特征的内容，而无需进一步训练新的生成模型。上图显示了重建结果。

作者使用这种反转方法改变掩蔽率并将结果与有监督的 ViT、DINO 和 MAE 进行比较。受监督的 ViT 无法生成任何有意义的内容。由于使用了颜色增强和归一化像素，DINO 和 MAE 都无法修复正确的颜色。

此外，DINO 对空间局部性也不准确。本文的结果在空间上是平滑的并且颜色准确。当掩蔽率非常大时，由于修复看不见的语义区域的能力有限，反转技术会受到影响。

Locality Properties

作者使用 k-最近邻技术来探测实例表示。首先通过对单个图像的空间和规模变化的图像裁剪进行随机采样来生成一个小型gallery集。然后，使用来自同一语义类别的另一个查询图像裁剪来对gallery集进行排名。作者将这些图像裁剪的大小调整为 224 × 224，并提取实例表示以测量相似度。

在上图中，顶部最近检索返回具有最接近空间和比例配置的图像边界框作为查询裁剪。结果表明，实例表示对空间和尺度变化很敏感，并且学习到的高级表示不是不变性的结果，而是更强大和泛化。

      04      

实验

作者首先使用单个mask来改变掩蔽率进行训练。在上表中，微调性能在 70% 到 90% 的宽范围内保持稳定，而linear probing性能在 80% 的比率处达到峰值。值得注意的是，90% 的极大掩蔽率也实现了相当好的性能，超过 90%之后性能下降。

对于每个图像实例，作者生成多个mask而不替换学生网络。损失是在多个mask输入上平均的。在上表2中，作者研究了75%、80% 和 90% 比率下的多重mask行为。随着更多的mask输入，微调性能不断提高。然而，当使用太多mask时，linear probing性能会下降，尤其是当掩蔽比变大时。

为了在保持快速学习的同时对抗过拟合，一个简单的解决方案是使用更大的数据集。因此，作者研究了 ImageNet22k 上的multi-masking码，它在总图像中大约大 10 倍。作者对模型进行了 30 个 epoch 的训练，它保持了有效的优化迭代次数，并仅从数据规模上揭示了影响。

微调和linear probing的评估均使用 ImageNet1k 进行。在表 3 中，在掩蔽率为 90% 时，ExtreMA 不再随着mask数量的增加而出现过拟合。模型性能也始终优于使用 ImageNet1k 训练数据。这表明本文的模型需要更多数据并从大规模表示学习中受益。

作者还检查添加额外数据增强的情况，其中教师分支的输入被增强，学生分支采用教师输入的随机mask样本，结果如上表4所示。上表5展示了不同cross-attention块数量的实验结果。

除了 BYOL，ExtreMA 还与其他孪生表示学习目标一起使用。在上表 6 中，作者使用 0.2 的对比温度提供了 MoCov3实现的结果。与 BYOL 目标相比，微调性能提高了 0.2%，linear probing下降了 2.0%。

上表展示了本文方法和其他自监督学习方法在ImageNet1K分类任务上的对比结果。

在上表中总结了代表性模型的wall clock时间。

上表展示了半监督分类的结果，ExtreMA 在使用 1% 的标签时比蒙版图像建模方法 MAE 和 BEiT 的性能高出 12% 和 30%。

上表展示了本文方法在语义分割上的实验结果。

      05      

总结

这项工作探索了mask作为孪生表示学习的一种新的增强方法。所研究的方法 ExtreMA 通过数据增强来学习强大的实例和分布式表示，而无需掩蔽建模监督。这项工作的灵感来自于掩蔽建模中的mask操作。但是， ExtreMA 的工作方式与掩蔽建模不同，因为学习目标不同。 

ExtreMA 表现出几个独特的特点： 

1. 使用极大的掩蔽率，75%-90%； 
2. 收敛速度快，具有multi-masking和大数据的可扩展性；
3. 计算资源消耗低，它为实例表示编码精确位置的能力可能会开辟新的可能性。

参考资料

[1]https://arxiv.org/abs/2206.04667

已建立深度学习公众号——FightingCV，欢迎大家关注！！！

ICCV、CVPR、NeurIPS、ICML论文解析汇总：https://github.com/xmu-xiaoma666/FightingCV-Paper-Reading

面向小白的Attention、重参数、MLP、卷积核心代码学习：https://github.com/xmu-xiaoma666/External-Attention-pytorch