【算法论文|自监督1】CNN-JEPA: Self-Supervised Pretraining Convolutional Neural Networks

CNN-JEPA：使用联合嵌入预测架构的自监督预训练卷积神经网络

摘要

自监督学习（SSL）已成为预训练大型神经网络的一种重要方法，可实现前所未有的模型和数据集大小的扩展。虽然I-JEPA等最新进展已经显示出Vision Transformers的良好结果，但将此类方法应用于卷积神经网络（CNN）带来了独特的挑战。在本文中，我们介绍了CNN-JEPA，这是一种新颖的SSL方法，成功地将联合嵌入预测架构方法应用于CNN。我们的方法结合了一个稀疏CNN编码器来处理掩码输入、一个使用深度可分离卷积的完全卷积预测器，以及一个改进的掩码策略。我们证明，CNN-JEPA在ImageNet-100上优于采用ViT架构的I-JEPA，使用标准ResNet-50编码器实现了73.3%的线性top-1精度。与其他基于CNN的SSL方法相比，CNN-JEPA在相同数量的epoch下需要的训练时间减少了17~35%，并接近BYOL、SimCLR和VICReg的线性和k-NNtop-1精度。我们的方法为现有的CNNSSL方法提供了一种更简单、更高效的替代方案，只需要最少的增强，并且不需要单独的投影仪网络。

关键词：自监督学习、表示学习、卷积神经网络、ImageNet、深度学习

图 1.比较CNN-JEPA（我们的方法）与I-JEPA[4]以及基于ImageNet-100上的线性top-1准确性和训练成本的常见SSL方法。标记的面积与模型中的参数数量成正比

1、引言

自监督学习（SSL）成为预训练大型神经网络的重要应用，它允许将模型和数据集大小扩展到以前不可能的水平。它是最近大型语言模型快速发展的驱动机制之一。自我监督预训练提出了一种新的范式，旨在学习一般表示，而不是学习特定于任务的特征。在计算机中，视觉设计良好的自我监督借口学习任务是一个具有挑战性的问题，文献中提出了许多方法。其中一种方法是掩码图像建模（MIM），其中输入图像的补丁被掩码，并且网络经过训练以预测掩码补丁。这种方法主要是针对视觉转换器（ViT）提出的[1]，但也被引入用于视觉神经网络（CNN）[2]、[3]。一种最近成功的方法是I-JEPA[4]（Image Joint Embedding Predictive Architecture），它基本上是通过从提供的上下文补丁中预测掩码补丁的潜在表示，在潜在空间中执行掩码建模。这种方法特别适合Vision Transformers，它可以为大型数据集提供良好的扩展，并允许简单地实现遮罩。
尽管 Vision Transformer越来越受欢迎，但卷积神经网络在许多任务上仍然提供有竞争力的性能，并在实践中得到广泛应用。它们特别高效，并且在CNN的归纳偏差具有优势的中小型数据集上提供更好的性能。然而，使I-JEPA适应CNN并非易事，因为掩码和密集特征图的预测都不是简单的。计算掩码输入的特征映射需要掩码卷积和掩码的仔细构造，这要考虑网络的下采样。预测掩蔽补丁也具有挑战性，因为它需要足够大的接受性field，能够根据提供的上下文预测潜在表征。典型的CNN通过堆叠许多卷积层来实现大的接受性f，这对于具有高深度的特征图来说可能是非常繁重的参数和计算成本高昂的。
在这项工作中，我们提出了一种新的自我监督学习方法，即 CNN-JEPA（基于卷积神经网络的联合嵌入预测架构），它解决了使I-JEPA适应CNN的挑战。Tian et al. [3] 提出的掩码图像建模方法使用基于网络下采样构建的掩码卷积和掩码。我们将这种稀疏CNN编码器引入I-JEPA以应对掩码的挑战。我们还引入了一种新颖的、完全卷积的预测器，它使用深度可分离卷积，可以预测掩码补丁的潜在表示，参数数量和计算成本低。此外，我们改进和简化了I-JEPA的掩蔽策略，方法是基于上下文区域作为预测因子的输入，将掩蔽斑块预测为单个区域，而不是像Assran等人[4]那样使用多个目标区域。我们的主要贡献是：
1）一种新颖的自我监督学习方法CNN-JEPA，它将成功的I-JEPA方法应用于卷积神经网络。
2）稀疏CNN编码器架构，可以正确处理掩蔽输入，并基于Tian等人提出的掩蔽CNN生成掩蔽特征图[3]。
3）引入了一种新颖的、使用深度可分离卷积的全卷积pre-dictor。
4）一种掩码策略，它遵循网络的下采样，并使用单个目标区域进行掩码预测任务。

支持本研究结果的源代码可在：https://github.com/kaland313/CNN-JEPA

2、相关工作

文献中已经提出了许多SSL方法。我们建议将与我们的工作最相关的方法分为两大类：实例判别方法和基于重建的方法。实例判别方法，如BYOL[5]、SimCLR[6]和VICReg[7]，通过最大化同一图像的增强视图之间的相似性来学习表示。这些方法采用不同的机制来避免琐碎的解决方案，包括使用负样本、不对称架构、动量编码器和正则化。
基于重建的方法主要依赖于掩码图像建模前置任务，其中输入图像的某些部分被掩码，并且网络被训练以预测掩码区域的像素。这些方法最初是针对视觉转换器[1]提出的，但也适用于复合神经网络[2]，[3]。早期的工作应用了其他基于重建的借口任务，例如从灰度版本预测图像的颜色[8];但是，遮罩图像建模已被证明更有效。
实例判别方法学习图像级特征，这些特征与用于生成视图的增强无关，旨在捕获图像的高级语义。另一方面，这些方法严重依赖于手工制作的增强，并且没有经过训练来捕捉局部细节。相比之下，基于重建的方法学习了更多的局部特征，这些特征捕获了解决重建任务所需的局部详细信息;但是，他们的Pixel Reconstruction目标函数可能无法捕获图像的高级语义。Assran等人[4]提出了I-JEPA，这种方法通过在潜在空间中执行掩蔽图像建模来结合两种方法的优点，其中网络被训练以预测掩蔽补丁的潜在表示。I-JEPA主要是为视觉转换者设计的，据我们所知，文献尚未对其进行改编为卷积神经网络。

3、方法

A. CNN-JEPA算法

所提出的CNN-JEPA算法的概述如图2所示。所提出的方法通过屏蔽输入图像的补丁并使用预测器预测被掩盖的补丁的潜在表示来训练编码器。目标潜在表示由教师编码器获得，其权重使用学生编码器权重的指数移动平均值进行更新。我们为学生编码器引入了一种稀疏CNN架构，它可以正确处理掩码输入并生成掩码特征图。对于训练，我们向学生和教师编码器提供相同的图像，掩码仅应用于学生编码器的输入。我们将被掩盖的图像和特征称为上下文，预测器在此基础上在被掩盖的位置产生潜在表示。正如Assran等人[4]所提出的那样，在将掩蔽的潜在表示提供给预测器之前，学生特征图中的掩蔽位置会填充一个可学习的掩码标记。然后，这些特征图将传递给预测器，该预测器是预测掩码补丁的潜在表示的CNN。训练目标是最小化预测的潜在表征和蒙蔽位置的目标潜在表征之间的L2损失。不会为未屏蔽的补丁计算损失，而是使用上下文掩码的补充作为损失计算的掩码。

图 2.CNN-JEPA方法概述。学生编码器和教师编码器共享一个通用架构，例如ResNet-50，学生编码器使用稀疏卷积。学习目标是使用预测器预测掩码补丁的潜在表示，该预测器也会训练学生编码器。

B. 学生和教师编码器

我们提出的方法中的学生和教师网络共享一个通用的架构，类似于其他采用不对称学生-教师机制的自我监督学习方法（如BYOL[5]、DINO[9]和I-JEPA[4]）。然而，我们方法的一个关键区别是，结构网络利用稀疏卷积，正如Liu等[10]所引入的，后来又被Tian等[3]用于掩蔽图像建模。这种稀疏卷积架构对于处理掩码输入和生成掩码特征图是必需的。在实践中，我们通过掩盖标准卷积层的输出来实现稀疏卷积，如IV-A节所述。这允许根据学生权重轻松更新teacher网络的权重。
教师网络采用标准卷积。根据Assran等人[4]，学生网络使用基于梯度的优化技术进行训练，而教师网络的权重则通过学生权重的指数移动平均值进行更新。此更新机制可确保教师网络为学生提供稳定的目标表示。我们还将 student 和 teacher 编码器分别称为上下文编码器和目标编码器。

C. 预测器

我们提出的方法中的预测器在特征图上运行并预测掩码补丁的潜在表示。特征图通常具有高深度（例如，ConvNeXts [11] 和 ViT-Base [12] 为 768，ResNet-50 为 2048 [13]），但空间分辨率较低（例如，在 224x224 图像上使用 ResNet 50 时为 7x7）。由于预测器的输出必须与特征图具有相同的维度和形状，因此其输出也具有这样的深度和空间大小属性。参数的数量随具有相同输入和输出形状的卷积层的深度呈二次方缩放。因此，对于高深度，它们需要大量参数并且计算成本高昂。为了解决这个问题，我们建议使用深度可分离卷积，它比标准卷积更高效，计算成本更低。使用标准卷积时，预测器中的参数数量很容易超过编码器中的参数数量，从而使训练效率低下，因为预测器上花费的资源比编码器上花费的资源多。我们使用深度可分离卷积有助于缓解这个问题，确保预测器和编码器之间的计算资源分配更加平衡。

D. 掩码标记

学生特征图中的掩码位置在传递给预测器之前，会用可学习的掩码标记填充，类似于掩码图像建模 [1] 和 I-JEPA [4] 使用掩码标记的方式。虽然这种掩码令牌的使用在视觉转换器中很常见，但 SparK [3] 也将其用于 CNN 的掩码自动编码。在我们的例子中，掩码标记是一个参数向量，其维度与特征图的深度相同，并且在小批量的所有掩码位置和样本之间共享。

E. 蒙版

在蒙版图像建模中，最常见的蒙版策略是随机蒙版，其中将补丁均匀采样到给定的蒙版比率。I-JEPA [4] 介绍了多块掩码方案，其中对多个块的块进行采样，块内的块被掩码，而块外的块形成预测器预测潜在表示的上下文。I-JEPA 提出了一项关于随机、单块和多块掩蔽的实证研究，并表明后者在很大程度上是最有效的。我们将多数据块掩码策略应用于 CNN，但我们处理掩码补丁的预测方式不同。虽然 I-JEPA 分别预测每个块的潜在特征，但我们将块的并集视为掩码区域，并预测整个区域的潜在表示。对于 CNN，后一种方法更易于实施，并且在学习方面相当。预测的上下文是我们的工作和 I-JEPA 中任何蒙面补丁之外的补丁。
当使用 CNN 的掩码图像建模时，掩码大小由网络的缩减采样决定。ResNets和ConvNeXts将分辨率减半特征图分为5个阶段（通过池化），导致 32 倍的下采样。为了保持掩码与特征图分辨率一致，我们使用补丁大小为 32×32 像素的掩码，这对应于最终特征图的 1×1 空间维度（“像素”）。对于分辨率为 224×224 的图像，特征图具有 7 × 7 个空间维度，具有匹配的 7×7 个 32×32 像素块的掩码。对于 ViTs，补丁大小由变压器的补丁大小决定，通常为 16 × 16 或 14 ×14 像素。

F. 与其他方法的比较

由于其联合嵌入预测架构，我们的方法学习了补丁级特征，这在图像级和像素级特征之间提供了很好的权衡。与实例判别算法相比，这种方法使模型能够捕获更多的局部特征，同时还提取了比使用掩码图像建模方法获得的更高级别的特征。
与常见的实例判别 SSL 方法相比，我们提出的方法的一个关键优势是它的简单性。它只需要最少的手工制作增强，仅依赖于随机调整大小的裁剪和蒙版。
此外，与大多数 SSL 方法不同，我们的方法不需要单独的 projector 网络，从而进一步简化了架构并减少了计算开销。许多 SSL 方法学习的特征不受各种增强的影响，因此需要在编码器之后使用 projector 网络。使用这种投影仪可以让编码器学习更多的通用特征，而投影仪可以学习特定于损失函数和 SSL 方法增强的映射。与 I-JEPA [4] 类似，我们的方法可以学习一般特征，而无需这样的投影仪网络。

图 3.使用多数据块和随机掩码策略生成的掩码图像示例。多块方法在矩形块中屏蔽多个连接的区域。随机蒙版从图像中均匀地对色块进行采样，直至达到给定的蒙版比率。我们采用两种掩码策略，对给定的小批量图像使用多块掩码的概率为 0.75。

4、实验

A. Encoder

我们使用 ResNet-50 [13] 编码器进行实验，因为它在 SSL 文献中很常见。稀疏卷积是通过使用适当升级到层激活维度的掩码将每个卷积层的输出归零来实现的。虽然这在计算效率上很低，但它很容易实现，并且可以使用针对许多 GPU 架构进行了优化的标准卷积层。使用稀疏卷积层需要一个专门的 GPU 库（例如 Minkowski 引擎 [14]），这将使我们方法的实现更加复杂且更难重现。此外，使用标准卷积可以简化教师编码器权重的指数移动平均更新。

B. Predictor

对于预测器，我们使用 3 个深度可分离的卷积层，内核大小为 3×3，批量归一化 [15] 和 ReLU 激活。预测器的层遵循以下模式：3×（DepthwiseConv-PointwiseConv-BatchNorm ReLU）。我们对预测器结构进行了消融研究，并表明这种配置是最有效的。

C. Datasets

我们在 ImageNet-100 和 ImageNet 1k [16] 数据集上进行了实验。对于缩小 32 倍的编码器架构（例如 ResNet-50），补丁大小必须至少为 32x32 像素，因此图像必须足够大于 32x32 像素，以便为掩码嵌入预测任务提供有意义的目标和上下文区域。因此，我们使用来自 ImageNet-1k 和 ImageNet-100 的 224x224 像素图像。此外，由于这些原因，我们不会使用较低分辨率的图像和数据集（如 CIFAR-10 或 STL-10）进行实验。

D. Classification Probes

为了评估学习到的表示的质量，我们在编码器的冻结特征之上使用线性和 k 最近邻 （k-NN） 分类探针。我们训练线性探针的全连接层 90 个 epoch，并报告最佳的 top-1 和 top-5 精度。对于 k-NN 分类，我们使用相同的特征，并报告 top-1 和 top-5 精度。对于这两个探针，我们使用最少的增强：确定性的大小调整和裁剪以及归一化。我们改编了 Lightly [17] 中两个探针的超参数。

E. Training Details

我们使用 AdamW 优化器，其恒定权重衰减为 0.01。使用带有预热的余弦学习率调度器10 个时期，达到 0.01 的峰值学习率。每个设备的批处理大小为 128，具有 4 个 GPU，因此总批处理大小为 512。我们使用其他 SSL 方法 [4]、[5] 中常见的指数移动平均方法更新教师编码器的权重，其中momentum参数从 0.996 增加到 1.0。在 ImageNet-100 上进行 200 个时期的训练，在 ImageNet-1k 上进行 100 个时期的训练。所有预训练运行都在集群上的 4NVIDIAA100 40GBGPU 上执行，每个节点具有 4 个 GPU、256GBRAM 和 64 个 CPU 内核。在此硬件上，在 ImageNet-100 上进行 200 个纪元的预训练需要 13 小时，而在 ImageNet-1k 上进行 100 个纪元的预训练需要 70 小时。我们使用 PyTorch、PyTorch Lightning [18] 和 Lightly[17]SSL库来实现基线、探测和our方法。我们从官方实现中获得 I-JEPAViT结果[4]。本研究中所有实验的累积运行时间为 10 000 小时，估计总排放量为 1000 kgCO2eq1

五、实验结果

A. 与其他方法的比较

我们的主要结果如图 1 和表 I.我们将我们的方法与使用 I-JEPA 训练的视觉转换器以及为 CNN 发布的常见 SSL 方法（如 SimCLR、BYOL 和 VICReg）进行了比较。图 1 显示，onImageNet-100 方法的性能优于 I-JEPAwithViT Small 和 ViT-Base，使用标准 ResNet-50 编码器实现了 73.3% 的线性 top-1 精度，这一结果与基于 CNN 的 SSL 方法具有竞争力。此外，由于更简单、无投影仪的架构和仅使用基本增强，我们的方法在计算上比其他基于 CNN 的方法更有效，与 SimCLR 和 BYOL 相比，需要的预训练时间分别减少了 17 到 35%.AsshownintableI，我们的方法也根据 formswellonImageNet-1k，在 100epochs 中实现了 54.23%的线性顶部精度。

B. 遮罩策略

我们实验了不同的掩蔽策略，并将我们的结果显示在表 II 中。随机掩码，常见的最 in-maskedautoencodingmethods（例如 [1]–[3]），性能不佳，这与 I-JEPA [4] 的发现一致。我们尝试了不同的掩蔽策略，并提出了Bojanowski等人[4]提出的多块掩蔽策略，该策略在卷积网络中表现良好。.此外，我们提出了一种混合掩码策略，其中我们为每个小批量在多块和随机掩码策略之间随机选择。为了主要使用多块掩码，我们选择使用多块掩码的概率来 be0.75.As 如表 II 所示，混合掩码策略的性能优于随机和多块掩码，优于后者的小边际线性和 k-NN 分类。

C. 预测器架构

我们在预测器架构上进行消融，并呈现结果见表III至V。表III清楚地显示了在预测器中使用深度可分离卷积的好处。深度可分离卷积不仅比标准卷积需要的参数少 90%，训练时间缩短了 17%，而且在 ImageNet-100 上将线性和 k NN top-1 精度分别提高了 0.62% 和 1.5%。
预测变量的深度也会对预训练性能产生重大影响。如表 IV 所示，两层或三层是最优的，即使此类预测变量的感受野比更深的预测变量小。在 ImageNet-100 上，线性和 k-NN top-1 精度表示不同的最佳深度;然而，在 ImageNet-1k 上，三层预测器提供了最高的准确性，具体来说是 6% 和 3%。表 V 显示了预测变量的核大小的结果。我们认为 3x3 内核提供最高的准确性;但是，所有测试的内核大小之间的差异小于 3%。

6、结论

在本文中，我们提出了 CNN-JEPA，这是一种新颖的自我监督学习方法，它成功地将 I-JEPA 应用于对话网络。我们的方法通过使用稀疏 CNN 编码器和精心设计的屏蔽策略，克服了将屏蔽输入应用于 CNN 的挑战。为了有效地处理联合嵌入预测任务，我们的方法利用了一个完全卷积的预测器，具有深度可分离的卷积。我们证明，CNN-JEPA 在 ImageNet-100 上优于 I-JEPA 和 ViT 架构，并且在 ImageNet-100 和 ImageNet-1k 上与其他基于 CNN 的 SSL 方法相比具有竞争力，分别实现了 73.26% 和 54.23% 的顶级线性准确率。我们提出的混合掩蔽策略，结合了多块和随机掩蔽，显示出比单独使用任何一种策略都有所改善。对预测器架构的消融研究揭示了使用深度可分离卷积和预测器的最佳深度设置的好处。CNN-JEPA 为 CNN 提供了一种更简单、更高效的自我监督学习方法，需要最少的增强，并且不需要单独的投影仪网络。