方法概述
模型架构
下图 1 为 R2D2 预训练框架的架构概览,它包含一个文本编码器、一个图像编码器和两个交叉编码器。其中,文本编码器和图像编码器分别将文本和图像转换为隐藏状态的序列。然后,文本和图像隐藏状态通过交叉注意力在两个交叉编码器中交互。
研究者使用 RoBERTa 作为文本编码器。给定文本序列,他们首先使用 RoBERTa-wwm-ext 的 tokenizer 进行 tokenize 处理。在这里,特殊的[CLS] token 被附加到 tokenized 文本的 head,同时[SEP] token 被附加到 tail。然后,tokenized 文本被输入到文本编码器中。
研究者使用 ViT 作为图像编码器。他们首先将输入图像缩放到标准尺寸,并将该图像分成多个 patch。然后每个 patch 进行线性映射并添加位置嵌入。此外,一个可学习的[CLS] token 与 patch 向量串联。最后,序列向量被输入到标准 Transformer 模型以获得图像隐藏状态向量。
研究者将图像和文本隐藏向量融合并输入到交叉编码器。具体来说,他们使用线性映射层来更改每个文本特征和图像特征的维度以使它们保持一致。多层 Transformer 借助交叉注意力融合了两种模态的特征信息,并产生了最终的跨模态输出。
预训练方法
为了探索图文对之间的匹配关系,研究者设计了一种预排序 + 排序的机制,即全局对比预排序(GCPR)和细粒度排序(FGR)。他们还采用掩码语言建模(MLM)有效地学习跨模态模型的表示。
首先是全局对比预排序。传统的对比学习旨在对齐多模态数据的表示(如成对图文),它最大化了正对的相似度分数并最小化了负对的分数。研究者则使用全局对比学习来完成预排序任务,他们在 k 个 GPU 上执行完整的反向传播。对于每个图像 I_i 和对应的文本 T_i,图文和文图的 softmax 归一化相似度分数可以如下定义:
全局对比预排序损失通过交叉熵损失 L_c(·)来计算,如下公式 (2) 所示:
然后是细粒度排序。如上所述,研究者使用全局对比预排序来获得图像和文本的单个表示。基于这些表示,进一步借助细粒度排序损失来执行细粒度排序任务。这是一个二元分类任务,旨在预测图文是否匹配。
研究者将 h_I_[CLS] 和 h_T_[CLS]作为两个交叉编码器的输出表示。给定一个图像表示 h_I_[CLS]和一个文本表示 h_T_[CLS],研究者将它们输入到一个全连接层 g(·)以得到各自的预测概率。令 y 表示为二元分类的 ground-truth 标签,研究者对细粒度排序损失进行如下的计算。
研究者将掩码语言建模损失应用于文图交叉编码器,以提高在 token 级别建模文本和图像之间关系的能力。15% 的文本 token 在输入中被遮盖,所有这些 token 被替换成了[MASK] token。
在研究者的模型中,掩码语言建模任务利用掩码文本和对应的图像一起进行去噪,从而增强了文本和图像之间的交互。由于细粒度排序严重依赖这种交互能力,因此研究者提出了增强训练(ET),它将掩码语言建模任务集成到了正向图文对的细粒度排序前向操作中。
双向蒸馏
大部分图文预训练数据由半自动化程序所收集,从而导致其数据带有噪声。不准确的标签可能会误导模型训练。为此,研究者提出目标导向型蒸馏(TgD),一种带软目标的基于教师 - 学生的蒸馏。为了进一步提高预训练模型的泛化性能,研究者引入了特征导向型蒸馏(FgD)。为方便起见,他们将这两种蒸馏的组合称为双向蒸馏(TwD)。
目标导向型蒸馏:为了降低从噪声标签中学习的风险,研究者建议采用动量更新编码器产生的软目标(soft target)。这里,动量更新编码器作为蒸馏的教师模型,由指数移动平均权重得到。
研究者通过系数 α 将相似度得分 s(·,·) 与 one-hot 标签 y(·,·) 相结合,以生成最终的 soft 标签。将和表示为最终 soft 标签。以为例,可定义为:
考虑到队列中特征的有效性随着时间步长的增加而降低,研究者还维护了一个加权队列 w 来标记相应位置特征的可靠性。具体来说,除了新传入的项之外,该研究每次迭代都将队列中的每个元素衰减 0.99 倍。因此,研究者将 替换为等式 2 中的加权交叉熵损失。使用目标导向型蒸馏, 定义为:
特征导向型蒸馏:与 TgD 类似,研究者采用教师 - 学生范式进行特征导向型蒸馏。以文本编码器为例,学生模型是文本编码器,而教师模型是通过动量更新的编码器。
为了进一步提高模型性能,研究者对输入采用掩码策略。在具体实现中,将完整的输入提供给教师模型,将经过遮盖的输入提供给学生。依靠动量机制,目标是让学生的特征更接近教师的特征。形式上,教师和学生模型的预测分布分别定义如下:
研究者利用交叉熵损失进行特征导向型蒸馏,损失 L_FgD 定义为:
最后用整体预训练目标进行模型训练:
实验结果
从下表 2 中可以看出,研究者提出的模型在绝大多数任务中都超越了此前的 SOTA,即使仅使用 2.3M 样本(约为 Wukong 数据大小的 2.3%)进行训练时也是如此。对 23M 样本进行预训练时,结果更好。在模型层面,R2D2ViT-L 也在所有数据集中都优于 R2D2ViT-B,表明随着预训练模型变大,算法效果会变好。
研究者还在所提出的下游数据集上进行实验,这些实验结果成为所提出数据集的基线。特别地,在 Flickr30k-CNA 上进行实验时,研究者使用 Flickr30k-CNA 的训练集精调预训练模型,并在 Flickr30k-CN 的测试集上测试以进行公平比较。从表 2 可以看出,在 Flickr30k-CNA 上微调的 R2D2 优于在 Flickr30k-CN 上微调的 R2D2,因为人工翻译的 Flickr30k-CNA 的质量远高于机器翻译的 Flickr30k-CN。
与图文检索不同,有关中文的图文匹配 (ITM) 任务的数据集很少。研究者提出了针对中文 ITM 任务的图像标题匹配数据集 (ICM) 和图像查询匹配数据集(IQM),并给出了相应的结果。
研究者使用了 AUC 作为评价指标。从下表 1 中可以看出,R2D2ViT-L 比 R2D2ViT-B 取得了更好的结果。此外,R2D2ViT-L (23M) 在 ICM 和 IQM 上的表现分别比 R2D2ViT-L (2.3M) 高出约 4.97% 、5.68%。这意味着更多的中文高质量数据能够提高 R2D2 的泛化能力。
为了进一步提高性能,研究者用从 50 亿样本中提取的 2.5 亿个图文对进行预训练。从表 2 可以看出,以最综合的评估指标 R@M 来衡量,该模型在 Flickr30k-CN、COCO-CN、AIC-ICC、MUGE、Flickr30k-CNA、ICR、IQR 等数据集的所有结果中都超过了 23M 数据训练的模型,这意味着增加数据量可以增强预训练模型的能力。同时,这些结果也显著超越了有公开结果批露的 WenLan2.0 以及 WuKong 的结果,成为最新的 SOTA。下表 1 中数据则展示了在图文匹配任务 ICM、IQM 中,更大量的数据训练的预训练模型会取得更好的结果。
为了表明各个机制的作用,研究者在 Zero-Corpus 的子集上(230 万图文预训练数据)进行消融实验。为方便起见,在消融实验中研究者将 R2D2ViT-L 定义为 R2D2。
细粒度排序 (FGR) 的效果。首先,研究者使用全局对比预排序(GCPR)和双向蒸馏(TwD)训练模型,并定义为 PRD2。PRD2 的设置类似于 CLIP。从下表 3 的第一行和第二行可以看出,R2D2 在图文检索任务上的表现显著优于 PRD2,可以推测出该结果显著好于 CLIP,这说明了提出的全局对比预排序 + 细粒度排序框架的有效性。
增强训练 (ET) 的效果。研究者对比了去掉增强训练后的结果。从下表 3 的第三行可以看出,R2D2 (带 ET) 在图文检索任务上使 recall@1 提高了 0.95%,AUC 从 80.27% 提高到 80.51%。R2D2 的另一个优点是比 R2D2 w/o ET 使用更少的计算资源。R2D2 需要 154.0 GFLOPs 计算量,能够以每秒 1.4 次迭代的速度运行,而没有增强训练的 R2D2 则需要 168.8 GFLOPs 计算量,每秒只能运行 1.1 次迭代。上述结果证明了增强训练的有效性。
双向蒸馏的效果。研究者提出的双向蒸馏包含了目标导向型蒸馏和特征导向型蒸馏。当移除双向蒸馏(TwD)时,R@M 从 74.06% 降到 73.10%,AUC 从 80.51% 降到 80.31%。当移除特征导向型蒸馏(FgD)时,R@M 从 74.06% 降到 73.29%,性能下降明显,说明在训练中进行特征对齐很重要。同时,移除目标导向型蒸馏(TgD)也会导致模型的性能下降。上述结果说明双向蒸馏是一种提升预训练模型泛化性的有效方法。
零样本任务。为了证明模型的泛化性能,研究者还进行了零样本迁移实验。从下表 4 中可以看出,与目前 SOTA 性能的 WukongViT-L 相比,R2D2ViT-L(23M)仅使用了不到 1/4 的数据,却在 Flickr30k-CN、COCO-CN 上都取得了更好的性能。当引入 2.5 亿级别的预训练数据时,R2D2 的精度获得进一步提升,相对 WukongViT-L,在 Flickr30k-CN 数据集上,R@M 提升到 85.6%(提升了 4.7%),在 COCO-CN 数据集上,R@M 提升到 80.5%(提升了 5.4%),在 MUGE 数据集上,R@M 提升到 69.5%(提升了 6.3%)。
基于实体的图像注意力可视化。在本实验中,研究者尝试在 COCO-CN 上可视化图像的注意力。具体来说,他们首先从中文文本中提取一个实体,并计算图像与实体对的注意力得分。下图 2 展示了四个不同实体在图像上的可视化解释。这表明 R2D2 很好地学会了将文本与图像中的正确内容对齐。
- 数据集地址:https://zero.so.com
- 代码和模型地址:https://github.com/yuxie11/R2D2