UNITER: UNiversal Image-TExt Representation Learning

Introduction

• UNITER: a UNiversal Image-TExt Representation, which can power heterogeneous downstream V+L tasks with joint multimodal embeddings.

Model Architecture

• 总体结构: 给定一个图片句子对，UNITER 将图片的 visual regions 和句子的 textual tokens 作为输入，分别利用 Image Embedder 和 Text Embedder 来提取出它们的 feature embedding。随后，提取出的 embedding 被送入 Transformer 得到一个跨模态的上下文 embedding。同时由于 Transformer 是不带位置信息的，因此还需要额外提供 tokens/regions 的 positions/locations 信息
• Image Embedder: 用一个预训练好的 Faster R-CNN 提取出图片中所有 region 的 pooled ROI features，同时用一个 7-$d$ 向量 $[x_1,y_1,x_2,y_2,w,h,w\ast h]$ (normalized top/left/bottom/right coordinates, width, height, and area.) 来表示每个 RoI 的 location。接着将视觉和位置信息分别送入不同的全连接层来映射到相同的嵌入空间并对它们分别施加 LN。然后将它们相加并通过 LN 层后，就得到了抽取出的 Image feature
• Text Embedder: 类似于 BERT，用 WordPiece 对输入句子进行分词，将 token 和 position 分别通过 embedding 层后得到 word embedding 和 position embedding。将它们相加后通过 LN 层得到 Text feature
• Modality embedding: Modality embedding 类似于 BERT 中的 segment embedding，用于帮助模型区分文本输入和视觉输入。在 Image Embedder 和 Text Embedder 最后的 LN 层之前都会加上 Modality embedding (Modality embedding 在图中被省略了)

---

2 model sizes

• UNITER-base: $L=12,H=768,A=12$ ($L$ 为堆叠的 Transformer 块的个数，$H$ 为隐藏层维数，$A$ 为多头注意力的头数), Total Parameters = $86M$
• UNITER-large: $L=24,H=1024,A=16$, Total Parameters = $303M$

原论文中分别使用了 882 和 3654 个 V100 GPU 来预训练 UNITER-base 和 UNITER-large (恐怖如斯…)

Pre-training tasks

Masked Language Modeling (MLM) $conditioned$ $on$ $image$

• MLM 与 BERT 的处理方式相似，区别在于预测被遮盖 token 时，除了使用未遮盖 token 信息，还使用了视觉信息。值得注意的是，MLM 并不会遮盖图像信息，这是为了防止遮盖的视觉信息正好对预测被遮盖 token 起关键作用 (This prevents potential miss-alignment when a masked region happens to be described by a masked word.)
  

---

Image-Text Matching (ITM)

• 在 ITM 中，新增了特殊符号 [CLS]，在它对应的输出 embedding 后增加一个全连接层，输出 image 和 text 的匹配得分。负样本是通过随机替换一对匹配文本图像中的文本/图像来创造的
  

---

Word-Region Alignment (WRA)

• WRA 利用 Optimal Transport (OT) 理论来显式地鼓励预训练期间单词和图像区域的细粒度对齐
• Intuitively, OT-based learning aims to optimize for distribution matching via minimizing the cost of transporting one distribution to another. In our context, we aim to minimize the cost of transporting the embeddings from image regions to words in a sentence (and vice versa), thus optimizing towards better cross-modal alignment.

---

Masked Region Modeling (MRM) $conditioned$ $on$ $text$

• MRM 与 MLM 类似，模型需要还原出随机遮盖的 image region 对应的视觉特征，遮盖的方式是将输入的 visual features 置零 (在 VL-BERT 中提到，其实应该在整张图片输入 Faster-RCNN 时就将 RoI 对应区域置零，以免被遮盖的 RoI 特征泄露在其他 RoI 的 visual features 中)。出于与 MLM 同样的原因，MRM 并不遮盖文本
• 由于 visual features 是高维连续特征，因此并不能像 MLM 一样采样 log loss。为此，UNITER 提出了 3 个变种：
  • (1) Masked Region Feature Regression (MRFR): 通过在 Transformer 输出后增加一个 FC 层来将输出转化至与 RoI pooled features 相同的维度，然后用 L2 loss 作为损失函数
  • (2) Masked Region Classification (MRC): 并不要求模型还原出 masked RoI pooled features，而是要求模型预测其类别 (FC + softmax + cross-entropy loss)，将 Faster R-CNN 的输出类别作为 ground-truth label
  • (3) Masked Region Classification with KL-Divergence (MRC-kl): MRC 使用 Faster R-CNN 输出的分最高的类别作为硬标签，而 MRC-kl 则是利用 Faster R-CNN 输出的类别概率分布作为软标签，在 Transformer 的输出经过 FC 后计算其与类别概率分布的 KL divergence 作为 loss

消融实验证明，MLM + ITM + MRFR + MRC-kl + WRA 这几个预训练任务的组合能带来最佳效果

在预训练 UNITER 时，每个 mini-batch 都随机选择一个预训练任务进行训练

Pre-training datasets

• UNITER 的预训练任务在 4 个 V+L 数据集的 image-sentence pair 上进行: COCO、Visual Genome (VG)、Conceptual Captions (CC) 和 SBU Captions
• 为了研究不同数据集对预训练的影响，我们将上述四个数据集分为 “In-domain” data 和 “Out-of-domain” data
  • “In-domain” data 由 COCO 中的 image captioning data 和 VG 中的 dense captioning data，称它们为 “In-domain” data 是因为大多数 V+L 任务都是构建在这两个数据集上。为了构建 a ‘fair’ data split，还将 COCO 的训练集和验证集进行了合并，同时清除了任何出现在下游任务验证/测试集中的图片。除此之外，我们还通过 URL 匹配删去了 COCO 和 Visual Genome 中与 Flickr30K 数据集重复的 222 张图片来确保 zero-shot image-text retrieval evaluation on Flickr 是公平的
    
  • “Out-of-domain” data 由 Conceptual Captions 和 SBU Captions 组成，同样有 URL matching 去除了 109 张图片
  • 这样就最终获得了 5.6M 个 image-text pairs 用于训练集，131K 个 image-text pairs 用于验证集
    

Experiments

• UNITER 通过把预训练模型在下游任务数据集上做 fine-tuning，在 6 个 V+L 任务上进行了测试 (new SOTA across six V+L tasks over nine datasets)：

对于需要输入两张图片和一段文本的任务，由于 UNITER 只在 image-sentence pair 上作预训练，如果直接用在这样的预训练任务上效果可能不太理想。为此，可以给每一个图片和文本提供单独的 embedding，然后在两张图片的 embedding 之间加上双向注意力层来更好地学习两张图片之间的关系

Implementation Details 可以参考原论文的 Appendix A.2

References

• UNITER: UNiversal Image-TExt Representation Learning (ICLR, 2020)