DeepMind对CLIP改进的一篇工作--SigLIP,发表在2023CVPR。
简单看下研究动机:传统的对比学习方法如CLIP等依赖于 softmax 归一化,这需要一个全局视角来计算成对相似度,从而限制了批处理大小的扩展能力,并且在小批处理大小下表现不佳。因此本文提出了一个简单的成对 Sigmoid 损失函数用于语言-图像预训练(SigLIP)。
01、方法介绍
用于语言图像预训练的Softmax损失
CLIP等经典对比学习方法通常使用基于 softmax 的损失函数来对齐图像和文本的表示。具体来说,给定一个包含图像-文本对的小批量,对比学习的目标是使匹配对
的嵌入对齐,同时将不匹配对
的嵌入推远。通常通过训练一个图像模型f(⋅) 和一个文本模型g(⋅)实现。损失函数如下:
其中:
由于 softmax 损失的不对称性,规范化分别在图像和文本上独立进行。这种损失函数需要计算所有成对相似度的全局归一化因子,在大规模批处理中会导致内存和计算开销。
用于语言图像预训练的sigmoid损失
为了解决 softmax 损失的局限性,本文提出了一种基于 Sigmoid 的损失函数。Sigmoid 损失函数独立处理每个图像-文本对,避免了全局归一化的需要。具体来说,Sigmoid 损失函数将学习问题转化为标准的二分类问题,其中匹配对标记为正样本,所有其他对
的嵌入推远,标记为负样本。Sigmoid 损失函数的形式如下:
其中:是标签,
如果是匹配对,则
,否则
。
是 Sigmoid 函数,t和b是可学习的温度和偏置参数,初始化为t=log10 和b=−10。
初始化时,由于负样本数量远多于正样本,损失函数会被负样本主导,导致初始优化步骤较大。为了缓解这一问题,引入了一个额外的可学习偏置项b,以确保训练从接近先验的地方开始,避免需要巨大的矫正步骤。
高效“分块”实现
在对比训练中,通常利用数据并行性来加速计算。当数据被分割到D个设备上时,计算损失函数需要收集所有嵌入向量,这通常涉及到昂贵的全聚集操作,并且需要在内存中构建一个大规模的成对相似性矩阵,这在内存消耗上是非常密集的。
- Sigmoid损失的优势:sigmoid损失函数不需要计算全局归一化因子,这使得它特别适合于内存高效、快速且数值稳定的实现。这种方法可以减少数据并行处理中的挑战。
- 分块实现方法:作者提出了一种“分块”方法,避免了全聚集操作,并且减少了内存使用。具体来说,每个设备的批量大小表示为
,损失函数被重新表述为:
D为设备数,
是sigmoid损失函数
02、计算过程
计算过程简单来说,首先计算与正对对应的损失分量,然后对设备间的表示进行排列,使得每个设备从其邻近设备中获取负样本。每个设备独立地计算其本地批量b的损失,并将所有设备的损失简单地求和。
着重介绍下SigLIP是如何进行分布式训练的,假设全局的batch size为B,一共有个GPU,那么每个GPU上的batch size为
,可以将Sigmoid 损失函数表示公式的损失拆解为上述公式所示,在Fig 1展示了整个过程的示意图,在初始化阶段,以第一个GPU为例子,其所包含的样本为:
此时GPU 1可以完成一次上述公式中C的计算,然后,交换GPU 1和GPU 2的文本编码器特征向量,即:
此时GPU 2完成一次上述公式中的B计算,以此类推,直到GPU 1遍历完所有样本为止,其他GPU也是如此操作的,最终把所有卡上的损失聚集即可,也就是A计算。这个轮流交换不同GPU之间数据的操作,可以称之为排列。不难发现,整个过程的通信成本来自于排列,一共需要D−1次聚集操作即可完成一次排列。
而在CLIP中需要对图文的编码器特征都进行聚集,因此需要2次全聚集操作。如果全聚集采用ring方法的话,那么一个全聚集就是D−1次聚集操作。由此得出一个SigLIP和CLIP的性能复杂度对比:
这种方法显著减少了内存消耗,从减少到
,并且由于b通常是固定的,通过增加加速器的数量来扩展|B|。这种方法使得在相对较少的设备上训练超过一百万个批量大小的模型成为可能。
03、实验结果
batch_size影响
实验结果显示,在batch size小于32k的时候,采用sigmoid的SigLIP的性能都会优于采用softmax的CLIP。在batch size足够大的时候,CLIP能够追上,甚至超越SigLIP的表现,但是最佳性能仍然是SigLIP@32k情况下得到,从实际应用来看,采用SigLIP能用更少的资源更快的训练速度得到更好的性能。从SigLiT的实验来看,随着batch size的尺度放大性能将会在32k batch size的情况下达到饱和,同时能观察到SigLiT在不同batch size下持续优于LiT。继续提高batch size将不能带来更好的收益,甚至会有细微的性能下降。
训练时间影响
另外,作者在SigLiT的设置下,训练了更长时间(也即是见了更多数据量),结果显示,在超大batch size,如262k的情况下,对比较小batch size(如8k)提供更长的训练时间的确能观察到性能的较大提升。并且也能观察到在超大batch size下,采用sigmoid和采用softmax的区别很小,但是在较小batch size(如8k)的情况下,差距明显。
因此,在资源受限的情况下,采用SigLIP更划算,所需资源少而且性能更强。同时,这个实验也说明了,超大的batch size并不意味着训练得更快,反而还需要更长的训练时间。
多语言数据集表现
在多语言数据集上,性能同样在32k batch size上达到了饱和
有限资源场景
在仅有四块 TPU-v4 芯片的情况下使用 SigLiT 模型的训练效果如下,在四个芯片上训练一天的结果显示,模型在ImageNet上的零样本分类准确率达到了79.7%,这在资源有限的情况下非常有竞争力。使用ViT-g/14模型作为视觉塔和大型文本塔,可以在两天内以20k的批量大小进行107k步的训练,进一步将零样本分类准确率提高到84.5%。
另外还探讨了在资源有限的情况下,特别是只有少量TPUv4芯片时,如何有效地训练SigLIP模型。结果显示,使用16个TPUv4芯片,以16k的批量大小训练了2.4B个样本,仅靠微调并不能立即提高模型的性能。这与之前的研究一致,即微调图像模型会降低视觉表示的质量。然而,通过冻结预训练权重的权重衰减,可以获得更好的结果。
两个微调经验
微调图像模型会降低视觉表示的质量。然而,通过冻结预训练权重的weight-decay,可以获得更好的结果
随着批量大小的增加,训练过程变得越来越不稳定。实验发现降低Adam和AdaFactor中的β2参数值(从默认的0.999降低到0.95)可以稳定训练过程。
sigmoid损失函数中正负样本比例的影响
对于sigmoid来说,它的loss是以pair为粒度计算的,positive和negative非常不平衡。以batch_size=16k为例,positive和negative的比率约为1:16k
文中尝试了4种掩码策略来保留最难的负样本(损失最高)或最易的负样本(损失最低)
- 随机(Random):随机选择负对进行掩码。
- 难(Hard):保留最难的负对。
- 易(Easy):保留最易的负对。
- 难+匹配总对数(Hard + matching total pairs seen):在固定数量的步骤中掩码样本会减少训练期间看到的总对数。因此,在匹配对的设置中,作者增加了掩码比例的步数,以保持看到的对数不变。
不做matched pair的情况下,用3种mask方式均会造成精度下降。影响程度:easy>random>hard。
结果显示,随机移除负样本以重新平衡会恶化性能。保留最易的样本根本不起作用,而保留最难的负样本几乎可以保持质量。通过保留最难的负样本,模型的性能得到了维持,这表明在负样本中,较难的样本对学习更为重要。
另外,还观察了学习到的偏置项的最终值以及正负样本的平均logit值,发现当正负样本的imbalance减弱时,learnable bias和pair的logit都在上升,说明了预设的learnable bias起到了积极的作用。
数据噪声对模型鲁棒性的影响
通过以下五种方法“污染”训练噪声:
- Image:以概率p将图文对的图片用均匀噪声替换;
- Text:以概率p将图文对的文本token序列用随机采样的等长token序列替换;
- Batch alignment: 随机将batch中的p%的sample的图文pair进行shuffle;
- Image & text: 同时进行1.和2.
- Image, text & batch: 同时进行3和4
从结果可见,sigmoid loss在“污染”数据的性能更好。
04、总结
总的来说,SigLIP是一个很棒的工作,作者采用sigmoid损失去取代对比学习中的softmax函数,以更小的资源开销带来了更好的模型表现,目前也被很多多模态大模型所采用,作为视觉端的编码器。