图文结合-UNITER

本文介绍一篇图文结合的经典论文，论文发布于2020年

论文信息

论文题目：

UNITER: UNiversal Image-TExt Representation Learning

论文地址：

https://arxiv.org/abs/1909.11740

代码地址：

https://github.com/ChenRocks/UNITER

主要内容

这篇论文比之前的ViLT和SOHO都要早，其提出的模型UNITER做的也更加复杂。

UNITER与其他模型的不同之处：

• 使用了条件mask。也就是当对文本进行mask时，要保证图片对应的语义没问题，反之亦然。
• 提出了word-region对齐的预训练任务

1、模型结构

模型结构如下图所示：

从模型的结构上看，，UNITER其实是一种单流模型。它分别对文本和图片进行编码，并将位置信息融合到编码中得到各自的特征向量，接着将两者的特征向量concat后输入交互层（多层Transformer），最后输出。

图片侧，UNITER使用Faster R-CNN去提取视觉特征（pooling后的ROI特征），然后用用一个七维的向量（$[x_1,y_1,x_2,y_2,w,h,w\times h]$）来编码每个区域的位置特征。视觉和位置特征编码都会经过一个FC层去映射到同样的embedding空间。接着，两者相加后经过LN输出。

文本侧，使用bert得到对应的token emb和pos emb。接着，使用和图片同样的处理方式得到对应的文本向量表征。

文本的embedding代码实现为：

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

图片的embedding代码实现为：

class UniterImageEmbeddings(nn.Module):
    def __init__(self, config, img_dim):
        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)

        # 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))
        transformed_pos = self.pos_layer_norm(self.pos_linear(img_pos_feat))
        embeddings = transformed_im + transformed_pos + type_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings

上面代码中的FusedLayerNorm实现的功能与LN是一致的，不过FusedLayerNorm的速度相比torch.nn.layernorm会快一点（这里仅指计算transformer的结构，也就是nlp和cv领域，看github上给出的是，计算(a,b,c)三维数组时，会快很多）。

2、预训练任务

（1）MLM

masked language modeling conditioned on image regions。

与bert的处理一致，不过增加了mask条件。即不会随机屏蔽两种模态，而是每次只屏蔽一种模态的数据，并且需要保证另一种模态不变。以此，可以防止潜在的错位。以下图为例：

当文本中的"dog"被mask后，conditional masking中，图片上的狗就不会被mask，但是，如果使用joint random masking就可能将图片上的狗mask，此时，模型就需要盲目的进行预测，可能导致错位。

同时，作者也对这一猜测进行了验证：

上图是在不同数据集上两种mask策略的对比，可以观察到conditional masking更好。

与bert不同的是，对文本中mask部分的预测，不仅会利用到上下文的信息，还会用到所有图片region的信息。其损失函数为：

$L_{MLM}(\theta )=-E_{(w,v)\in D}log{P}_{\theta }(w_m|w_{/m},v)$

其中图片region为$v=v_1,...,v_K$，输入文本为$w=w_1,...,w_T$，mask的索引为$m\in N^M$，$w_m$表示文本中被mask的位置，$w_{/m}$表示文本中未被mask的上下文。

（2）MRM

masked region modeling conditioned on input text。

region mask通过将视觉特征向量进行全0替换来实现。其mask方式与MLM一致。

作者认为，视觉向量表征与文本不同，它是一种高维的、连续的形式，因此不能和文本一样直接用类别的最大似然来求解。基于此，作者提出了三个不同的变体，从不同角度对masked region进行拟合。

并且，三个变体的loss求解都以下式作为基础：
$L_{MRM}(\theta )=E_{(w,v)\in D}{f}_0(v_m|v_{/m},w)$

masked region feature regression(MRFR)

学习$v_m$在transformer输出到视觉特征的回归。进一步说，使用FC层将transformer层的输出转化为$h_0(v_m^{(i)})$，转化后的向量维度与输入的视觉特征（$r(v_m^{(i)})$）的维度一致。最后，对两者执行岭回归（L2回归），其计算方式如下：

$f_0(v_m|v_{/m})={\sum }_{i=1}^M||h_0(v_m^{(i)})-r(v_m^{(i)})|{|}_2$

masked region classification(MRC)

该任务根据语义用来预测每个masked region的类别。将transformer的输出经过FC层，预测每个类别的得分，然后经过softmax将得分进行归一化。因为数据没有提供对应的标签，所以作者将Faster R-CNN检测到的对象类别，作为标签，并将其转化为one-hot向量。最后，需要最小化交叉熵函数，其计算方式为：
$f_0(v_m|v_{/m},w)={\sum }_{i=1}^{M}CE(c(v_m^{(i)}),g_{\theta }(v_m^{(i)}))$
其中$c(v_m^{(i)})$表示one-hot标签向量，$g_{\theta }(v_m^{(i)})$表示经过softmax归一化后的得分输出。

masked region classification with kl-divergence(MRC-KL)

作者认为，MRC将目标检测模型中最有可能的目标类别作为硬标签，并假设检测到的目标类别是该区域的真实标签。但是，可能该标签不是真实的，因此，作者提出了MRC-KL的任务。

该任务将Faster R-CNN的输出作为软标签，通过KL散度，最小化两者之间的分布差异。其计算公式为：
$f_0(v_m|v_{/m},w)={\sum }_{i=1}^M{D}_{KL}(\tilde{c}(v_m^{(i)})||g_{\theta }(v_m^{(i)}))$

（3）ITM

image-text matching。

输入需要添加一个特殊令牌[CLS]，用来表示两种模态的融合表征。其中，负样本通过用从其他样本中随机选择的图像或文本替换配对样本中的图像或文本来构建。该任务的损失函数为：

$L_{ITM}(\theta )=-E_{(w,v)\in D}[ylog{s}_0(w,v)+(1-y)log(1-s_0(w,v))]$

（4）WRA

word-region alignment。

通过optimal transport（将image emb与text emb对应的最小代价）学习细粒度的word-region对齐信息。之所以选择OT是有以下原因：

• 自归一化。T的所有元素总和为1
• 稀疏性。当精确求解时，OT最多产生(2r-1)个非0解，$r=max(K,T)$，使得文本和图像对齐的更加鲁棒，以及更具有可解释性
• 与传统的线性规划求解相比，OT可以很容易的使用迭代过程获得，因此只需要进行矩阵乘积，适合大规模模型预训练。

该任务的损失函数为：
$L_{WRA}(\theta )=D_{ot}(\mu ,v)=mi{n}_{T\in \prod (a,b)}{\sum }_{i=1}^T{\sum }_{j=1}^K{T}_{ij}\cdot c(w_i,v_j)$

其中 ∏ ( a , b ) = { T ∈ R + T × K ∣ T 1 m = a , T T 1 n = b } \prod(a,b)=\{T\in{R_{+}^{T\times{K}}|T1_m=a,T^T1_n=b}\} ∏(a,b)={T∈R+T×K​∣T1m​=a,TT1n​=b}，$l_n$表示n维的全1向量，$c(w_i,v_j)$表示代价函数，用来计算$w_i$和$v_j$之间的距离。在实验过程中，作者使用的代价函数是$c(\cdot )=1-\frac{w_i^T{v}_j}{||w_i|{|}_2||v_j|{|}_2}$。T表示transport plan，用来解释两个模态之间的对齐情况。

由于很难精确计算T的最小值，所以一般使用IPOT算法去近似OT距离。最后，OT距离作为WRA的loss，对参数进行更新。

作者对每个预训练任务做了消融实验，实验结果如下所示：

可以看到，预训练任务的数量越多，最后在下游任务的表现也越好。同时，作者对比了没有预训练时的情况，如上表所示第一行，表示没有预训练，第二行表示基于文本，使用MLM进行了预训练，可以看到用经过MLM预训练的模型参数，对UNITER进行初始化，可以很好的提升UNITER的基线水平，即使该初始化参数并没有任何的图像信息。同时，实验发现NRC-kl的效果优于MRC。因此，作者最后使用的经过MLM训练的参数对模型进行初始化。

3、下游任务

作者在做下游任务实验时，选择的表现得最好的预训练任务的组合，并且结合了领域内和领域外的数据。

作者对了各个模型在不同的下游任务上的表现，实验结果如下：

同时，针对VCR任务，作者提出了两阶段预训练的方法：

1. 在标准的数据集上预训练
2. 在VCR数据上预训练

作者，发现第一阶段的预训练对VCR数据集的指标提升有限，第二阶段则有很大提升。这多半是因为第一阶段用于预训练的数据集中，并没有VCR这种结构的数据，不过，在加入类似的数据训练后，模型就可以很好的处理这种类型的数据了。

实验结果如下表4所示：

除此之外，NLVR2数据集会输入两张图片，这与预训练的情况不同，如果直接进行微调，结果较差。因此，作者提出了以下三种改进的方法：

1. triplet：将图片对和问题的embedding联合
2. pair：问题分别与每个图像进行组合，也就是形成两个输入，一个pair
3. pair-biattn：在pair的基础上添加了双向注意力网络，去学习成对图片之间的交互

实验结果如上表5展示，可以看到第三种方法是最优的。

作者的下游任务实验表明，通过对UNITER的顶层进行改动，可以使其适应与预训练任务截然不同的新的下游任务。

另外，作者对UNITER中注意力的工作模式进行了观察，发现注意力机制以模态间和模态内的方式运作。具体可以分为以下几种：

• vertical:attention to special tokens [CLS] or [SEP]
• diagonal:attention to the token/region itself or preceding/following tokens/regions
• vertical+diagonal:minture
• block:intra-modality attention,i.e,textual self-attn ad visual self-attn
• heterogeneous:diverse attentions that cannot be categorized and is highly dependent on actual input
• reversed block:inter-modality attention,i.e,text-to-image and image-to-text attention

如下图所示：

作者对模态间的注意力进行了演示，如下图所示：

总结

作者从模型结构、预训练任务等多个方向进行了创新，同时，还解决了模型在下游任务上的兼容性问题。显然，由于预训练的任务比较多，作者耗费了大量的资源来进行预训练。模型的规模如下：

UNITER-base: L=12, H=768, A=12, Total Parameters=86M. UNITER-large:
L=24, H=1024, A=16, Total Parameters=303M (L: number of stacked Transformer
blocks; H: hidden activation dimension; A: number of attention heads). 882 and 3645
V100 GPU hours were used for pre-training UNITER-base and UNITER-large.

不过，看上面模型的参数量好像还没bert的多，bert里的词表大小为30522，作者用的此表大小为28996，以及在实际实验时，UNITER选择的是动态句子长度（也就是没有固定max_seq_len），这应该是导致最后UNITER参数量略小于bert的原因。

另外，不得不说的是，UNITER这篇文章的代码写的真的清晰啊，读起来太舒服了，不像某些论文，给的代码不忍直视。不过UNITER进行了混合精度加速，所以在模型训练这一块，如果不了解amp的话，可能读起来有点费劲。

使用amp进行混合精度训练的代码示例如下：

# Declare model and optimizer as usual, with default (FP32) precision
model = torch.nn.Linear(D_in, D_out).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

# Allow Amp to perform casts as required by the opt_level
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
...
# loss.backward() becomes:
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
...

更加详细的可以参考官方文档：

https://nvidia.github.io/apex/amp.html