对比语言图像预训练（CLIP）因其卓越的零样本性能和优秀的迁移能力而受到了广泛关注。然而，训练这样的大型模型通常需要大量的计算和存储，这对一般用户（拥有消费级计算机）来说是一个障碍。
为了应对这一观察，本文探讨了如何在仅使用一块 Nvidia RTX3090 GPU 和一兆字节存储空间的情况下实现竞争性能。
一方面，作者简化了 Transformer 块结构，并将权重继承与多阶段知识蒸馏（WIKD）相结合，从而减少了参数并提高了训练和部署期间的推理速度。另一方面，面对小数据集带来的收敛挑战，作者为每个样本生成合成字幕作为数据增强，并设计了一种新颖的配对匹配（PM）损失，以充分利用正负图像文本对之间的区分。
大量实验表明，作者的模型可以在数据量-参数-精度之间实现新的最先进权衡，这可能进一步在相关研究社区中普及 CLIP 模型。
1 Introduction
预训练的大型图像文本基础模型，如对比语言图像预训练（CLIP）模型[28]，最近在计算机视觉和自然语言处理领域引起了广泛关注。这些模型在广泛的下游任务上表现出出色的零样本性能和鲁棒性，例如图像文本检索和分类（朱等，2023 年）。然而，CLIP 类模型的巨大计算和存储成本阻碍了它们的进一步普及。例如，MobileCLIP [33]在 256xA100 GPU 上训练，全局批量为 65,536，相应的数据集 DataCompDR-1B 需要 140 TB 的本地存储空间。此外，巨大的参数大小（例如，CLIP-B/16 模型[28]包含 86.2M 个图像编码器参数和 63.4M 个文本编码器参数）导致了推理延迟的增加，这为部署在计算资源有限的设备上带来了挑战。这些缺陷为没有足够计算资源和数据集的一般用户参与大规模模型的训练和部署设置了障碍。
在实际应用中，消费级计算机的 GPU 内存通常不超过 24GB（例如，Nvidia RTX 3090），存储容量可能小于 1TB。在这样资源限制的背景下训练 CLIP 类似的模型，需要解决两个主要问题。首先，必须尽可能减少需要训练的参数数量，同时保留尽可能多的现有模型知识。其次，小型数据集需要适当扩充，并需要开发更有效的方法，充分利用图像文本对在有限样本内的内部关联。
在这篇论文中，作者研究了如何使用仅有一个 RTX3090 显卡和 1TB 存储的轻量级 CLIP 模型训练方法，从而普及在消费级计算机上的 CLIP 类模型研究。为此，作者首先提出将传统的 Transformer 块简化为 SAS-P 块，并采用权重共享策略。然后，通过继承现有模型的权重并提取知识，可以进一步减少训练所需的参数数量。在数据集方面，作者选择广泛使用的 CC12M [1]作为基础。该数据集不仅规模较小，而且标签质量低，这都为模型训练过程的收敛带来了困难。为了解决这个问题，作者对 CC12M 中的每个图像样本添加多个文本标签，创建了新的 CC12M-SYN。此外，为了从这样的小数据集中提取有价值的信息，作者引入了 Pair Matching (PM)损失，以帮助模型捕捉正负图像-文本对之间的区分。这些方法在作者的广泛实验中显著提高了模型训练的收敛速度。最后，通过在 38 个数据集上的性能比较（如图 1 所示），作者的提出的 SiCLIP 框架实现了新的数据量-参数-准确率权衡的最先进水平。
作者的贡献：本工作的贡献可以概括如下：
作者提出了一个系统性的框架，用于在消费级计算机上训练轻量级 CLIP 模型，包括数据集构建和相应的训练过程，称为 SiCLIP。在 SiCLIP 中，计算和存储成本得到了降低，同时保持了与其他大规模模型竞争的性能。
作者通过在 SAS-P 块之间共享权重来简化 CLIP 模型的结构，并将权重继承与多阶段知识蒸馏（WIKD）相结合，从而在训练和部署过程中降低了内存需求。
一种名为 PM 损失的新损失函数被设计出来，该函数在训练过程中预测图像-文本对是否匹配。结合作者增强的 CC12M-SYN 数据集，PM 损失可以利用正负图像-文本对之间的区分。实验结果显示，新的数据集和 PM 损失都可以显著提高训练效率，同时略微增加数据集大小。
2 Related Work
Efficient Training for CLIP
自从 CLIP 作为一种大规模图像文本基础模型，在各种下游任务上展现出惊人的零样本性能以来，旨在提高其训练效率和模型大小的相关研究已经有很多。例如，包括细粒度图像文本对齐 [42]，数据增强 [23, 20, 33]，单模态自监督 [23, 20]，以及图像文本标签空间中的对比学习 [39]。此外，Zhai 等人 [45] 提出了一种对偶的 Sigmoid 损失作为对比损失的简单替代方案，在小型批量训练时证明了其有效性。然而，由于需要计算所有正负图像文本对之间的匹配对数，其可能导致二次计算复杂度。Li 等人 [19] 使用了细粒度图像文本匹配（ITM）损失作为对比损失的补充，但 ITM 需要一个基于多层 Transformer 的编码器来编码多模态细粒度特征，这并不适合轻量级模型。
基于权重继承（WI）和知识蒸馏（KD）[13]的方法也被采用以实现高效的训练。TinyCLIP[36]通过跨模态亲和力模拟和 WI 训练紧凑的 CLIP 模型。杨等人[38]探讨了不同的 KD 方法对 CLIP 训练的有效性。
高质量的数据集对于有效的训练也非常重要。方等人在[6]中利用过滤方法去除噪声样本，Gadre 等人[7]也提出了同样的方法。然而，剩余的描述性仍然不够。最近的研究[40; 18]表明，从预训练的描述生成模型中合成的合成描述可以提高数据集质量。
Simplifying the Transformer Architecture
近年来，随着 Transformer 在各种任务中的显著成功，许多努力都致力于简化 Transformer 架构以提高其训练和推理效率。余等人[44]证明了 Transformer 块的一般结构对其性能更为关键，因此可以消除基于注意力的 Token 混合器，这些混合器通常由于多头自注意力（MHSA）在长序列表示上的平方复杂度而变得 prohibitively 昂贵。此外，CNN 和 Transformer 的前期研究都表明，浅层主要关注局部模式，而更深层则倾向于捕捉高级语义或全局关系[14; 37; 5]，因此，在早期阶段通常不需要通过 MHSA 建模全局关系。基于这些事实，刘等人[22]提出了一种分层的 Transformer ，并采用移位窗口来限制自注意力计算限于非重叠的局部窗口，同时允许跨窗口连接，从而带来更高的效率。在另一项工作线中，潘等人[26]和郭等人[9]将卷积层引入到 Transformer 的前几层。在这些工作的基础上，瓦苏等人[32]提出了 RepMixer 作为 Token 混合器，它使用结构重参化通过在网络中删除 Short-Cut 来降低内存访问成本。
作为一种简单但有效轻量级方法，权重共享策略已在许多基于 Transformer 的模型中得到应用。Dehghani 等人[3]首次提出了利用 Transformer 层进行自然语言处理任务的不同动机：他们将重复的网络层视为 Transformer 引入循环归纳偏置的一种互补方式，并观察到他们的方法在多个任务上优于原始 Transformer 。Jaegle 等人[17]在多模态预训练中采用了交叉注意力层权重共享。Hernandez 等人[12]探讨了在不同粒度 Level 上共享符合器[8]的不同部分，受模型大小硬约束。最近，He 等人[10]通过信号传播理论研究了标准 Pre-LN Transformer 块[34]，并提出了一个配备形状注意力[25]的简化并行结构 Transformer 块作为 Token 混合器，称为简化注意力子块并行（SAS-P），在降低模型参数数量的同时，不损失语言下游任务的表现。作者的工作是首次尝试将 SAS-P 扩展到多模态领域，并通过在相邻块之间共享 Token 混合器的权重进一步简化。
3 Methods
在本节中，作者首先介绍了作者的简化模型结构，该结构通过共享 SAS-P 块的权重来实现。然后，作者介绍了一种名为 WIKD 的高效训练方法。接下来，作者介绍了一种新的损失函数，称为配对匹配（PM）损失，以进一步提高训练性能。最后，作者还通过在训练模型所使用的 CC12M 数据集中添加合成字幕，以提高数据多样性和数据质量，同时最小化附加存储空间。新的数据集称为 CC12M-SYN。图 2 显示了作者的方法的整体框架。
Simplifying Model Structure by Sharing Weights Among SAS-P Blocks
作者基于最新的 MobileCLIP-S0 模型[33]构建了作者的架构，并通过多种方式对其进行了增强。MobileCLIP-S0 框架具有图像编码器（image encoder）和文本编码器（text encoder）的混合结构，其中包含基于卷积（convolution-based）和 MHSA（Mobile Hyperspectral-to-Spatial Augmentation）的块的协同排列。然而，对于每个基于 MHSA 的块，MobileCLIP-S0 仅采用标准的 Pre-LN 块，其中 MHSA 作为混合器（token mixer）[34]，如图 3（左）所示。
作者首先从每个预连接块内的 Shortcut 参数开始减少。这些连接在内存访问和推理速度上造成了 Bottleneck ，因此，使用轻量级的基于 MHSA 的块设计变得至关重要。此外，已经证明，注意力模块的馈送层可以无缝集成到 Transformer 层中，而不会降低 Transformer 层的性能[31; 24]。
如图 3 右侧说明了 SAS-P（He 和 Hofmann，2024 年），这是一种简化的并行 Transformer 模块，在消除跳接连接的同时，还消除了值和投影参数。它使用形状注意力[25]作为其 Token 混合器，以防止在消除跳接连接后信号降解，使注意力矩阵更具标识性，从而保持良好的信号传播。形状注意力的注意力矩阵由：
SAS-P 在多个语言任务上表现出令人印象深刻的性能，同时比 Pre-LN 更快地推理速度，并使用较少的参数。 为了进一步简化模型结构，作者评估了相邻基于 MHSA 的块之间的 Jensen-Shannon（JS）离差（见图 4）。低 JS 离差意味着在这些矩阵之间进行权重共享不会降低性能。因此，在训练期间使用 KD 之前，作者的"学生"模型用 SAS-P 块替换了所有的预 LN 块，并在这些块之间进行权重共享。结果，作者模型的图像编码器与 MobileCLIP-S0 相比大约有 14%更少的参数，而在 OpenAI-B/16[28]中只有 11%。
为了从小型数据集中受益，一个广泛使用的方法是使用与任务相关的预训练 Backbone 网络，并添加一些任务特定的层[15]。受到使用 Backbone 网络思想的启发，作者采用 WI[36]来训练 CLIP 在小型数据集上。在实际应用中，由于作者修改了 MobileCLIP-S0 结构的基于 MHSA 的模块，而保持基于 RepMixer 的模块不变（这些模块已经高效），作者可以直接从已经在大型数据集上预训练得很好的 MobileCLIP-S0 继承这些模块的权重。在这种情况下，继承的模块可以被视为“ Backbone ”。
然后，作者冻结这些继承的层，只在新添加的 SAS-P 块上在一个非常小的数据集上进行训练。应用上述方法可以减少梯度存储，从而允许作者使用更大的批量大小来保持对比学习的性能。 此外，作者认为作者的模型是 MobileCLIP-S0 的移动模型，并在训练过程中进行多阶段知识蒸馏，从而进一步提高性能。具体而言，作者在单模态特征空间（阶段 1）、对比关系空间（阶段 2）和交互对比关系空间（阶段 3）上应用知识蒸馏。对于一个图像文本对齐的批处理，学生模型首先通过优化特征蒸馏损失（LFD）模拟教师图像和文本特征分布：
作者的最终蒸馏损失定义为：
λ1, λ2, λ3 均为超参数。
Pair Matching (PM) Loss
小型数据集训练的 CLIP 模型通常在零样本性能方面表现不佳[38]。作者认为这种现象的一个可能原因是，使用较少数据训练的模型在区分语义上相似的图像-文本对方面存在更大的困难。因此，作者提出构建一个辅助超平面来帮助模型判断图像-文本对是否匹配。请参阅图 5。
关于负样本对，对于每个图像，作者根据图像到文本相似性矩阵（如果其嵌入相似，则选择一个负样本文本的概率较高）选择一个负文本。同样地，这个过程也应用于每个文本。因此，负匹配 logits 被定义为：
因此，作者提出的 PM 损失函数可以表示为：
CC12M-SYN Dataset
用于训练 CLIP 模型的图像文本数据集大多来源于互联网，其中包含固有的噪声样本，这些样本的描述性不足。当使用小型数据集时，数据样本的多样性和质量变得尤为重要。增加合成字幕是一种既便宜又有效的方法，可以提高多样性和质量。作者采用广泛使用的数据集 CC12M [1]，并使用 coca [43]为该数据集中的每个图像生成多个合成字幕，得到 CC12M-SYN。图 6 展示了 CC12M-SYN 中的一些带有合成字幕的示例。在训练过程中，作者随机选择一组原始和合成字幕中的一个文本。因此，CC12M-SYN 中的一个样本包括一个图像和一个合成或原始的描述。
4 Experiments
Implementation Details
作者在第一万次训练迭代中采用了一个 Warm up 策略。作者使用了 AdamW 优化器，并将批量大小设置为 1536，权重衰减设置为 0.1。作者在 Nvidia RTX3090 上训练模型 32 个周期，学习率设置为 0.001。在消融研究中，作者将周期数设置为 9。作者采用 MobileCLIP-S0 作为 WIKD 的教师。对于超参数，作者设置 λ1=4000，λ2=λ3=1，λ4=0.1。其他设置遵循 CLIP-KD [38]。
作者在多个数据集上评估了零样本性能。具体来说，作者使用 ImageNet-1k [4]，ImageNet-V2 [29]，ImageNet-R [11]和 ImageNet-S [35]来评估零样本图像分类性能。对于零样本图像文本检索，作者使用 MSCOCO [21]和 Flickr30k [27]。默认为图像分类中的 Top-1 准确率（acc1）和图像文本检索中的 R@1。
4.1.1 Data Augmentation.
作者应用了随机缩放（RandomResizedCrop）和随机裁剪（RandAugment）进行图像增强。在 RandomResizedCrop 中，作者将缩放比例设置为（0.08, 1.0），以在原始图像上进行强大的增强。然后，作者在处理后的图像上应用 RandAugment，通过随机采用默认的 31 种增强方法（[2]）进一步增加图像的多样性。
Main Results
4.2.1 Zero-shot image-text retrieval.
表 1 报告了在 MSCOCO 和 Flickr30k 上的零样本图像文本检索性能。与相似规模的训练数据集（最多 2000 万样本）上的模型相比，作者的模型在 MSCOCO 上超过了所有其他工作。对于 Flickr30k，作者的模型在使用了较少的参数的同时，也实现了当前最先进模型 TinyCLIP 的性能。与在大规模数据集（400M-1B）上训练的模型相比，作者的模型在竞争性能上取得了良好的表现，并超越了许多现有工作。例如，与最先进的 MobileCLIP-S0 相比，作者的模型在使用了大约 3%的训练样本和 14%更少的图像编码器参数的情况下，T2I 性能仅低约 1%。此外，作者的模型在两个数据集的 I2T 和 T2I 指标上都超过了 DataComp-B/32、OpenAI-X（除 B/16）和 LAION-B/32。
4.2.2 Zero-shot image classification on ImageNet.
表 2 报告了零样本分类性能。与其他在相似规模的数据集上进行训练的模型相比，作者的模型在所有报告的数据集上都优于其他工作，这证明了作者的方法的有效性。对于大规模数据集，尽管与最新的 state-of-the-art DataComp-B/16 相比并非最佳，但作者与几项现有工作相比仍取得了一些具有竞争力的结果。
4.2.3 Inference Speed.
为了评估推理速度，作者在 CPU（Intel(R)-Xeon(R)-Silver-4314-CPU@2.40GHz）上进行了模拟实验，并比较了作者的模型与最先进的 MobileCLIP 系列[33]的平均推理速度。如表 3 所示，当输入序列包含 1000 张图像时，作者模型的处理速度达到 39.5 张/秒，略高于最先进的 MobileCLIP-S0（38.2 张/秒）。这表明采用 SAS-P 块的好处。
Ablation Studies
4.3.1 Training Efficiency of CC12M-SYN.
为了证明 CC12M-SYN 的训练效率提升，作者在 CC12M-SYN 和 CC12M 上分别训练作者的模型 20 个周期。作者报告了前 9 个周期的平均损失曲线，以及最后周期的 IN-1k 和 Flickr30k 上的零样本性能。图 7 报告了 CC12M 和 CC12M-SYN 的损失曲线，显示在 CC12M-SYN 上训练导致损失更快减少。表 4 表明，在 CC12M-SYN 上训练的模型在零样本分类和零样本图像文本检索方面具有更好的性能。这些结果表明合成标签对于数据多样性和质量提升具有益处。
4.3.2 Analysis of WIKD and PM loss.
作者通过比较有无 WIKD 和 PM 损失的训练（ Baseline ）以及仅 WI、仅 WIKD 和同时 WIKD 和 PM 损失的训练效果，来探索 WIKD 和 PM 损失的有效性。结果见表 5。它表明，仅 WI 的训练可以提高零样本分类（+13.0 和+6.1/+0.1 在分类准确率 acc1 上）和图像文本检索（R@1 的检索率+15.9/+15.2）。而当训练仅 WIKD 时，性能更高（+25.4 和+15.9/+15.2，分别）。当同时训练 WIKD 和 PM 损失时，模型达到最高性能。这些结果明显证明了 WIKD 和 PM 损失的有效性。
5 Conclusion
在这项工作中，作者提出了一系列技术，使得在消费级计算机上进行 CLIP 模型的训练和推理成为可能，同时实现了具有竞争力的性能。这对于将基础模型的卓越成果带到边缘设备上至关重要。
作者减少了模型结构，提高了推理速度。
此外，作者提出了 WIKD 和 PM 损失函数，这些损失函数对性能提升做出了贡献，并且可以用于简化其他领域中的模型。
最后，在增强的 CC12M-SYN 数据集上训练后，作者的模型在参数更少、训练数据集更小的情况下，与现有工作相比仍实现了具有竞争力的性能。
参考文献
[0]. Simplifying CLIP: Unleashing the Power of Large-Scale Models on Consumer-level Computers.
END
作者：小书童
来源：集智书童
推荐阅读
欢迎大家点赞留言，更多 Arm 技术文章动态请关注极术社区嵌入式 AI 专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
