[人工智能]UNITER多模态预训练模型原理加代码解读

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'的词表size
                 hidden_size=768,                 #编码层和池化层的size
                 num_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 encoder
                 hidden_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`的词表size
                 initializer_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] = value
        elif isinstance(vocab_size_or_config_json_file, int):
            self.vocab_size = vocab_size_or_config_json_file
            self.hidden_size = hidden_size
            self.num_hidden_layers = num_hidden_layers
            self.num_attention_heads = num_attention_heads
            self.hidden_act = hidden_act
            self.intermediate_size = intermediate_size
            self.hidden_dropout_prob = hidden_dropout_prob
            self.attention_probs_dropout_prob = attention_probs_dropout_prob
            self.max_position_embeddings = max_position_embeddings
            self.type_vocab_size = type_vocab_size
            self.initializer_range = initializer_range
        else:
            raise ValueError("First argument must be either a vocabulary size "
                             "(int) or the path to a pretrained model config "
                             "file (str)")

    @classmethod
    def from_dict(cls, json_object):
        """Constructs a `UniterConfig` from a
           Python dictionary of parameters."""
        config = UniterConfig(vocab_size_or_config_json_file=-1)
        for key, value in json_object.items():
            config.__dict__[key] = value
        return config

    @classmethod
    def 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 output

    def 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的encoder
        self.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 output

    def _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 output

    def _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 input
        gather_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_output

    def 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 mask
        extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
        extended_attention_mask = extended_attention_mask.to(
            dtype=next(self.parameters()).dtype)  # fp16 compatibility
        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0

        # embedding layer
        if input_ids is None:
            # image only
            embedding_output = self._compute_img_embeddings(
                img_feat, img_pos_feat, img_masks, img_type_ids)
        elif img_feat is None:
            # text only
            embedding_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 file
        self.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内，默认等于2048
        super().__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 norm
        self.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 + mask

        transformed_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_embeddings
        embeddings = 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