同济提出简化 Transformer结构：在RTX3090上实现CLIP的轻量级训练 ！

对比语言图像预训练（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 输入， 和 分别是 Query 和键矩阵， 是模型维数， 是单位矩阵， 是可学习的参数。 是中心矩阵，其中每个元素都设置为 是输入 Token的数量。在初始化时， 设置为 分别设置为 1 ，导致 和 ，这对于良好的信号传播是有效的。

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）上应用知识蒸馏。对于一个图像文本对齐的批处理，学生模型首先通过优化特征蒸馏损失（）模拟教师图像和文本特征分布：

其中 和 分别表示教师模型和学生模型的（图像，文本）特征，其中 为批量大小。然后，它计算对比关系蒸馏损失 和交互对比关系损失 来模拟对比关系空间和交互对比关系空间的图像－文本相似矩阵分布，定义如下：

作者的最终蒸馏损失定义为：

其中是一个可学习的温度参数，Sim 表示图像特征和文本特征之间的相似矩阵。

λ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.