融合图像文字输入的Transformer：ViT-BERT
摘要
Transformer是一种基于自注意力机制的深度神经网络，于2017年横空出世。近年来，不同应用领域的Transformer模型，在模型结构和学习方法方面已逐渐趋同于非常相似的范式：以预训练-微调的方式进行训练。基于Transformer的模型通常在大型图像数据集上进行预训练，然后根据特定任务（如NLP的情感分析、计算机视觉的图像分类、目标检测和图像分割等下游任务）进行微调，就能快速实现SOTA的性能，大大提高了模型训练效率。在此基础上，Google与UCLA联合提出了基于Transformer的多模态模型ViT-BERT，将Bert和ViT的表征能力蒸馏到一个模型，使得该模型能同时表征文本和图像特征，大量实验证明这种统一模型的有效性。本文深入分析和探讨这种新型的多模态表征及训练等方法。
特定模态token生成器
ViT-BERT模型[2]只考虑图像和文本两种输入方式。于视觉任务，受到Vision Transformer（ViT）[1]模型的启发，文中将一张图像 切片为一系列大小为 的块图像 ， 表示块图像的个数。然后这些块图像经过flatten，再经过线性投影层 ，最后在patchembedding向量中添加类别 和位置编码 。形式上，图像被处理成如下形式的一系列token：
同理，文本被处理成如下形式的一系列token： 其中， 表示词嵌入矩阵， 是文本位置编码、 是文本段编码。
共享Transformer编码器
共享Transformer编码器（Shared Transformer Encoder）由多个多头自注意力层和MLP层堆叠组成。其中MLP 包含两个具有 GELU 非线性激活函数的全连接层，并且Layer Normalization（LN）操作应用于每个 MSA 或 MLP层之前。如图1所示，共享Transformer编码器的输入为经过特定模态token生成器得到的embedding向量。
特定任务预测模块
作者在这项工作中处理的所有任务，不论是预训练任务还是下游任务，都转换为分类问题。共享Transformer编码器的输出作为特定任务预测模块的输入。文中采用具有 GeLU 激活函数的两层 MLP分类器预测输出类别，形式上有：
公式（3）中， 、 和 、 分别表示第一层MLP的权重和偏置和第二层MLP的权重和偏置。公式（4）中的表示预测输出类别的概率分布。
联合预训练属于多任务学习的范畴，然而由于梯度冲突等优化问题，使得多任务学习极具挑战。联合训练神经网络因同时执行多个任务，通常需要仔细校准以确保任一任务在训练中达到损失平衡。在ViT-BERT模型的联合预训练中同样要考虑上述问题。作者对此提出两个解决方法：（1）采用知识蒸馏为联合训练提供额外的监督。（2）设计了一种迭代梯度掩码策略，以适应来自不同任务的潜在冲突梯度。下面将详细介绍这两种方法。
知识蒸馏
将原始的ViT和BERT模型[4]分别在图像和文本上进行预训练后，作为ViT-BERT模型的两个教师模型。知识蒸馏的核心是最大限度地减少教师模型和学生模型的概率分布之间的Kullback-Leibler (KL) 差异。 表示教师模型的预测概率， 表示学生模型的预测概率， 表示ground-truth标签。蒸馏操作的目标函数可由下式表示：
迭代梯度掩码
但简单地将这两种损失函数相加，自发忽略梯度冲突，这可能会减慢训练速度，且使模型处于次优状态。受到神经网络剪枝方法的启发，文中提出了一种新的梯度掩蔽策略，以协调来自文本预训练和图像预训练的潜在冲突梯度。该方法的主要思想是：保留文本预训练的最重要梯度的子集，而忽略其他梯度，为图像预训练留出空间。设 为Shared Transformer Encoder中的参数， 和 分别是文本预训练和图像预训练任务的损失。作者提出一种自适应掩码 整合文本和图像预训练过程中的梯度，计算过程如下：
其中， 。
该过程将根据 的大小自适应地生成掩码 ，主动为文本预训练选择最重要的梯度，并将其余梯度留给图像预训练。掩码矩阵 在训练过程中会动态更新若干次，自适应产生掩码矩阵的算法如下图2所示。
表1展示了基于纯文本任务和纯视觉任务的微调结果。通过对比表1的实验结果，可得到以下结论：文中提出的ViT-BERT模型可以很好地学习纯文本任务和纯视觉任务的特征表示，但跨模态微调将给下游任务的平均精度带来负面影响；联合训练基线在纯视觉和纯文本任务上都表现出色。
此外，如表2和图3所示的消融实验结果显示，不论是知识蒸馏还是迭代梯度掩码都有利于图像和文本的联合训练。
作者还做了如图4所示的定性实验，发现在训练一段时间后，底层的网络仅对图像或者问文本任务中的一种有强响应，而高层网络则对两种模态都有强响应。该实验结果与解释神经网络主流观点相一致：即随着神经网络的深入，提取到的特征更加抽象，且高层的网络不再关注与特定的任务格式，能够处理更抽象的特征信息。
目前多模态统一模型是深度学习一个重要研究内容。谷歌联合UCLA提出了一个基于Transformer的基础模型——ViT-BERT，该模型能够同时处理文本、图像下游任务。ViT-BERT模型主要结构独立于下游任务，只有输出层与具体下游任务相关，从而实现参数共享。为了提升模型训练的稳定性，文[2]利用知识蒸馏和迭代梯度掩码策略来提高基础模型联合预训练的性能。前者使用已经训练好的单模态模型（BERT和ViT）作为教师模型，在训练过程中，学生模型通过拟合教师模型的输出结果；后者是根据模态更新参数，从而使得训练更加有效，进而防止多模态和大量数据导致的模型拟合困难。广泛实验结果表明，结合预训练的ViT-BERT模型已经为纯视觉和纯文本任务学习了高质量的模型表示。消融研究验证了所提出的知识蒸馏和梯度掩蔽策略的重要性，同时定性研究表明，Transformer编码器顶部block的表示空间在图像和文本之间可以高度共享。该文章构思了一个可以处理各种模态的统一基础模型，但此项工作仅关注图像和文本。还有更多的模态需要考虑，比如视频和音频。同时，其他相关多模态预训练（如，vision-language预训练 [3]）或下游多模态推理任务（例如，视觉问答或图像字幕 [5]）的技术也可能会引起研究热情。
[1] Austin Derrow-Pinion, Jennifer She, David Wong, et al. ETA Predictionwith Graph Neural Networks in Google Maps. 2021
[1] Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., & Houlsby, N. (2021). An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. ArXiv, abs/2010.11929.
[2] Li, Q., Gong, B., Cui, Y., Kondratyuk, D., Du, X., Yang, M., & Brown, M. (2021). Towards a Unified Foundation Model: Jointly Pre-Training Transformers on Unpaired Images and Text. ArXiv, abs/2112.07074.
[3] Yen-Chun Chen, Linjie Li, Licheng Yu, Ahmed El Kholy, Faisal Ahmed, Zhe Gan, Yu Cheng, and Jingjing Liu. Uniter: Universal image-text representation learning. In ECCV, pages 104–120. Springer, 2020. 2, 8.
[4] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. NAACL, 2019. 1, 2, 5.
[5] Xia, Q., Huang, H., Duan, N., Zhang, D.,Ji, L., Sui, Z., Cui, E., Bharti, T., Liu, X., & Zhou, M. (2021). XGPT:Cross-modal Generative Pre-Training for Image Captioning. NLPCC.
往期推荐
壁仞科技研究院作为壁仞科技的前沿研究部门，旨在研究新型智能计算系统的关键技术，重点关注新型架构，先进编译技术和设计方法学，并将逐渐拓展研究方向，探索未来智能系统的各种可能。壁仞科技研究院秉持开放的原则，将积极投入各类产学研合作并参与开源社区的建设，为相关领域的技术进步做出自己的贡献。