这篇文章的主要内容可以总结如下:
-
研究背景与问题:
-
CLIP(对比语言-图像预训练)模型在零样本任务和下游任务迁移中表现出色,但训练这类大规模模型需要大量计算和存储资源,限制了其在消费级计算机上的应用。
-
本文旨在解决如何在资源有限的设备(如单块RTX3090 GPU和1TB存储)上训练轻量级CLIP模型,并保持竞争力。
-
-
主要贡献:
-
简化模型结构:通过共享SAS-P(简化注意力子块并行)块的权重,减少了模型参数,提升了推理速度。
-
权重继承与多阶段知识蒸馏(WIKD):通过继承现有模型的权重并结合多阶段知识蒸馏,进一步减少训练所需的参数,提升模型性能。
-
配对匹配(PM)损失:设计了一种新的损失函数,帮助模型更好地区分正负图像-文本对,提升训练效率。
-
数据集增强:通过为CC12M数据集生成合成标题,创建了CC12M-SYN数据集,增加了数据多样性和质量。
-
-
实验结果:
-
在多个数据集上的实验表明,本文提出的SiCLIP框架在数据规模、参数数量和准确性之间实现了新的最优权衡。
-
与现有模型相比,SiCLIP在零样本图像分类和图像-文本检索任务上表现优异,且推理速度更快。
-
-
结论:
-
本文提出的方法使得在消费级计算机上训练和部署CLIP模型成为可能,显著降低了计算和存储需求,同时保持了与大规模模型相媲美的性能。
-
核心创新点:简化模型结构、权重继承与知识蒸馏、新的PM损失函数、数据集增强。
这里是自己的论文阅读记录,感兴趣的话可以参考一下,如果需要阅读原文的话可以看这里,如下所示:

摘要
对比语言-图像预训练(CLIP)因其卓越的零样本性能和出色的下游任务迁移能力而引起了广泛关注。然而,训练这种大规模模型通常需要大量的计算和存储资源,这对使用消费级计算机的普通用户构成了障碍。基于这一观察,本文研究了如何在仅使用一块Nvidia RTX3090 GPU和1TB存储空间的情况下实现具有竞争力的性能。一方面,我们简化了Transformer块结构,并结合权重继承与多阶段知识蒸馏(WIKD),从而减少了参数数量并提高了训练和部署时的推理速度。另一方面,面对小数据集带来的收敛挑战,我们为每个样本生成合成标题作为数据增强,并设计了一种新的配对匹配(PM)损失函数,以充分利用正负图像-文本对之间的区分度。大量实验表明,我们的模型在数据规模-参数-准确性之间实现了新的最先进权衡,这可以进一步推动CLIP模型在相关研究社区的普及。

1 引言
预训练的大规模图像-文本基础模型,尤其是对比语言-图像预训练(CLIP)模型,最近在计算机视觉和自然语言处理领域引起了广泛关注。这些模型在广泛的零样本任务中表现出色,并且在图像-文本检索和分类等下游任务中表现出强大的鲁棒性。然而,训练CLIP类模型所需的大量计算和存储成本阻碍了其进一步普及。例如,MobileCLIP在256块A100 GPU上训练,全局批量大小为65,536,相应的数据集DataCompDR-1B需要140TB的本地存储空间。此外,巨大的参数量(例如,CLIP-B/16模型的图像编码器有86.2M参数,文本编码器有63.4M参数)导致推理延迟增加,这为在计算资源有限的设备上部署模型带来了挑战。这些缺点使得计算资源和数据集不足的普通用户难以参与此类大规模模型的训练和部署。
在实践中,消费级计算机的GPU内存通常不超过24GB(例如,Nvidia RTX3090),而存储容量可能小于1TB。在这种资源限制下训练CLIP类模型时,必须解决两个主要问题。首先,必须尽量减少需要训练的参数数量,同时尽可能保留现有模型的知识。其次,小规模数据集需要适当增强,并且需要开发更有效的方法来充分利用有限样本中图像-文本对的内部相关性。
本文研究了如何仅使用一块RTX3090 GPU和1TB存储空间来训练轻量级CLIP模型,从而在消费级计算机上普及CLIP类模型的研究。为此,我们首先提出将传统的Transformer块简化为SAS-P块,并结合权重共享策略。然后,通过继承权重并从现有模型中蒸馏知识,进一步减少训练所需的参数数量。在数据集方面,我们选择广泛使用的CC12M作为基础数据集。该数据集不仅规模小,而且标签质量低,这都给模型训练过程的收敛带来了困难。为了解决这个问题,我们为CC12M的每个图像样本添加了多个文本标签,创建了新的CC12M-SYN数据集。此外,为了从这种小数据集中提取有价值的信息,我们引入了配对匹配(PM)损失函数,以帮助模型捕捉正负图像-文本对之间的区分度。这些方法通过大量实验显著提高了模型训练的收敛速度。最后,通过在38个数据集上与几种流行方法的性能比较(图1),我们提出的SiCLIP框架实现了新的数据规模-参数-准确性权衡。
我们的贡献: 本工作的贡献可以总结如下:
-
我们提出了一个系统框架,用于在消费级计算机上训练轻量级CLIP模型,包括数据集构建和相应的训练过程,称为SiCLIP。在SiCLIP中,计算和存储成本都得到了降低,同时与其他大规模模型相比保持了竞争力的性能。
-
我们通过共享SAS-P块之间的权重简化了CLIP模型结构,并结合权重继承与多阶段知识蒸馏(WIKD),从而减少了训练和部署时的内存需求。
-
我们设计了一种新的损失函数,称为PM损失,用于预测图像-文本对是否匹配。结合我们增强的数据集CC12M-SYN,PM损失可以充分利用正负图像-文本对之间的区分度。实验结果表明,新的数据集和PM损失都可以显著提高训练效率,同时仅略微增加数据集大小。

图2:我们提出的SiCLIP框架概述。(a) 我们通过为每张图像添加多个合成标题来增强广泛使用的CC12M数据集。在训练过程中,对于每个图像-文本对,图像通过RandomResizedCrop和RandAugment进行增强,文本则从原始和合成标题集合中随机选择(RC)。(b) 训练过程的示意图。我们简化了模型结构,继承了教师的权重,并在训练过程中同时优化对比损失、PM损失和多阶段知识蒸馏(KD)损失。
2 相关工作
2.1 CLIP的高效训练
自从CLIP作为一种具有出色零样本性能的大规模图像-文本基础模型被引入以来,已有许多工作致力于提高其训练效率和模型大小。例如,细粒度图像-文本对齐、数据增强、单模态自监督学习以及图像-文本-标签空间中的对比学习。此外,Zhai等人提出了成对Sigmoid损失作为对比损失的简单替代方案,证明了在小批量训练时的有效性。然而,由于所有正负图像-文本对之间的匹配对数计算,可能会导致二次计算复杂度。Li等人使用细粒度图像-文本匹配(ITM)损失作为对比损失的补充,但ITM需要基于多层Transformer的编码器来编码多模态细粒度特征,这不适合轻量级模型。
权重继承(WI)和知识蒸馏(KD)方法也被用于高效训练。TinyCLIP通过跨模态亲和力模仿和WI训练紧凑的CLIP模型。Yang等人探索了不同KD方法在CLIP训练中的有效性。
2.2 简化Transformer架构
随着Transformer在各种任务中的显著成功,近年来有许多努力致力于简化Transformer架构以提高其训练和推理效率。Yu等人证明了Transformer块的通用结构对性能更为重要,使得可以消除基于注意力的令牌混合器,这些混合器通常由于多头自注意力(MHSA)在长序列表示上的二次复杂度而变得昂贵。此外,先前在CNN和Transformer中的研究表明,浅层主要关注局部模式,而深层倾向于捕捉高级语义或全局关系,因此在早期阶段通常不需要通过MHSA建模全局关系。基于这些事实,Liu等人提出了分层Transformer,并使用移位窗口将自注意力计算限制在非重叠的局部窗口内,同时允许跨窗口连接,从而提高了效率。在另一项工作中,Pan等人和Guo等人将卷积层引入早期Transformer层。Vasu等人提出了RepMixer作为令牌混合器,通过去除网络中的跳跃连接来降低内存访问成本。
作为一种简单但有效的轻量级方法,权重共享策略已被许多基于Transformer的模型采用。Dehghani等人首次提出了在自然语言处理任务中重用Transformer层的想法,他们认为重复的网络层是引入递归归纳偏差的补充方式,并且观察到他们的方法在多个任务上优于普通Transformer。Jaegle等人在多模态预训练中采用了交叉注意力层权重共享。Hernandez等人探索了在不同粒度级别共享Conformer的不同部分,以硬约束内存中的模型大小。最近,He等人通过信号传播理论研究了标准的Pre-LN Transformer块,并提出了一种简化的并行结构Transformer块,配备形状注意力作为令牌混合器,称为简化注意力子块并行(SAS-P),减少了参数数量并提高了模型的吞吐量,而不会在语言下游任务中损失性能。我们的工作是首次将SAS-P扩展到多模态领域,并通过在相邻块之间共享令牌混合器的权重进一步简化它。
3 方法
在本节中,我们首先介绍了我们简化的模型结构,该结构在SAS-P块之间共享权重。然后,我们介绍了一种称为WIKD的高效训练方法。接下来,我们介绍了一种新的损失函数,称为配对匹配(PM)损失,以进一步提高训练性能。最后,我们还通过添加合成标题来改进我们用于训练模型的CC12M数据集,以提高数据多样性和数据质量,同时仅略微增加存储空间。新的数据集称为CC12M-SYN。图2展示了我们方法的整体框架。
3.1 通过共享SAS-P块权重简化模型结构
我们的架构基于最先进的MobileCLIP-S0模型,我们在多个方面对其进行了增强。MobileCLIP-S0框架为图像编码器和文本编码器采用了混合结构,结合了基于卷积和基于MHSA的块的协同排列。然而,对于每个基于MHSA的块,MobileCLIP-S0简单地采用了标准的Pre-LN块,其中MHSA作为令牌混合器。
我们首先减少了每个Pre-LN块中跳跃连接的参数。这些连接的存在导致了内存访问和推理速度的瓶颈,因此轻量级的基于MHSA的块设计变得至关重要。此外,研究表明,前馈层可以集成到注意力模块中,而不会降低Transformer层的性能。
图3的右侧展示了SAS-P(He和Hofmann,2024),这是一种简化的并行Transformer块,消除了跳跃连接以及值和投影参数。它使用形状注意力作为其令牌混合器,以防止在去除跳跃连接后信号退化,使注意力矩阵更类似于身份矩阵,以保持良好的信号传播。形状注意力的注意力矩阵由以下公式给出:

为了进一步简化模型结构,我们评估了相邻MHSA块之间的Jensen-Shannon(JS)散度(见图4)。低JS散度意味着可以在这些矩阵之间使用权重共享而不会降低性能。因此,在训练期间应用KD之前,我们的“学生”模型将所有Pre-LN块替换为SAS-P块,并在这些块之间应用权重共享。结果,我们的模型的图像编码器比MobileCLIP-S0减少了约14%的参数,仅为OpenAI-B/16的11%。

3.2 权重继承与多阶段知识蒸馏(WIKD)
为了从小规模数据集中受益,广泛使用的范式是采用与任务相关的预训练骨干网络并添加一些任务特定的层。受使用骨干网络的想法启发,我们采用WI在小规模数据集上训练CLIP。在实践中,由于我们修改了MobileCLIP-S0结构的基于MHSA的块并保持基于RepMixer的块不变(这些块已经高效),我们可以直接从MobileCLIP-S0继承这些模块的权重,这些模块已经在大型数据集上进行了良好的预训练。在这种情况下,继承的模块可以被视为“骨干网络”。然后,我们冻结这些继承的层,仅在更小的数据集上训练新添加的SAS-P块。应用上述方法可以减少梯度存储,从而允许我们使用更大的批量大小来保持对比学习的性能。
此外,我们将我们的模型视为MobileCLIP-S0的学生模型,并在训练期间执行多阶段KD,从而进一步提高性能。具体来说,我们在单模态特征空间(阶段1)、对比关系空间(阶段2)和交互对比空间(阶段3)上应用KD。给定一批图像-文本对,学生模型首先通过优化特征蒸馏损失(
)来模仿教师的图像和文本特征分布,公式如下:

其中τ是可学习的温度参数,Sim表示图像和文本特征之间的相似度矩阵。我们的最终蒸馏损失定义为:


3.3 配对匹配(PM)损失
在小规模数据集上训练的CLIP模型通常具有不令人满意的零样本性能。我们认为这种现象的一个可能原因是,使用较少数据训练的模型更难区分语义上相似的图像-文本对。因此,我们提出构建一个辅助超平面来帮助模型确定图像-文本对是否匹配。图5展示了一个示意图。



3.4 CC12M-SYN数据集
用于训练CLIP模型的图像-文本数据集大多从互联网收集,包含固有的噪声样本,这些样本的描述性不够。在使用小规模数据集时,数据样本的多样性和质量变得更加重要。添加合成标题是一种廉价但有效的方法,可以提高多样性和质量。我们采用广泛使用的数据集CC12M,并使用coca为数据集中的每个图像生成多个合成标题,获得CC12M-SYN。图6展示了CC12M-SYN中带有合成标题的一些示例。在训练期间,我们从原始和合成标题的集合中随机选择一个文本。结果,CC12M-SYN中的一个样本由一张图像和一个合成或原始标题组成。
4 实验
4.1 实现细节
我们在前10000次训练迭代中采用预热策略。我们使用AdamW优化器,并将批量大小设置为1536,权重衰减设置为0.1。我们在Nvidia RTX3090上以学习率0.001训练我们的模型32个epoch。在消融研究中,epoch数设置为9。我们采用MobileCLIP-S0作为WIKD的教师模型。对于超参数,我们设置λ1=4000,λ2=λ3=1,λ4=0.1。其他设置遵循CLIP-KD。
我们在多个数据集上评估零样本性能。具体来说,我们使用ImageNet-1k、ImageNet-V2、ImageNet-R和ImageNet-S来评估零样本图像分类性能。对于零样本图像-文本检索,我们使用MSCOCO和Flickr30k。默认情况下,我们报告图像分类的top-1准确率(acc1)和图像-文本检索的R@1。
4.1.1 数据增强
我们应用RandomResizedCrop和RandAugment进行图像增强。我们在RandomResizedCrop中将比例设置为(0.08,1.0)以对原始图像进行强增强,然后在处理后的图像上应用RandAugment,通过随机采用默认的31种增强方法中的两种来进一步增强图像。

4.2 主要结果
4.2.1 零样本图像-文本检索
表1报告了在MSCOCO和Flickr30k上的零样本图像-文本检索性能。与在类似规模数据集(最多20M样本)上训练的模型相比,我们的模型在MSCOCO上优于所有其他工作。对于Flickr30k,我们的模型也达到了当前最先进模型TinyCLIP的性能,同时使用了更少的参数。与在大规模数据集(400M-1B)上训练的模型相比,我们的模型实现了具有竞争力的性能,并优于许多现有工作。例如,与最先进的MobileCLIP-S0相比,我们的模型在T2I性能上仅低约1%,同时使用了约3%的训练样本和14%更少的图像编码器参数。此外,我们的模型在I2T和T2I指标上均优于DataComp-B/32、OpenAI-X(除了B/16)和LAION-B/32。
4.2.2 ImageNet上的零样本图像分类
表2报告了零样本分类性能。与在类似规模数据集上训练的模型相比,我们的模型在所有报告的数据集上优于其他工作,证明了我们方法的有效性。对于大规模数据集,虽然不如最新的最先进DataComp-B/16,但我们仍然与几个现有工作相比取得了一些具有竞争力的结果。
4.2.3 推理速度
为了评估推理速度,我们在CPU(Intel(R)-Xeon(R)-Silver-4314-CPU@2.40GHz)上进行了模拟实验,并与最先进的MobileCLIP系列进行了平均推理速度的比较。如表3所示,给定1000张图像的输入序列,我们的模型的处理速度达到了39.5 images/sec,略优于最先进的MobileCLIP-S0(38.2 images/sec)。这证明了采用SAS-P块的好处。
4.3 消融研究
4.3.1 CC12M-SYN的训练效率

为了证明CC12M-SYN的训练效率改进,我们分别在CC12M-SYN和CC12M上训练我们的模型20个epoch。我们报告了前9个epoch的平均损失曲线以及最后一个epoch在IN-1k和Flickr30k上的零样本性能。图7报告了CC12M和CC12M-SYN的损失曲线,显示在CC12M-SYN上训练导致损失更快减少。表4显示,在CC12M-SYN上训练的模型在零样本分类和零样本图像-文本检索上具有更好的性能。这些结果表明了合成标签对数据多样性和质量改进的好处。
4.3.2 WIKD和PM损失的分析
我们通过比较在没有WIKD和PM损失的情况下训练(基线)、仅使用WI训练、使用WIKD训练以及同时使用WIKD和PM损失训练的性能,探索了WIKD和PM损失的有效性。结果如表5所示。结果表明,仅使用WI训练可以在零样本分类和图像-文本检索上提供改进(分类的acc1提高了+13.0,检索的R@1提高了+6.1/+0.1)。当使用WIKD训练时,性能会更高(分别提高了+25.4和+15.9/+15.2)。当同时使用WIKD和PM损失训练时,模型达到了最高的性能。这些结果明显支持了WIKD和PM损失的有效性。
5 结论
在这项工作中,我们提出了多种技术,使得在消费级计算机上训练和推理CLIP模型成为可能,同时实现了具有竞争力的性能。这对于将基础模型的令人印象深刻的成果带到边缘设备至关重要。我们减少了模型结构,提高了推理速度。此外,我们提出了WIKD和PM损失,这些方法有助于性能改进,并且可以用于简化其他领域中的模型。最后,在增强的CC12M-SYN数据集上训练,我们的模型与现有工作相比实现了具有竞争力的性能,尽管具有更少的参数并且在更小的数据集上训练。

总结
本文通过简化模型结构、引入权重继承与多阶段知识蒸馏、设计新的损失函数以及增强数据集,成功地在消费级计算机上实现了轻量级CLIP模型的训练,并取得了与大规模模型相媲美的性能。这一研究为在资源受限的设备上部署大规模模型提供了新的思路和方法。
1376






