论文阅读：All4One=Symbiotic Neighbour Contrastive Learning via Self-Attention and Redundancy Reduction

ICCV2023: Imanol G. Estepa et.al.
solo-learn 上刷榜的最新自监督表征学习

我必须要吐槽这篇论文的公式写得依托… log exp不正体，括号不分大小，矩阵向量元素大小写不分

摘要

基于最近邻的方法已经被证明是最成功的自监督学习方法之一，因为它们具有很高的泛化能力。然而，当考虑多近邻时，它们的计算效率会降低。在本文中，我们提出了一种新颖的对比式SSL方法，我们称之为All4One，它基于Self-attention机制创建的“质心”来优化降低近邻表征之间的距离。我们使用质心对比损失以及1近邻对比和特征对比损失。质心有助于从多个邻居中学习上下文信息，而近邻对比使得直接从近邻中学习表征成为可能，而特征对比则允许学习特征唯一的表示。这种组合使得All4One在流行的基准数据集上的线性分类评估中胜过流行的实例区分方法超过1％，并获得了最先进的结果。最后，我们展示了All4One在嵌入维度和增强方面的鲁棒性，超过了NNCLR和Barlow Twins超过5％的低维度和弱增强设置。源代码可在https://github.com/ImaGonEs/all4one 中获得。

图1：All4One的简化架构。 All4One使用三种不同的目标函数，对比不同的表示：中心对比从多个近邻中提取的上下文信息，而近邻对比确保多样性。此外，特征对比度量生成特征的相关性并增加它们的独立性。

1. 简介

自监督学习（SSL）可以在不需要大量标注数据的情况下训练深度模型。最流行的SSL方法依赖于实例辨别学习，这种策略训练模型对应用于单个图像的扭曲是不变的，这些图像被定义为正样本。由于所有视图都属于同一图像，因此它们属于同一语义类别。将它们放置在相同的特征空间中鼓励模型为相似的图像创建相似的表示。如SimCLR和BYOL的鼓舞人心的最初结果推动了多个改进，解决了实例辨别中常见的问题，如样本之间缺乏多样性和模型崩溃。

邻域对比学习依赖于一个事实，即数据增强在选择正样本时提供的多样性不足，因为所有正样本都是从同一初始图像中提取的。为了解决这个问题，最近邻对比学习（NNCLR）提出了使用最近邻（NN）来增加正样本之间的多样性，从而提高模型的泛化能力。具体来说，它们增加了一个增广样本与另一个增广样本的最近邻之间的接近程度。然而，完全依赖于最近邻会阻碍方法的真正潜力。MSF提出使用k-近邻来增加模型的泛化能力。然而，由于需要为每个最近邻计算目标函数（k次），MSF的计算量很大。除了正样本的多样性不足之外，实例辨别方法还面临模型崩溃的问题，即模型学习了一个常数平凡解的情况。Barlow Twins 提出了一种基于冗余减少的方法，通过测量生成图像表示中特征之间的相关性，自然地避免了崩溃。然而，这种避免崩溃的方法需要在高维空间中投影嵌入。

在我们的工作中，我们通过避免多次计算目标函数，以更高效的方式对多个近邻的信息进行对比。为此，我们提出使用由自注意力机制构建的新嵌入，例如Transformer编码器，将提取的近邻特征组合成一个包含有关它们所有的上下文信息的单个特征。因此，我们能够在单个目标计算中对所有邻居的信息进行对比。我们利用一个支持集来主动存储训练过程中计算得到的表示，以便我们可以提取所需的邻居。此外，我们将我们的方法与冗余减少方法集成。使计算得到的交叉相关矩阵接近于单位矩阵，减少了相同图像表示的特征冗余，同时使它们对增广具有不变性。这个想法与其他实例辨别方法完全不同地对比了表示。因此，我们通过将邻域对比方法与冗余减少目标相结合，直接对比编码器生成的特征并旨在增加它们的独立性，从而增加了模型学习到的表示的丰富程度。此外，由于我们的SSL目标组合，减少冗余的特征对比方法对高维嵌入的需求得到了缓解。

在本文中，我们介绍了一种新的共生型自监督学习方法，称为All4One，它利用邻域对比学习的思想，并将其与特征对比方法相结合。All4One集成了三种不同的目标，它们相互促进，并提供了更好的表征学习。我们的贡献如下：(i) 我们定义了一种新颖的目标函数，称为质心对比，它基于通过自注意机制在一个新的潜在空间中对样本邻居进行投影。(ii) 我们的提议All4One基于质心对比、邻域对比和特征对比目标的组合，超越了单一邻域对比，同时避免了多次计算目标函数；(iii) 我们演示了使用基于InfoNCE 的不同表示对比和特征对比目标有助于整体性能，并减轻了诸如特征对比方法对高维嵌入的依赖等个别缺陷；(iv) 我们证明了All4One通过对比邻居和多种表示的上下文信息，在低增强设置和低数据情况下优于单一最近邻自监督学习，证明了对增强的依赖性较小，增加了泛化能力；(v) 我们证明了All4One在不同公共数据集和使用不同的骨干网络下，比单一邻域对比方法等方法在各种公开数据集上表现更好，提高了1%以上。

图2：邻域对比。 虽然常见的邻域对比方法只对比最近邻，但我们创建的表示包含来自k个最近邻的上下文信息，并在单个目标计算中对其进行对比。

2. All4One

2.1 符号

我们在图3中展示了我们提出的All4One流程，其中显示了三个不同的目标：第一个目标是近邻对比目标（绿色路径）；第二个目标是质心对比目标（紫色路径），通过应用自注意机制来实现；最后一个目标是基于冗余减少的特征对比目标（红色路径）。在训练过程中，初始表示根据所跟随的路径进行转换，以使它们适应路径的目标。

图3：All4One架构。 特征对比、质心对比和邻域对比目标函数分别用红色、紫色和绿色表示。

该架构由一对编码器$f$以及一对投影器$g$组成。投影器由一个基本的MLP组成，它将编码器的输出转换为新的表征。在这种情况下，我们应用了一个动量编码器，它是在线编码器的平滑版本，类似于BYOL 。对于每个图像，使用数据增强管道生成两个增广图像$X^1,X^2$。两路增广图像都被送入编码器，然后送入投影器。编码-投影器的这个顺序根据所使用的编码器定义为动量或在线分支。我们将动量分支称为包含动量投影器的编码-投影器序列，反之亦然。动量分支和在线分支的输出分别可以定义为$Z^1=g_{\xi }(f_{\xi }(X^1))$和$Z^2=g_{\theta }(f_{\theta }(X^2))$。

2.2 近邻对比

NNCLR是最流行的近邻对比方法之一。它使用简单的KNN运算符从一个队列或支持集中提取第一个扭曲图像的最近邻，并使用 InfoNCE 的一种变体将该近邻与第二路增广图像进行对比。对于批次中的每一对图像，损失定义如下：

$$L_i^{NNCLR}=-\log \left(\frac{\exp (n{n}_i^1\cdot p_i^2/\tau )}{{\sum }_{k=1}^N\exp (n{n}_i^1\cdot p_k^2/\tau )}\right)$$

其中，$n{n}_i$ 是$z$的第 $i$个最近邻，$p_i$是第二个扭曲图像，$\tau$是温度常数，$N$ 是批次中的样本数。通过使用更多样化的样本，这种方法提高了模型的泛化性能。需要注意的是，在我们的情况中，第二个预测器$q_{\theta }^{nn}$用于近邻对比目标。

2.3 质心对比

使用单个最近邻可能限制了该方法的潜力。事实上，他们表明，将第二路增广图像与多个最近邻进行对比可以提供一个更好的自监督学习框架，该框架在选择最近邻方面也能获得更高的准确性。然而，这会影响模型的计算效率，因为它们需要 k 次计算目标函数。遵循使用多个最近邻的想法，我们介绍了一个替代性提议：从提取的 k 个最近邻中汇编相关信息，创建一对表示，称为“质心”，其中包含有关所有最近邻的上下文信息，并将它们应用在特征空间中聚集起来。这样，模型的泛化能力就可以得到改善，而无需逐个对比多个最近邻。

一旦从一对增广图像中计算出$Z^1$和 $Z^2$我们就计算每个 $Z$与支持集 $Q$ 之间的余弦相似度，支持集 $Q$存储了计算出的表征。接下来，我们提取 $K$ 个最相似表征的序列作为该增广图像中的所有表征序列：
$$n{n}_i^1=KNN(z_1^1,Q)$$
对于每个 mini-batch的$Z^1$和 $Z^2$，我们获得它们各自的最相似表示序列$N{N}^1$和 $N{N}^2$。由于我们尽量避免逐个对比近邻，我们向pipeline中引入了一个新元素：一个 transformer 编码器$\psi$。给定$N{N}^1$和 $N{N}^2$，我们将每个序列输入$\psi$以计算自注意力。

给定邻居表示序列 $n{n}_i^1$，我们获得一个单一表示 $c_1$并尽可能多地包含有关输入序列$n{n}_i^1$ 的信息。在计算自注意力时，我们以加权的方式将输入序列的表示混合，从而得到一个新的丰富表示向量。在训练期间，对于$N{N}^1$ 中的每个序列，该过程由以下步骤组成：

1. 对于 $N{N}^1$中的每个序列$Se{q}_i$，我们添加正弦位置编码
2. 然后，我们用$Se{q}_i$输入 transformer 编码器$\psi$
3. 在 transformer 编码器内部，计算自注意力并返回新序列 $Se{q}_i^c$
4. 最后，我们选择返回序列$Se{q}_i^c$中的第一个表示$Se{q}_{i,1}^c$作为我们的质心$c_i$，因为我们的目标是对比单个表示。

在在线分支上，我们遵循稍微不同的过程。类似于投影头的第二个MLP，作为预测器，用于改变当前mini-batch$Z^2$特征空间，并获得 $P^{c2}=q_{\theta }^c(Z^2)$的转换表示。更具体地说，我们将$Z^2$通过质心预测器$q_{\theta }^c$。然后，我们将 $N{N}^2$的每个序列 $n{n}_i^2$中的第五个近邻替换为$p_i^{c2}$。最后，我们重新排序每个序列，使$p_i^{c2}$成为第一个元素。我们将此过程定义为 Shift 操作。这样做是为了将扭曲的图像引入序列中，从而影响反向传播。创建修改后的$N{N}^2$后，我们通过传递它们到transformer来获得$C^2$，方法与前一段的过程相同。最后，我们使用 InfoNCE损失函数的变体来对比$C^1$和$C^2$，旨在将一对增广近邻的质心拉近。质心损失可以定义为：

$$L_i^{centroid}=-\log \left(\frac{\exp (c_i^1\cdot c_i^2/\tau )}{{\sum }_{n=1}^n\exp (c_i^1\cdot c_n^2/\tau )}\right)$$

2.3 减少冗余：特征对比

应用冗余减少原则来增加特征的独立性是自监督学习领域最成功的方法之一 。Barlow Twins的主要思想是，不专注于图像，而是通过计算特征的交叉相关矩阵$C_{i}j$来直接对比特征。它旨在通过将对角线元素$c_{ii}$设置为1（不变性项），同时通过减少不同特征之间的相关性来降低特征之间的冗余。受到这种方法的启发，我们通过引入一个特征对比目标函数来丰富我们的框架，该目标函数衡量特征之间的相关性，并旨在增加它们的独立性。通过鼓励特征之间相互不相关，我们的模型学习到更多样化和信息丰富的表示，从而提高了在自监督学习任务中的性能。

为此，我们将动量投影头和在线投影头的输出构造为$(Z{)}_{ij}$矩阵，其中每个元素表示单个增强图像表示 $i$ 的确切特征 j。然后，我们在mini-batch维度上对两个矩阵进行 ${\ell }_2$ 归一化，并计算转置矩阵 $Z^{1\top }$和$Z^2$之间的余弦相似度，以获得交叉相关矩阵$C{C}^1$。我们对称地计算$C{C}^2$：这是通过交换分支并计算相似度来完成的。最后，特征对比目标函数计算如下：

正如可以注意到的那样，第一项增加了扭曲图像表示中相同特征的元素之间的相关性（对角线元素），而第二项降低了不同特征的表示元素之间的相关性（非对角线元素）。

$$\begin{array}{r}{L}_{Redu.}=\frac{1}{2}\sqrt{\frac{1}{2D}\sum _{i=1}^D((1-c{c}_{ii}^1{)}^2+(1-c{c}_{ii}^2{)}^2)}\\ +\frac{1}{2}\sqrt{\frac{1}{2D(D-1)}\sum _{i=1}^D\sum _{j\ne i}^D((c{c}_{ij}^1{)}^2+(c{c}_{ij}^2{)}^2)}\end{array}$$

2.4 All4One 目标

一旦所有目标都计算完成，最终损失函数由之前定义的各个目标的总和组成。All4One目标定义为：
$$L_{All4One}=\sigma L_{NNCLR}+\kappa L_{Centroid}+\eta L_{Redu.}$$
其中，$\sigma$、$\kappa$和$\eta$通过损失的变化来确定，分别设为0.5、0.5和5。根据数据集的不同，给每个目标赋予更重要的权重可能会提高模型的性能。然而，用于验证的常见基准数据集是平衡的，因此我们保持三个All4One目标的权重均匀。通过结合不同的目标，All4One目标改善了表示学习。我们在下面展示了与其他方法相比的改进。

3 实验

3.1 实现细节

架构。 动量分支包括通常的动量投影器$g_{\xi }$。在线分支则包括投影器$g_{\theta }$和双预测器，$q_{\theta }^{nn},q_{\theta }^c$ 。抛去backbone不谈，MLP的投影头由大小为$[2048,2048,256]$的3个全连接层组成，而MLP预测器使用维度为$[4096,256]$的2个全连接层。除了最后一层，所有的全连接层都跟着一个批归一化。

SSL方法中的MLP用于过滤编码器生成的特征，保留对下游任务有用的特征。在All4One中，应用了多个完全不同表示的目标，因此所涉及的预测器应该独立地过滤特征。基于这个原因，我们提出使用双分支预测器，分别针对每个目标单独工作。最后，我们的方法引入了一个小型Transformer $\psi$，它应用于两个分支，并且一个共享的k-NN($\cdot$,$\cdot$)运算符。给定一个query $z$, k-NN从支持集$Q$中提取图像表示。

训练。
我们在CIFAR-10、CIFAR-100、ImageNet100（使用ResNet-18骨干网络）和完整的ImageNet（使用ResNet-50骨干网络）上训练All4One，而无需任何类标签或注释。在训练期间，所有的骨干网络都使用默认的Solo-learn初始化。所有的MLP组件都使用PyTorch默认的初始化。Transformer编码器使用每个8个头的三个Transformer编码器层。按照SimCLR的做法，通过使用lr * batch size / 256的公式调整lr，并且对每一层使用LARS进行重新调整（仅针对CNN）。对于ResNet-18和ResNet-50，我们使用SGD优化器，对于ViT-Small骨干网络，我们使用AdamW。我们发现将基础学习率分别设置为1.0、0.00015和0.1对于CNN骨干网络、ViT和内部Transformer编码器逐渐减小的温和余弦退火调度器是最优的。当使用AdamW时，我们将Transformer编码器的学习率分离，将其设置为常数0.1。对于完整的ImageNet实验，我们将学习率分别调整为1.5和0.45，用于All4One和NNCLR。所有实验都使用相同的批量大小和相同数量的epochs进行运行。最后，队列或支持集的大小设置为98304，遵循NNCLR的设置。模型的其余超参数直接从NNCLR中copy。

3.2 实验验证

3.3 消融实验

目标函数。 如前所述，我们的方法引入了冗余减少和新颖的邻域对比目标。因此，我们发现逐个分析它们的重要性非常有趣。在表3中，我们报告了框架中每个目标的性能。此外，我们还研究了EMA (Exponential Moving Average)（版本2与版本3）的重要性。我们发现EMA的添加提高了模型的整体性能。

从不同版本的All4One可以看出，所有三个目标都对整体性能至关重要。直觉上，NNCLR和基于质心的目标都专注于对比邻域图像表示，因此在训练过程中，两个目标可能部分重叠。而冗余减少不同，它专注于特征，因此其移除比其他目标更为关键。此外，我们设计冗余目标使用表示$z_i^1$和$z_i^2$来计算损失，而不是使用邻域。这一事实为最终损失函数增加了更多丰富性，因为总共使用了三种不同的表示（原始邻域$n{n}_i^1$、从邻域派生的质心$c_i^1$和图像表示$z_i^1$）进行统一目标。

另一方面，我们发现质心目标是在NN检索准确率上增加最多的目标，当与NNCLR和EMA架构（实验v3）结合时达到了86.16%。正如预期的那样，对比来自多个邻域的上下文信息鼓励模型创建更好的表示，只需使用简单的KNN运算符即可更容易地区分。然而，引入不同的目标，如冗余减少目标，迫使模型更广泛地推广，因此在下游任务上表现更好。

维度。 诸如Barlow Twins之类的冗余减少方法高度依赖于嵌入的维度。其他方法，如Opt-SSL，需要高维嵌入才能提供SOTA结果。然而，我们的方法却能够以更低维的嵌入表现出色，因为它不仅依赖于冗余减少。即使嵌入的维度较低，所形成的共生关系也避免了性能下降，如表4a所示。考虑使用低维度的另一个因素是$KNN(\cdot ,\cdot )$受到维度灾难的影响，当使用高维度嵌入时，其效率会降低。

数据增广。 SSL框架严重依赖于增广形式以创建难以区分的正样本 。然而，邻域对比方法经验性地证明它们在去除增强方面更具鲁棒性 。NNCLR提出除了随机裁剪增强外，移除所有增强，以将NNCLR框架的鲁棒性与SimCLR和BYOL框架进行比较。我们也遵循相同的方法来检查All4One相对于NNCLR和Barlow Twins的鲁棒性。正如表4b所示，我们的方法相对于Barlow Twins和NNCLR在去除增强方面表现出更高的鲁棒性。直观地说，添加一个使用多个邻域的项会增加损失函数的丰富程度，使增强变得不那么重要。

3.4 讨论

尽管All4One取得了令人期待的结果，但我们确定了一些局限性。

1. 计算效率。 由于引入了三个不同的目标函数，All4One比使用多个邻居的NN方法更有效，但比仅对比单个邻居的NNCLR [14]效率低。这在低计算约束下是一个限制。我们在补充材料中提供了计算复杂度分析。
2. 超参数数量增加。 最终的All4One目标函数引入了3个额外的超参数进行调整。此外，变压器编码器使用不同的学习率，这也为整体框架增加了一个额外的超参数。
3. 变压器编码器参数。 在训练变压器方面有几项进展。变压器编码器在使用不同设置时的稳定性与其他组件相比尚不清楚。

4. 结论

我们提出了一种共生方法，通过自注意力有效地对比来自多个邻居的上下文信息来利用NN对比学习。此外，我们整合了有利于整体框架的特征对比目标函数。All4One证明了更好地泛化并提供更丰富的表示，通过整合不同目标的方法优于以前的SoTA对比方法。这突显了其在低数据情境、低维度情境和弱数据增强设置下的出色性能。在未来，我们计划将All4One扩展到更复杂的主干网络，并研究其在各种下游任务中的应用，如目标检测和实例分割。