ALBEF：基于动量蒸馏的视觉语言表示学习 《Align before Fuse：Vision and Language Representation Learning with Momentum Distillation》

论文地址：https://arxiv.org/pdf/2107.07651.pdf

相关博客：
【自然语言处理】【多模态】多模态综述：视觉语言预训练模型
【自然语言处理】【多模态】CLIP：从自然语言监督中学习可迁移视觉模型
【自然语言处理】【多模态】ViT-BERT：在非图像文本对数据上预训练统一基础模型
【自然语言处理】【多模态】BLIP：面向统一视觉语言理解和生成的自举语言图像预训练
【自然语言处理】【多模态】FLAVA：一个基础语言和视觉对齐模型
【自然语言处理】【多模态】SIMVLM：基于弱监督的简单视觉语言模型预训练
【自然语言处理】【多模态】UniT：基于统一Transformer的多模态多任务学习
【自然语言处理】【多模态】Product1M：基于跨模态预训练的弱监督实例级产品检索
【自然语言处理】【多模态】ALBEF：基于动量蒸馏的视觉语言表示学习
【自然语言处理】【多模态】VinVL：回顾视觉语言模型中的视觉表示
【自然语言处理】【多模态】OFA：通过简单的sequence-to-sequence学习框架统一架构、任务和模态
【自然语言处理】【多模态】Zero&R2D2：大规模中文跨模态基准和视觉语言框架

一、简介

​ 视觉语言预训练$\text{(Vision-and-Language Pre-training,VLP)}$的目标是从大规模image-text对中学习多模态表示，用于改善下游的视觉语言任务$(\text{Vision-and-Language,V+L})$。许多现有的$\text{VLP}$方法依赖于预训练的目标检测器来抽取基于图像特征的区域，并利用一个多模态编码去来将图像特征与单词特征进行融合。多模态编码器被训练来解决那些需要联合理解图像和文本的任务，例如：masked language modeling和image-text matching。

​ 虽然有效，但是这些$\text{VLP}$框架存在着几个关系的限制：(1) 图像特征和单词嵌入都处于自己的空间中，这使得多模态编码器学习建模他们的交互更具挑战性；(2) 目标检测器的标准和计算都很昂贵，因为其需要在预训练的时候人工标注bounding box，并且在推断时为高分辨率图像；(3) 广泛被使用的image-text数据集都是从网络上收集的并且存在大类噪音，现有像$\text{MLM}$这样的预训练目标可能会过拟合噪音文本，并降低模型的泛化性能。

​ 作者提出了$\text{ALBEF(ALign BEfore Fuse)}$，一个新的$\text{VLP}$框架来解决这些限制。首先会使用一个无需检测器的图像编码器和文本编码器来独立编码图像和文本。然后，多模态编码器通过跨模态注意力机制来融合图像特征和文本特征。作者引入了一个中间的image-text对比损失函数$(\text{ITC})$，应用在单模态编码器的表示上，其有三个目的：(1) 对齐图像特征和文本特征，使得多模态编码器更容易执行跨模态学习；(2) 改善单模态编码器来更好的理解图像和文本的语义；(3) 其能够学习一个共同的低维空间来嵌入图像和文本，通过对比难样本挖掘来发现更具信息量的样本。

​ 为了改善在噪音监督下的学习，作者提出了动量蒸馏$\text{MoD}$，一个简单的使模型能够利用较大的噪音数据集的方法。在训练过程中，通过对模型参数进行平均来维护一个动量版本的模型，并使用动量模型来生成伪标签作为额外的监督。使用$\text{MoD}$，模型不会应为生成不同于网络标注的合理输出而受到惩罚。$\text{MoD}$不仅能够改善预训练，也能够改善下游任务。

​ 作者从最大互信息的角度提供了$\text{ALBEF}$的理论分析。特别地，$\text{ITC}$和$\text{MLM}$最大化了image-text对不同视图的互信息下边界，这些视图是通过采用每个对的部分信息生成的。从这个角度出发，动量蒸馏能够被解释为生成语义相同的新视图。因此，$\text{ALBEF}$能够学习对语义表示不变的视觉语言表示。

​ 作者在各种下游$\text{V+L}$任务上证明了$\text{ALBEF}$的有效性，包含image-text检索、视觉问答、视觉推理、视觉蕴含以及弱监督的visual grounding。$\text{ALBEF}$比现有的state-of-the-art方法实现了显著的改进。在image-text检索中，其比那些在更大一个量级数据集上预训练的方法要好($\text{CLIP}$和$\text{ALIGN}$)。在$\text{VQA}$和$\text{NLVR}$上，相较于state-of-the-art方法$\text{VILLA}$，其实现了2.37%和3.84%的改进，并且具有更快的推理速度。此外，作者还使用$\text{Grad-CAM}$对$\text{ALBEF}$进行了定性和定量的分析。

二、$\text{ALBEF}$预训练

1. 模型架构

​ 如上图所示 ，$\text{ALBEF}$包含一个图像编码器、一个文本编码器 和一个多模态编码器。使用12层的$\text{ViT-B/16}$作为图像编码器，然后使用在$\text{ImageNet-1K}$上预训练得到的权重来初始化。一个输入图像$\text{I}$被编码为嵌入序列： { v c l s , v 1 , … , v N } \{\textbf{v}_{cls},\textbf{v}_1,\dots,\textbf{v}_N\} {vcls​,v1​,…,vN​}，其中$v_{cls}$是[CLS]的嵌入向量。使用6层的$\text{Transformer}$作为文本编码器和多模态编码器。文本编码器使用${\text{BERT}}_{base}$的前6层进行初始化，多模态编码器层使用${\text{BERT}}_{base}$的后6层进行初始化。文本编码器将输入文本$T$转换为嵌入向量序列 { w c l s , w 1 , … , w N } \{\textbf{w}_{cls},\textbf{w}_1,\dots,\textbf{w}_N\} {wcls​,w1​,…,wN​}，其会被输入至多模态编码器。在多模态编码器的每层会通过注意力机制将图像特征和文本特征进行融合。

2. 预训练目标

​ 使用三个目标函数预训练$\text{ALBEF}$：单模态编码器上的image-text对比学习$(\text{ITC})$，多模态编码器上的遮蔽语言模型$(\text{MLM})$和image-text匹配$(\text{ITM})$。此外，这里还使用在线难负样本挖掘对比来改善$\text{ITM}$。

2.1 $\text{Image-text}$对比学习$(\text{ITC})$

​ 该损失函数的目标是在融合之前更好的学习单模态表示。其会学习一个相似函数$s=g_v({\text{v}}_{cls}{)}^{\top }{g}_w({\text{w}}_{cls})$，使得并行的image-text对具有更高的相似分数。$g_v$和$g_w$是将[CLS]嵌入向量映射为规范化低维度表示的线性变换。受$\text{MoCo}$启发，维护两个队列来存储来自动量单模态编码器中最近的$M$个image-text表示。来自动量编码器的规范化特征表示为$g_v^{\prime }({\text{v}}_{cls}^{\prime })$和$g_w^{\prime }({\text{w}}_{cls}^{\prime })$。定义$s(I,T)=g_v({\text{v}}_{cls}{)}^{\top }{g}_w^{\prime }({\text{w}}_{cls}^{\prime })$且$s(T,I)=g_w({\text{w}}_{cls}{)}^{\top }{g}_v^{\prime }({\text{v}}_{cls}^{\prime })$。

​ 对于每个图像和文本，计算image-to-text和text-to-image的相似度为：
$$p_m^{i2t}(I)=\frac{exp(s(I,T_m)/\tau )}{{\sum }_{m=1}^{M}exp(s(I,T_m)\tau )},p_m^{t2i}=\frac{exp(s(T,I_m)/\tau )}{{\sum }_{m=1}^{M}exp(s(T,I_m)/\tau )}\text{\hspace{1em}(1)}$$
其中，$\tau$是可学习temperature参数。令${\text{y}}^{i2t}(I)$和${\text{y}}^{t2i}(T)$表示真实的one-hot相似度，其中负样本对具有概率0且正样本对的概率为1。image-text对比损失函数被定义为$\text{p}$和$\text{y}$的交叉熵$H$：
$${\mathcal{L}}_{itc}=\frac{1}{2}{\mathbb{E}}_{(I,T)\sim D}[H({\text{y}}^{i2t}(I),{\text{p}}^{i2t}(I))+H({\text{y}}^{t2i}(T),{\text{p}}^{t2i}(T))]\text{\hspace{1em}(2)}$$

2.2 遮蔽语言模型$(\text{MLM})$

​ $\text{MLM}$会利用图像和文本预测被遮蔽的单词。以15%的概率随机遮蔽输入的tokens，并使用特殊的[MASK] token进行替换。令$\hat{T}$表示为遮蔽的文本，并且${\text{p}}^{msk}(I,\hat{T})$表示模型对遮蔽的token的预测概率。$\text{MLM}$最小化交叉熵损失：
$${\mathcal{L}}_{mlm}={\mathbb{E}}_{(I,\hat{T})\sim D}H({\text{y}}^{msk},{\text{p}}^{msk}(I,\hat{T}))\text{\hspace{1em}(3)}$$
其中，${\text{y}}^{msk}$是一个one-hot词表分布，其中真实的token概率为1。

2.3 $\text{Image-Text}$匹配$(\text{ITM})$

​ $\text{ITM}$预测图像和文本对匹配或者不匹配。使用多模态编码器对于[CLS]输出的嵌入向量作为image-text对的联合表示，并通过在全链接层跟一个softmax来预测两种类别的概率$p^{itm}$。$\text{ITM}$损失函数为：
$${\mathcal{L}}_{itm}={\mathbb{E}}_{(I,T)\sim D}H({\text{y}}^{itm},{\text{p}}^{itm}(I,T))\text{\hspace{1em}(4)}$$
其中，${\text{y}}^{itm}$是一个二维的one-hot向量表示。

​ 此外，作者提出了一个针对$\text{ITM}$任务的难负样本采样策略。如果image-text对共享相似的语义但在细粒度细节上不同，那么就可以认为是难负样本。利用等式$(1)$中的对比相似性来寻找batch内部的难负样本。对于batch内的每个图像，从同一个batch中按照对比相似度分布采样一个负文本，其中文本与图像越相似则被采样的机会更高。类似地，为每个文本采样一个难负图像。

• $\text{ALBEF}$的全部预训练目标函数为：
  $$\mathcal{L}={\mathcal{L}}_{itc}+{\mathcal{L}}_{mlm}+{\mathcal{L}}_{itm}\text{\hspace{1em}(5)}$$

3. 动量蒸馏

​ 用于预训练image-text对的数据主要从网络上收集，并且包含了噪音。正样本对通常具有弱相关性：文本中可能包含与图像不相关的单词，或者图像中可能包含文本中未描述的实体。对于$\text{ITC}$学习，一个图像的负文本也可能匹配图像中的内容。对于$\text{MLM}$，存在着一些不同于标注的单词能够更好的描述图像。然而，$\text{ITC}$和$\text{MLM}$的one-hot标签惩罚所有的负预测，并且忽视了这些正确性。

​ 为了解决这个问题，作者提出了通过动量模型生成的伪目标进行学习。动量模型是由单模态编码器和多模态编码器的指数移动平均版本组成的、不断进化的教师模型。在训练过程中，训练基础模型使其预测与动量模型的预测相匹配。特别地，对于$\text{ITC}$，首先使用动量单模态编码器的特征来计算image-text相似度为$s^{\prime }(I,T)=g_v^{\prime }({\text{v}}_{cls}^{\prime }{)}^{\top }{g}_w^{\prime }({\text{w}}_{cls}^{\prime })$和$s^{\prime }(T,I)=g_w^{\prime }({\text{w}}_{cls}{)}^{\top }{g}_v^{\prime }({\text{v}}_{cls}^{\prime })$。然后，通过替换等式$(1)$中的$s$和$s^{\prime }$来计算伪标签${\text{q}}^{i2t}$和${\text{q}}^{t2i}$。${\text{ITC}}_{MoD}$损失函数定义为：
$${\mathcal{L}}_{itc}^{mod}=(1-\alpha ){\mathcal{L}}_{itc}+\frac{\alpha }{2}{\mathbb{E}}_{(I,T)\sim D}[\text{KL}({\text{q}}^{i2t}(I)\parallel {\text{p}}^{i2t}(I))+\text{KL}({\text{q}}^{t2i}(T)\parallel {\text{p}}^{t2i}(T))]\text{\hspace{1em}(6)}$$
类似地，对于$\text{MLM}$，令${\text{q}}^{msk}(I,\hat{T})$表示动量模型对于遮蔽token的预测概率，${\text{MLM}}_{MoD}$损失函数为：
$${\mathcal{L}}_{mlm}^{mod}=(1-\alpha ){\mathcal{L}}_{mlm}+\alpha {\mathbb{E}}_{(I,\hat{T})\sim D}\text{KL}({\text{q}}^{msk}(I,\hat{T})\parallel {\text{p}}^{msk}(I,\hat{T}))\text{\hspace{1em}(7)}$$
上图中，展示了伪目标的top-5候选，其有效的捕获一个图像的相关单词/文本。

​ 作者将$\text{MoD}$应用在下游任务。每个任务的最终损失函数是原始任务损失函数的加权合并，以及模型预测和伪标签的$\text{KL}$散度。为了简单，对于所有的预训练和下游任务设置权重$\alpha =0.4$。

4. 预训练数据集

​ 遵循$\text{UNITER}$，使用两个网络数据集$(\text{Conceptual Captions},\text{SBU Captions})$和两个领域内数据集$(\text{COCO},\text{Visual Genome})$。唯一图像的数量是4M，并且image-text对的数量是5.1M。为了展示本文方法在大规模网络数据的扩展性，作者也引入了更多噪音的$\text{Conceptual 12M}$数据集，总的图像数量增加至14.1M。

5. 实现细节

​ 本文的模型是由具有123.7M参数的${\text{BERT}}_{base}$和具有85.8M参数的$\text{ViT-B/16}$。在8块NVIDIA A100 GPUs上使用batch size为512的方式预训练模型30个epochs。使用具有梯度衰减为0.02的AdamW优化器。在前1000个迭代中学习率预热至$1e^{-4}$，然后按照余弦调度衰减至$1e^{-5}$。在预训练过程中，采用随机图像裁剪分辨率$256\times 256$作为输入 ，并且应用$\text{RandAugment}$。在微调过程中，增加图像分辨率至$384\times 384$，并且为图像patches插入位置编码。更新动量模型的动量参数被设置为0.995，用于image-text对比学习的队列尺寸被设置为65536。在第一个epoch中蒸馏权重$\alpha$中0线性增加至0.4。

三、互信息最大化视角

​ 在本小节中，提供一个$\text{ALBEF}$的可选视角，并展示了其是最大化image-text对不同视角的互信息下边界。$\text{ITC}$、$\text{MLM}$和$\text{MoD}$能够被解释为生成视图的不同方式。

​ 正式来说，定义两个随机变量$a$和$b$为一个数据点的两个不同的视角。在自监督学习中，$a$和$b$是同一图片的两个增强样本。在视觉-语言表示学习中，考虑$a$和$b$是image-text的捕获相同语义的不同变体。目标是学习不随视角变化的表示。这能够通过最大化$a$和$b$的互信息来最大化。在实践中，通过最小化$\text{InfoNCE}$损失函数来最大化$\text{MI(a,b)}$的下边界。
$${\mathcal{L}}_{NCE}=-{\mathbb{E}}_{p(a,b)}\left[\text{log}\frac{\exp (s(a,b))}{{\sum }_{\hat{b}\in \hat{B}}\exp (s(a,\hat{b}))}\right]\text{\hspace{1em}(8)}$$
其中，$s(a,b)$是一个评分函数，$\hat{B}$包含正样本$b$和$|\hat{B}-1|$个负样本。

​ 本文的$\text{ITC}$损失函数能够被重写为：
$${\mathcal{L}}_{itc}=-\frac{1}{2}{\mathbb{E}}_{p(I,T)}[\log \frac{\exp (s(I,T)/\tau )}{{\sum }_{m=1}^M\exp (s(I,T_m)/\tau )}+\log \frac{\exp (s(T,I)/\tau )}{{\sum }_{m=1}^M\exp (s(T,I_m)/\tau )}]\text{\hspace{1em}(9)}$$
​ 最小化${\mathcal{L}}_{itc}$能够看作是最大化对称版本的$\text{InfoNCE}$。因此，$\text{ITC}$将两个独立的模态作为image-text对的两个视图，并训练单模态编码器来最大化正样本对中图像和文本视角的$\text{MI}$。

​ $\text{MLM}$也能够被解释为遮蔽单词与其上下文的最大互信息。具体来说，可以重写$\text{MLM}$损失函数为
$${\mathcal{L}}_{mlm}=-{\mathbb{E}}_{p(I,\hat{T})}\left[\log \frac{\exp (\psi (y^{msk}{)}^{\top }f(I,\hat{T}))}{{\sum }_{y\in \mathcal{V}}\exp (\psi (y{)}^{\top }f(I,\hat{T}))}\right]\text{\hspace{1em}(10)}$$

其中，$\psi (y):\mathcal{V}\to {\mathbb{R}}^d$是多模态编码器输出层的lookup函数，映射单词token $y$至一个向量，并且$\mathcal{V}$是整个词表的集合，并且$f(I,\hat{T})$是一个返回多模态编码器对应的遮蔽上下文的最终hidden state。因此，$\text{MLM}$将image-text对看作是两个视图：(1) 一个随机选择的单词token；(2) 图像+被遮蔽单词的上下文；

​ $\text{ITC}$和$\text{MLM}$通过从image-text对采取部分信息来生成视图。本文的动量蒸馏可以看作是从整个分布来生成可选视图。以等式$(6)$的${\text{ITC}}_{MoD}$为例，最小化$\text{KL}({\text{p}}^{i2t}(I),{\text{q}}^{i2t}(I))$等价于最小化下面的目标函数
$$-\sum _m{q}_m^{i2t}(I)\log p_m^{i2t}(I)=-su{m}_m\frac{\exp (s^{\prime }(I,T_m)/\tau )}{{\sum }_{m=1}^M\exp (s^{\prime }(I,T_m)/\tau )}\log \frac{\exp (s(I,T_m)/\tau )}{{\sum }_{m=1}^M\exp (s(I,T_m)/\tau )}\text{\hspace{1em}(11)}$$
其会最大化与文本共享相似语义的图像$I$的互信息$\text{MI}(I,T_m)$，因为这些文本会有较大的$q_m^{i2t}(I)$。类似地，${\text{ITC}}_{MoD}$也能够最大化与图像相似的$T$的$\text{MI}(I_m,T)$。可以遵循相同的方式，${\text{MLM}}_{MoD}$为遮蔽单词$y^{msk}$生成可选视图$y^{\prime }\in \mathcal{V}$，并最大化$y^{\prime }$和$(I,\hat{T})$的最大化信息$\text{MI}$。因此，动量蒸馏可以看作是在原始视图上执行数据增强。动量模型生成了与原始image-text不同的视图，并鼓励基础模型学习视图不变语义信息的表示。

四、下游 $\text{V+L}$任务

​ 在下游五个$\text{V+L}$任务上应用预训练模型。下面介绍每个任务以及微调策略。

1. $\text{Image-Text}$检索

​ $\text{Image-Text}$包含了两个子任务：image-to-text检索$(\text{TR})$和text-to-image检索$(\text{IR})$。在$\text{Flickr30K}$和$\text{COCO}$基准上评估$\text{ALBEF}$，并且从每个数据集上使用训练样本进行微调预训练模型。对于$\text{Flickr30K}$上的zero-shot检索，在$\text{COCO}$上微调模型进行评估。在微调过程中，联合优化$\text{ITC}$损失函数和$\text{ITM}$损失函数。$\text{ITC}$学习基于单模态相似度的image-text评分函数，而$\text{ITM}$建模图像和文本的细粒度交互来预测匹配分数。由于下游数据集中每个图像包含多个文本，改变$\text{ITC}$中的真实标签来考虑队列中的多个正样本，每个正样本的概率都为1。在推断的过程中，为所有的image-text对计算特征相似分$s_{itc}$。然后，采用$\text{top-k}$作为候选并计算他们的$\text{ITM}$分数$s_{itm}$用于排序。由于$k$能够被设置的非常小，推断的速度会快很多。

2. $\text{Visual Entailment}$

​ $\text{Visual Entailment}$用于预测图片和文本是否具有蕴含、等价或者相反关系的细粒度视觉推理任务。遵循模型$\text{UNITER}$并考虑将$\text{Visual Entailment}$作为三分类问题，然后在[CLS]多模态编码器表示上基础上使用$\text{MLP}$来预测类别概率。

3. $\text{Visual Question Answering(VQA)}$

​ 给定一个图像和一个问题，$\text{VQA}$需要模型预测一个答案。不同于现有的方法将$\text{VQA}$作为多答案分类问题，作者将$\text{VQA}$作为一个答案生成问题。具体来说，使用6层的$\text{Transformer}$解码器来生成答案。如上图$(a)$所示，自回归答案解码器接收多模态嵌入，然后将[CLS]的向量用作解码器的初始输入token。同样的，[SEP]会被追加至解码器输出的后面表示生成的完成。答案解码器使用多模型编码器的预训练权重进行初始化，并且使用条件语言损失函数进行微调。为了与现有的方法进行公平的比较，在推理的过程中约束解码器仅能从3192个后续答案中生成。

4. 用于视觉推理的自然语言$\text{NLVR}$

​ $\text{NLVR}$需要模型判断一条文本是否是一对图像的描述。作者扩展多模态解码器来使其能够在两个图像上进行推理。如上图$(b)$所示，多模态编码器的每层都会被重复为两个连续的$\text{Transformer}$块，每个块都包含一个自注意力层、一个交叉注意力层和一个前向传播层。每层中的两个块会使用相同的预训练权重进行初始化，两个交叉注意力能共享相同的线性投影权重。在训练模型中，两个块接收图像对的两个嵌入集合。在多模态编码器的[CLS]表示上追加一个$\text{MLP}$分类器用于预测。

​ 对于$\text{NLVR}$，执行额外的预训练步骤来为编码图像对准备新的多模态编码器。作者设计了一个文本分配任务$(\text{text-assignment,TA})$：给定一个图像和文本对，模型需要将文本分配给第一个图像、第二个图像、或者都不分配。作者将其定义为一个三分类问题，并在[CLS]表示上使用$\text{FC}$层来预测分配。在$4M$图像上使用$\text{TA}$预训练1个epoch。

5. $\text{Visual Grounding}$

​ $\text{Visual Grounding}$的目标是定位图像中与特定文本描述相关的区域。作者研究了弱监督设置，也就是没有标注的$\text{bounding box}$。作者在$\text{RefCOCO+}$数据集上执行实验，并使用与image-text检索相同策略形成的image-text监督来微调模型。在推断过程中，作者扩展$\text{Grad-CAM}$来获取热图，并使用它们对检测到的对象进行排序。

五、实验

略