UNITER多模态预训练模型原理

1. 数据

过去的5年中，Vision+NLP的研究者所使用的主要数据集如下展示：

本文中所使用到的4种数据集如下图所示,Conceptual Captions和SBU Captions数据集是免费的，但是不是太clean：

2. 算法

2.1 整体结构

这里采用的是two-stage training pipeline训练流程。在一些大量的，含有噪声但是比较cheap的数据上，我们设计一些预训练任务，但是这些预训练任务需要被小心的design，因为我们要获得一个鲁棒的training。然后，把这些预训练好的模型通过fine-tining适应到下流的任务中，同时在下流任务中的数据一般是规模比较小的而且clean的。同时之前的任务一般最多只能适应2-3个task，但是这里我们同时在下游fine-tine9个model，取得了9个SOTA。具体如下图所示：

原始训练数据的形式是如下图所示：

首先通过目标检测模型把image转换成一系列的region，然后把句子tokenize之后转化成一个序列，之后把它们作为输入转化到UNITER中：

UNITER的结构由三个部分构成，第一个部分是Image Embedder，具体来说就是用Faster-RCNN抽取每个region的region feature，同时用一个7维的向量（可能是四个坐标位置加上宽高或者面积等信息构成的7维向量，UNITER并没有直接使用检测器针对每个region输出的类别信息）来encode每一个region的position，将region feature和location feature通过fc层再相加构成了一个Image feature。

在Text Embedder中，就和bert中类似，就是一个token embedding和position embedding相加之后经过一个Layer Norm层，得到了Text Feature。把前面两部分特征经过transfomer就得到了UNITER Model。

除此之外UNITER还设计了三个预训练的任务去训练UNITER模型。第一个是Masked Language Modeling（MLM），即随机mask掉一些词，然后训练模型尝试让他去recover这个词。相同的MRM表示的是我们随机mask掉图片中提取出来的一些区域，然后recover这些缺失的区域。第三个任务是Image-Text-Matching(ITM)，在这个任务中我们抽取整样本图文pair或者负样本图文pair，让模型去预测输入时正还是负。

2.2 具体细节

下面让我们看一下每一个预训练任务具体的细节。

（1）第一个预训练任务MLM，输入是图片和文本对，随机的mask掉一些位置的词语，目标是在这些mask掉的位置让UNITER去recover原本的token，所以这里用到的损失函数是负的log似然函数，具体如下图：

（2）第二个预训练任务是MRM，被mask掉的不是word，而是图片中的region，希望能够recover那些mask掉的region，这里提出了三种方法（对应三种loss）去重构。第一种是让模型recover出的feature对应的vector1和真实的feature对应的vector2之间做L2范数损失。

第二种是Faster-RCNN得到的是class label和bbox，所以我们针对一个region可以让UNITER去预测它的分类，我们希望ground truth的class可以得到最高的置信度分数，具体来说如下图所示，橙色的是ground的label，绿色的是UNITER的输出置信度，我们可以用交叉熵损失，从而让他学习到每一个mask类型的分类。

第三种是我们可以用Faster-RCNN输出的类别分布做损失，橙色的是Faster-RCNN的输出，我们希望UNITER输出的分布尽可能接近橙色的分布，具体见下图：

（3）第三个预训练任务是ITM，每当我们输入一个image-text pair，我们希望通过随机mask掉image或者text来建造一个negative pair，我们的目标是预测输入的image-text pair是不是存在对应关系，所以这是一个二分类问题。

前面介绍了这么多预训练任务和方法，具体哪些方法对模型性能提升有效果，这个在实验中具体介绍。

3.计算

3.1 计算资源

计算基本资源如下所示，每个virtual machine（VM）是有16个机器，每个机器上有4块V100的卡。对于UNITER base来说，预训练4.5天即可，加上finetune基本上10天时间确定一个模型的performance足够，UNITER large需要4个VM，预训练需要10天时间。

3.2 Optimization on Limited GPUs

本文中一共用了三种计算优化方法,下面依次介绍一下。

3.2.1 Dynamic Batching

第一种就是动态batching，因为transformer的计算复杂度是和输入序列的长度的二次方呈线性相关，这里的长度在我们的vision+language里面是词的数量加上region的长度。一般的batching是像左边图这样，把输入padding至相同的长度，所以可以看见浪费很多的计算空间，而我们的做法是把输入按照长度进行排序，然后把长度相近的放在同一个batch里，这样我们就减少了padding从而节省了计算时间。

3.2.2 Gradient Accumulation

UNITER large里面用了四个Virtual Machine（VM），而时间主要花费在network communication overhead上，我们用Gradient Accumulation减少交流的频率，从而增加总吞吐量。

3.2.3 Mixed-precision Training

这里同时应用了16bit和32bit的图像表示，节省了空间，从而batchsize可以增大，从而减少训练时间。

4.下游任务具体实现过程

4.1 Visual Question Answering

对于一个图片，回答和图片相关的问题。我们把图片和问题输入到UNITER中，输出是答案的分布，得到最大概率的答案作为我们的预测答案，我们这里用交叉熵损失监督这个下有任务的fine-tuning。

4.2 Visual Entailment

问题描述为图片是premise，文本为Hypothesis，目标是预测文本text是不是Entail图片。有三种label，Entailment（确定）、Neutral（中立）、Contradiction（反对），输出是三个类别，用交叉熵监督。

4.3 Natural Language for Vision Reasoning

问题描述为输入两张图片和一个text描述，判断文字是否为两张图片的描述。这里结构做了一定的改动，我们把输入的两张图片分别put到两个UNITER结构种，并且把text描述double一下，使得两边是一样的描述，这样就得到了联合特征表达，把这个联合特征表达经过一个叫Bi-Attention层输出预测的label。

4.4 Visual Commonsense Reasoning

问题描述为对于一个问题有四个备选答案，model必须从四个备选答案中选择出一个作为prediction，并且从四个备选的理由中选出选择这个答案的理由。所以在训练中，我们把问题和每个备选的答案concat到一块，分别和图片一起输入到UNITER中得到了四个score，最高score的答案作为最后的预测，同理在选择理由的时候我们也可以把问题答案和理由的选项concat到一起，输入到UNITER中，然后用交叉熵去监督训练。

4.5 Referring Expression Comprehension

问题描述为输入是一个句子，模型要在图片中找到对应的区域。之前有讲到，其实图片部分是表征为一系列的region集合，所以可以每一个region都输出一个score，得到最高score的region作为预测的region。

4.6 Image-Text Retrieval

这个问题描述是给定一个text，找出对应的图片，或者给定一个图片，找出对应的text。和前面介绍到的图文匹配很像，通过选取positive还是negative pair的方式来训练model，UNITER的目标是判断输入的pair是positive还是negative。

5.实验

5.1 消融实验

为了验证前面哪种方式对模型性能提升最大，此处针对五个下游任务做了消融实验，通过是将五个任务的结果合并成一个Meta-Sum的评价指标用来衡量模型的好坏。同时在消融实验中得到了几个结论如下所示：

1、Language Pre-training is better than No pretraining.

2、Vision + Language together > Vision/Language

3、Best combination : MLM + ITM + MRC-kl + MRFR

4、More data the better

5.2 实验结果

多项任务取得SOTA

6.核心代码阅读

UNITER模型的核心在根目录下的model文件夹下的model.py中，其中包含UniterConfig类、UniterPreTrainedModel类、UniterModel类（主模型类）。下面就介绍一下UniterConfig类和UniterModel类，同时介绍一下构成UniterModel类的一些子类。github代码网址链接：https://github.com/ChenRocks/UNITER

6.1 UniterConfig类

该类主要是配置Uniter模型的一些超参数，具体的超参数含义在如下代码后面标识。

class UniterConfig(object):"""Configuration class to store the configuration of a `UniterModel`."""def __init__(self,vocab_size_or_config_json_file,  #在`UniterModel`中'inputs_ids'的词表sizehidden_size=768,                 #编码层和池化层的sizenum_hidden_layers=12,            #transformer encoder中隐藏层的数量num_attention_heads=12,          #attention head的数量对于每个transformer encoder隐藏层intermediate_size=3072,          #size of "intermediate" (i.e.feed-forward) layer in the Transformer encoderhidden_act="gelu",               #激活函数，swish和relu同样支持hidden_dropout_prob=0.1,         #在embedding，编码和池化中的全连接层drop概率attention_probs_dropout_prob=0.1,#attention中的drop概率max_position_embeddings=512,     #最大序列长度，通常设置为(e.g., 512 or 1024 or 2048)type_vocab_size=2,               #`token_type_ids`的词表sizeinitializer_range=0.02):         #初始权重矩阵的标准差if isinstance(vocab_size_or_config_json_file, str):with open(vocab_size_or_config_json_file,"r", encoding='utf-8') as reader:json_config = json.loads(reader.read())for key, value in json_config.items():self.__dict__[key] = valueelif isinstance(vocab_size_or_config_json_file, int):self.vocab_size = vocab_size_or_config_json_fileself.hidden_size = hidden_sizeself.num_hidden_layers = num_hidden_layersself.num_attention_heads = num_attention_headsself.hidden_act = hidden_actself.intermediate_size = intermediate_sizeself.hidden_dropout_prob = hidden_dropout_probself.attention_probs_dropout_prob = attention_probs_dropout_probself.max_position_embeddings = max_position_embeddingsself.type_vocab_size = type_vocab_sizeself.initializer_range = initializer_rangeelse:raise ValueError("First argument must be either a vocabulary size ""(int) or the path to a pretrained model config ""file (str)")@classmethoddef from_dict(cls, json_object):"""Constructs a `UniterConfig` from aPython dictionary of parameters."""config = UniterConfig(vocab_size_or_config_json_file=-1)for key, value in json_object.items():config.__dict__[key] = valuereturn config@classmethoddef from_json_file(cls, json_file):"""Constructs a `UniterConfig` from a json file of parameters."""with open(json_file, "r", encoding='utf-8') as reader:text = reader.read()return cls.from_dict(json.loads(text))def __repr__(self):return str(self.to_json_string())def to_dict(self):"""Serializes this instance to a Python dictionary."""output = copy.deepcopy(self.__dict__)return outputdef to_json_string(self):"""Serializes this instance to a JSON string."""return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"

6.2 UniterModel类

UniterModel的主模块，首先是在构造函数中先进行UniterTextEmbeddings、UniterImageEmbeddings、UniterEncoder等类的初始化。分仅有text，仅有image或者既有text又有image三种情况构造embedding layer。

class UniterModel(UniterPreTrainedModel):""" Modification for Joint Vision-Language Encoding"""def __init__(self, config, img_dim):super().__init__(config)self.embeddings = UniterTextEmbeddings(config)  # textEmbedding类self.img_embeddings = UniterImageEmbeddings(config, img_dim) # imageEmbedding类self.encoder = UniterEncoder(config)            # encoder类使用的就是bert的encoderself.pooler = BertPooler(config)self.apply(self.init_weights)def _compute_txt_embeddings(self, input_ids, position_ids,txt_type_ids=None):output = self.embeddings(input_ids, position_ids, txt_type_ids)return outputdef _compute_img_embeddings(self, img_feat, img_pos_feat, img_masks=None,img_type_ids=None):if img_type_ids is None:img_type_ids = torch.ones_like(img_feat[:, :, 0].long())img_type_embeddings = self.embeddings.token_type_embeddings(img_type_ids)output = self.img_embeddings(img_feat, img_pos_feat,img_type_embeddings, img_masks)return outputdef _compute_img_txt_embeddings(self, input_ids, position_ids,img_feat, img_pos_feat,gather_index, img_masks=None,txt_type_ids=None, img_type_ids=None):txt_emb = self._compute_txt_embeddings(input_ids, position_ids, txt_type_ids)img_emb = self._compute_img_embeddings(img_feat, img_pos_feat, img_masks, img_type_ids)# align back to most compact inputgather_index = gather_index.unsqueeze(-1).expand(-1, -1, self.config.hidden_size)embedding_output = torch.gather(torch.cat([txt_emb, img_emb], dim=1),dim=1, index=gather_index)return embedding_outputdef forward(self, input_ids, position_ids,img_feat, img_pos_feat,attention_mask, gather_index=None, img_masks=None,output_all_encoded_layers=True,txt_type_ids=None, img_type_ids=None):# compute self-attention maskextended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype)  # fp16 compatibilityextended_attention_mask = (1.0 - extended_attention_mask) * -10000.0# embedding layerif input_ids is None:# image onlyembedding_output = self._compute_img_embeddings(img_feat, img_pos_feat, img_masks, img_type_ids)elif img_feat is None:# text onlyembedding_output = self._compute_txt_embeddings(input_ids, position_ids, txt_type_ids)else:embedding_output = self._compute_img_txt_embeddings(input_ids, position_ids,img_feat, img_pos_feat,gather_index, img_masks, txt_type_ids, img_type_ids)encoded_layers = self.encoder(embedding_output, extended_attention_mask,output_all_encoded_layers=output_all_encoded_layers)if not output_all_encoded_layers:encoded_layers = encoded_layers[-1]return encoded_layers

6.3 UniterModel模块中的关键子模块

这里就介绍UniterTextEmbeddings、UniterImageEmbeddings、UniterEncoder三个关键子模块的代码。

6.3.1 UniterTextEmbeddings类

此类的作用是将训练语料中的text信息转化成embedding向量，embedding向量的构成分为三个部分words_embeddings（词嵌入向量）、position_embeddings（位置嵌入向量）、token_type_embeddings（是否为前后语句对嵌入向量）。

class UniterTextEmbeddings(nn.Module):def __init__(self, config):super().__init__()self.word_embeddings = nn.Embedding(config.vocab_size,config.hidden_size, padding_idx=0)self.position_embeddings = nn.Embedding(config.max_position_embeddings,config.hidden_size)self.token_type_embeddings = nn.Embedding(config.type_vocab_size,config.hidden_size)# self.LayerNorm is not snake-cased to stick with TensorFlow model# variable name and be able to load any TensorFlow checkpoint fileself.LayerNorm = FusedLayerNorm(config.hidden_size, eps=1e-12)self.dropout = nn.Dropout(config.hidden_dropout_prob)def forward(self, input_ids, position_ids, token_type_ids=None):if token_type_ids is None:token_type_ids = torch.zeros_like(input_ids)words_embeddings = self.word_embeddings(input_ids)position_embeddings = self.position_embeddings(position_ids)token_type_embeddings = self.token_type_embeddings(token_type_ids)embeddings = (words_embeddings+ position_embeddings+ token_type_embeddings)embeddings = self.LayerNorm(embeddings)embeddings = self.dropout(embeddings)return embeddings

6.3.2 UniterImageEmbeddings类

Uniter模型中Image转化为embedding向量是关键，img_feat和img_pos_feat是如何生成的可以去前面2.1小节看一下，简单来说就是使用Faster-RCNN检测出多个region，region feature特征经过reshape或者dnn等方式转为img_feat，region的location转为img_pos_feat。然后分别经过fc层至相同维度，相加再经过LN层就生成了既含有图片特征又含有位置信息的embedding。

class UniterImageEmbeddings(nn.Module):def __init__(self, config, img_dim):  # img_dim在utils文件夹下的const.py内，默认等于2048super().__init__()self.img_linear = nn.Linear(img_dim, config.hidden_size)self.img_layer_norm = FusedLayerNorm(config.hidden_size, eps=1e-12)self.pos_layer_norm = FusedLayerNorm(config.hidden_size, eps=1e-12)self.pos_linear = nn.Linear(7, config.hidden_size)self.mask_embedding = nn.Embedding(2, img_dim, padding_idx=0) # 第一维度为2代表是否mask# tf naming convention for layer normself.LayerNorm = FusedLayerNorm(config.hidden_size, eps=1e-12)self.dropout = nn.Dropout(config.hidden_dropout_prob)def forward(self, img_feat, img_pos_feat, type_embeddings, img_masks=None):if img_masks is not None:self.mask_embedding.weight.data[0, :].fill_(0)mask = self.mask_embedding(img_masks.long())img_feat = img_feat + masktransformed_im = self.img_layer_norm(self.img_linear(img_feat)) # img_feat做全连接transformed_pos = self.pos_layer_norm(self.pos_linear(img_pos_feat)) # img_pos_feat做全连接embeddings = transformed_im + transformed_pos + type_embeddingsembeddings = self.LayerNorm(embeddings)embeddings = self.dropout(embeddings)return embeddings

6.3.3 UniterEncoder类

Uniter模型中的Encoder部分使用的就是原始Bert中的encoder层，具体的可以去看pytorch版bert源码。

class UniterEncoder(nn.Module):def __init__(self, config):super().__init__()layer = BertLayer(config)self.layer = nn.ModuleList([copy.deepcopy(layer)for _ in range(config.num_hidden_layers)])def forward(self, input_, attention_mask,output_all_encoded_layers=True):all_encoder_layers = []hidden_states = input_for layer_module in self.layer:hidden_states = layer_module(hidden_states, attention_mask)if output_all_encoded_layers:all_encoder_layers.append(hidden_states)if not output_all_encoded_layers:all_encoder_layers.append(hidden_states)return all_encoder_layers