Sigmoid Loss for Language Image Pre-Training

3 方法

我们介绍一种成对的sigmoid损失并且讨论他的有效应用。

给定一个小批次B = {(I1, T1), (I2, T2), ...}，对比学习的目的就是让匹配的图文对 (Ii, Ti)嵌入向量尽可能近，让不匹配图文对(Ii, Tj ≠ i)的嵌入向量相互远离。在实际操作中，我们假设对于所有图像 i,以及与另一个图形 j所关联的文本之间是毫无关联的。

3.1 图文预训练的Softmax损失

当使用softmax损失对目标公式化时，我们训练图像模型f(⋅)与文本模型 g(⋅)以将以下目标函数最小化。

$$ -\frac{1}{2|B|}\displaystyle\sum_{i=1}^{|B|}(\overbrace{log\frac{e^{tx_i\cdot y_i}}{\displaystyle\sum_{j=1}^{|B|}e^{tx_i\cdot y_j}}}^{image\to text\ softmax}+\overbrace{log\frac{e^{tx_i\cdot y_i}}{\displaystyle\sum_{j=1}^{|B|}e^{tx_j\cdot y_i}}}^{text\to image\ softmax}) $$

式中$x_i=\frac{f(I_i)}{||f(I_i)||_2}$（即图像Ii经过图像编码器f(⋅)后进行L2归一化）,$y_i=\frac{g(T_i)}{||g(T_i)||_2}$（即文本Ti经过文本编码器g(⋅)后进行L2归一化）。在这篇论文中，图像编码器采用Vision Transformer架构，文本编码器采用transformer架构。值得注意的是，由于softmax损失具有非对称性，我们需要在图像维度和文本维度分别进行两次归一化。标量温度参数t被参数化为exp(t′),其中，t′是一种可学习的全局变量。

3.2 图文预训练的Sigmoid损失

我们提出了一种不需要进行全局归一化计算的简单替代品。基于sigmoid的损失独立地处理每一个图文对，有效地将学习问题转化为在成对数据集上地标准二分类。在图文对组成的数据集中，匹配的图文对(Ii, Ti)为正类，其他图文对 (Ii, Tj ≠ i)为负类。损失函数定义如下：

$$ -\frac{1}{|B|}\displaystyle\sum_{i=1}^{|B|}\displaystyle\sum_{j=1}^{|B|}\underbrace{log\frac{1}{1+e^{z_{ij}(-tx_i\cdot y_j+b)}}}_{\mathcal{L}_{ij}} $$

zij是给定输入图文的标签，当图文成对时为1，否则为−1。在训练初期，大量的负值会导致严重的不平衡，导致我们需要大量初始优化步骤尝试纠正这些偏差。为了减缓这个问题，我们引入一种类似于温度t的可学习偏置项b。我们分别将t′tb初始化为log10与−10。这样能够让训练从更加合理的先验状态出发，不需要过度的初始纠正。

3.3 高效“分块”实现

对此学习采用数据并行方式。当数据分布在D个不同的设备上，计算损失时需要通过all-gathers操作将数据汇总，这具有高昂的代价。更重要的是需要构建内存占用为|B|×|B|的成对相似矩阵。

sigmoid损失适合于内存效率高，计算快速，并且具有稳定的数字化实现方式，这种实现方式有效改善了这两个问题。每台设备的批次大小记为$b=\frac{|B|}{D}$,loss值如下：

$$ -\frac{1}{|B|}\quad\underbrace{\displaystyle\sum^{D}_{d_i=1}}_{A:\forall device\space d_i}\quad\overbrace{\displaystyle\sum^{D}_{d_j=1}}^{B:swap\space negs\space across\space devices}\quad \overbrace{\underbrace{\displaystyle\sum^{b(d_i+1)}_{i=bd_i}}_{all\space local\space positives}\underbrace{\displaystyle\sum^{b(d_j+1)}_{j=bd_j}}_{negs\space from\space next\space device}\mathcal{L}_{ij}}^{C:per\space device\space loss} $$

这对于sigmoid损失很简单，因为每一个图文对都是独立的。

image.png

通俗地讲，我们首先计算正样本的损失值，然后计算每个正样本对应的b-1个负样本的损失值。然后我们在设备之间交换表示（嵌入向量），让每个设备从临近的设备中获取新的负样本（B步骤的下一轮）。这时在当前分块上计算相应的损失（这是C步骤）。这些在每台设备上独立完成，每台设备只计算本地的小批次b。将各个设备的损失直接求和（步骤A）。独立的collective permutes（B步骤）计算是快速的（实际上，D次collective permutes运算是快于D台设备间的all-gather运算的），任一时刻的内存消耗由|B|2缩小为b2（C步骤）。通常我们保持b不变，通过增加设备数量（加卡）来扩展整体batch size B。由于原始loss计算的内存需求随batch size二次增长，它很快就会成为扩展训练规模的瓶颈。这个分块的方法能够在相当有限的几台设备上训练参数超过一百万的模型。

4 结果

本节，我们对提出的SigLiT模型和SigLIP模型在多个批次大小上进行评估。我们讨论使用这两种方法在较少加速芯片的情况下能达到什么样的效果。我们也会简要分析批次大小对多语言图文预训练的影响。我们通过消融实验验证大批次稳定性调整的重要性，引入了可学习偏置项，研究了sigmoid中正负样本所占比例有何影响。最后，我们讨论了SigLIP面对数据噪音的鲁棒性。

为了验证我们的模型，我们报告了在ImageNet数据集上的零样本迁移学习结果，还有在XM3600数据机上跨36种语言的检索结果。我们的所有实验默认使用Scaling ViT-Adafactor优化器。

4.1 SigLiT：将批次大小最大化

我们使用 ViT-g 图像模型预先计算好的图像嵌入（embeddings），并在 LiT 图文数据集上，用相同的超参数从零开始训练一个 base 尺寸的文本编码器。

我们对于不同的批次大小进行研究，从512到1M，演示了批次大小对对比实验的影响，结果在图二钟展示。

图二：预训练批次大小的影响。左侧：SigLiT，

图二：预训练批次大小的影响。左侧：SigLiT，

当批次大小小于16K时，sigmoid损失表现得比softmax损失好很多。随着批次大小的增加，我们观察到softmax损失的性能迅速追赶并有望在批次大小足够大时稍稍超过sigmoid损失。总而言之，我们推荐在面对大批次时使用SigLIP，因为其简单，节约计算，并且直接高效操作内存。

目前普遍认为，大批次对于对比学习是有益的。然而大多数现有研究都止步于64k的批次大小。为了探究对比学习的极限，我们成功训练出了一个批次大小为100万的SigLiT模型。让我们惊讶的是，性能在32k时达到了饱和，继续增大batch size仅仅得到有限提升，最终在256k的batch size下达到最大值。我们最好的SigLiT模型达到了84.7%的零样本迁移准确率，使用B-size文本编译器。而LiT原始论文使用g-size文本编译器，仅仅得到了稍高一些的85.2%。

图三 不同训练时长对于SigLiT ImageNet下的零学习迁移准确率的影响。大batch size展现出了明显的性能提升，但是需要足够长的时间爬升。

图三 不同训练时长对于SigLiT ImageNet下的零学习迁移准确率的影响。大batch size展现出了明显的性能提升，但是需要足够长的时间爬升。

图像展示了如果训练时间足够长，262k的大batch size显著优于8k的小batch size。需要注意的是，在训练时间较短时，大batch size会减少更新步数，因此需要更多时间爬升。

4.2 SigLIP:Sigmoid损失适用于图文预训练

我们在WebLI dataset上预训练了SigLIP模型，仅使用英文图文对。我们将使用标准softmax损失在WebLI上训练的CLIP基线模型称为CLIP(WebLI)。我们使用了中等规模的模型：图像编码器为 ViT-B/16，文本编码器为 B 级 transformer。输入图像大小为224*224大小。文本使用一个 32k 词表的 sentencepiece 分词器进行分词（该分词器在英文 C4 数据集上训练而成），最多保留 16 个 token。图二中间图像展示了SigLIP的结果，在小于32k的批次大小下，SigLIP的性能优于CLIP(WebLI)基线。另一方面，空间高效的sigmoid损失能够支持更大的batch size。例如，四块TPU-v4芯片的情况下，使用SigLIP能够达到4096的batch size，而使用CLIP模型仅能达到2048。这两个优势结合在一起，使得sigmoid损失在图文预训练上展现了明显优势，在4.5中我们将继续讨论。

随着batch size增加，sigmoi与softmax之间的差距会减小。SigLIP在batch size为32k时性能最佳，而softmax损失需要98k才能展现出最佳性能，并且仍然无法超过基于sigmoi的模型。进一步扩大batch size如307k会导致两种模型的性能都下降。

4.3 mSigLIP：多语言预训练

我们继续通过保留WebLI数据集的100种语言来增大训练数据。处理读语言数据时，通常需要更大规模的跨国词汇。我们首先验证两个分词器的影响：一个32k tokens的小型多语言词表以及一个250k tokens的大型多语言词表。我们在总共 9 亿个样本上训练了 B 级规模的 ViT 与文本模型，发现使用更大词表时模型性能提升略超过 1%。

然而 ，当词表非常大时，token的嵌入矩阵会变得非常大。标准设置下，我们需要储存一个N × W的token嵌入查找表来训练多语言模型，其中N是上边提到的词表大小，W是文本模型嵌入向量的维度。为了节省内存，我们提出使用瓶颈型(bottlenecked)token嵌入。我们使用N × K嵌入矩阵以及一个K × W投影，其中K远小于W。

在实验中，观察到使用大型带瓶颈多语言词表可扩展性如小型多语言词表一样高效。具体而言，在Base模型架构中，W设置为768，K设置为96，我们在 ImageNet 的 zero-shot 分类任务中，模型性能仅下降约 0.5%，与使用完整的 250k 嵌入相比几乎无明显损失。

表二

表二

在不同的batch size下训练mSigLIP模型，训练了300亿个样本。表2与图2右侧展示了训练结果。我们原先认为更大的batch size能够提升多语言模型训练效果，因为模型在同一个mini-batch中会看到来在同一语言的更多难负样本。然而，在batch size大于32k后我们没有观察到明显的提升。对于多语言设置而言，32k已经足够。在XM3600跨模态检索任务中，发现batch size大于32k时会是性能整体下降，而对于ImageNet的零学习迁移任务则基本保持不变。mSigLIP使用Base级别的模型，在XM3600文本→图像检索任务中取得了最先进的性能。我们最好的成绩是34.9%，比之前使用四十亿参数ViT-e的标准LiT模型的28.5%高出了6%。

4.4 使用四块TPU-v4的SigLiT

对于很多参与者，对重要的问题是“怎样在算力受限的情况下训练模型？”。我们探讨在仅使用四块TPU-v4的情况下训练SigLiT模型，因为内存高效的sigmoid函数很适合这种算例受限的场景。

我们使用与4.1中相同的设置。使用公开的ViT-AugReg-B/8模型作为冻结视觉主干，并且预计算嵌入向量来加速训练。文本编码器使用大型Transformer，但是深度只有12层而不是常规的24层。使用LION优化器，设置解耦式权重衰减为1 × 10−7,学习率在6500步时达到峰值1 × 10−4,以余弦衰减减小到0。我们将batch size为32k的模型训练65000步，仅仅花费了不到一天的时间。表一展示了模型在四块芯片上训练一天的结果，ImageNet分类零学习准确率达到了79.7%。将视觉模型换为 ViT-g/14 [58] 并搭配 Large 文本编码器，我们可以在 4 块芯片上以 batch size = 20k 进行 107,000 步训练，用时不足两天。这将 zero-shot ImageNet 分类准确率进一步提升到 84.5%。

4.5 少量TPU-v4条件下的SigLIP

从零开始训练一个CLIP模型是很耗费资源的，而使用SigLIP可以在更少的芯片上容纳更大的batch size。在本节，我们研究如何借助预训练权重来训练SigLIP模型。使用预训练权重初始化的图像模型来加速预训练。我们使用公开的ViT-AugReg-B/16模型初始化视觉主干并在与SigLIP同样的WebLI英文数据上进行微调。在所有实验中，图象主干使用0.1学习率以适应微调。

图四与图五

图四与图五

图四展示了解锁微调模型与随机初始化训练模型的对比结果。我们使用16块TPU-v4芯片，batch size为16k训练24亿样本。我们发现，微调方案在默认配置下表现不好，这与先前的发现一致，微调图像模型会削弱视觉表示能力。这一点在ImageNet 10-shot线性分类中得到验证，如图四所示微调方案的表现仅仅略好于从头开始训练的基线。