LayoutLMv2: 面向视觉丰富文档理解的多模态预训练

1. 引言

   LayoutLMv2 [1] 是LayoutLM [2] 的后续改进版本，不熟悉的读者可以参考笔者的上一篇博文[3]。LayoutLM进行文本、版面的预训练，只在下游任务加入图像特征，而LayoutLMv2进一步将图像特征纳入到统一的多模态预训练框架中。具体地，如图1所示，LayoutLMv2将CNN图像特征展开为一维的token序列，并与文本token序列进行拼接，同时输入到Transformer [4]网络进行预训练。LayoutLMv2除了采用已有的掩码视觉语言模型预训练任务，还提出了新的文本-图像对齐与文本-图像匹配任务，从而令模型更好地学习到跨模态的交互作用。此外，LayoutLMv2在Transformer网络中引入了空间感知的自注意力机制，令模型更好地理解文档中不同文本块的相对位置关系。LayoutLMv2大幅优于LayoutLM等基线模型，在多个视觉丰富文档理解任务上取得了领先的结果。

2. 方法

2.1 模型结构

  LayoutLMv2采用Transformer作为骨干网络，并接受三种模态的输入：文本、图像与版面。每种模态的输入都转化为一个特征序列，并由模型进行特征融合。

  文本特征。文档内容由现有的OCR工具或PDF解析器进行获取，并按照从左到右、从上到下的顺序进行序列化。文本序列采用WordPiece方法进行分词，获取token序列，且每个token分配有一个块标记 $s_i\in\{[A], [B]\}$。然后，[CLS] token与[SEP] token被分别添加到token序列的起始与结束位置。token序列的长度被限制于最大长度 $L$以内；若序列长度小于 $L$，则填充[PAD] token至[SEP] token之后直到最大长度 $L$。最终的token序列表示如下：

$$S=\{[\text{CLS}],w_1,w_2,...,[\text{SEP}],[\text{PAD}],[\text{PAD}],...\},|S|=L$$

最终的文本特征是token特征、1D位置特征与块特征这三个特征之和，具体如下：

$$\bold{t}_i=\text{TokEmb}(w_i)+\text{PosEmb1D}(i)+\text{SegEmb}(s_i), 0\le i<L$$

  视觉特征。LayoutLMv2采用ResNetXt-FPN [5,6] 作为视觉编码器的骨干网络。给定一张文档图像 $I$，将其缩放到224x224大小并输入视觉网络，得到宽度为 $W$、高度为 $H$的特征图，并序列化为长度为 $WH$的视觉token序列。论文中取 $W=H=7$。为了统一文本与视觉特征的维度，论文采用了一个线性层对视觉特征进行映射。由于CNN特征无法捕捉位置信息，同样地为其添加1D位置特征。此外，对于块特征，每个视觉token都分配有一个视觉块标记[C]。因此，视觉特征具体表示如下：

$$\bold{v}_i=\text{Proj}(\text{VisTokEmb}(I)_i)+\text{PosEmb1D}(i)+\text{SegEmb([C])}, 0\le i<WH$$

  版面特征。版面特征指的是文档中token包围盒所表示的空间位置信息。与LayoutLM一致，LayoutLMv2将坐标归一化、离散化至[0, 1000]之间，并采用两个嵌入层分别用于获取x-轴特征与y-轴特征。记第 $i$个文本/视觉token的包围盒为 $\text{box}_i=(x_0,x_1,y_0,y_1,w,h)$，则版面特征具体表示如下：

$$\bold{l}_i=\text{Concat}(\text{PosEmb2D}_x(x_0,x_1,w),\text{PosEmb2D}_y(y_0,y_1,h)), 0\le i<WH+L$$

对于视觉token，它们被视为在原图上均匀分割的网格，故其包围盒坐标能够容易计算得到。对于特殊的[CLS]、[SEP]与[PAD]等token，它们的包围盒置为 $\text{box}_{\text{PAD}}=(0,0,0,0,0,0)$。

  结合空间感知自注意力机制的多模态编码器。编码器首先将视觉特征 $\{\bold{v}_0,...,\bold{v}_{WH-1}\}$与文本特征 $\{\bold{t}_0,...,\bold{t}_{L-1}\}$进行拼接，得到统一的特征序列 $X$，并与版面特征 $\{\bold{l}_0,...,\bold{l}_{WH+L}\}$进行叠加，得到第一层的输入 $\bold{x}^{(0)}$:

$$\bold{x}_i^{(0)}=X_i+\bold{l}_i, 其中X=\{\bold{v}_0,...,\bold{v}_{WH-1},\bold{t}_0,...,\bold{t}_{L-1}\}$$

输入特征经过经过Transformer的多头自注意力机制进行多模态特征学习。应注意的是，原始的自注意机制缺乏相对位置信息: 对于query向量 $\bold{x}_i$、key向量 $\bold{x}_j$，原始的自注意机制计算如下的注意力得分：

$$\alpha_{ij}=\frac{1}{\sqrt{d_{head}}}(\bold{x}_i\bold{W}^Q)(\bold{x}_j\bold{W}^K)^T$$

为了更侧重局部特征，LayoutLMv2引入了相对位置偏置项，得到如下的空间感知注意力得分：

$$\alpha'_{ij}=\alpha_{ij}+\bold{b}^{(1D)}_{j-i}+\bold{b}^{(2D_x)}_{x_j-x_i}+\bold{b}^{(2D_y)}_{y_j-y_i}$$

其中 $\bold{b}^{(1D)}、\bold{b}^{(2D_x)}、\bold{b}^{(2D_y)}$分别为可学习的1D与2D相对位置偏置。最终的输出特征为：最终的输出特征为：

$$\mathbf{h}_{i}=\sum_{j} \frac{\exp \left(\alpha_{i j}^{\prime}\right)}{\sum_{k} \exp \left(\alpha_{i k}^{\prime}\right)} \mathbf{x}_{j} \mathbf{W}^{V}$$

2.2 预训练任务

  预训练任务#1：与LayoutLM类似，采用掩码视觉语言模型MVLM，此处不再赘述。

  预训练任务#2：文本-图像对齐（Text-Image Alignment, TIA）。在TIA任务中，随机遮掩文档图像的部分文本行，并令模型预测文本序列中每个token对应的位置是否被遮掩。TIA任务帮助模型对齐文本和图像的位置信息。

  预训练任务#3：文本-图像匹配（Text-Image Matching, TIM）。在TIM任务中，随机替换文档图像，并令模型预测文本序列与文档图像是否匹配。TIM任务帮助模型对齐文本和图像的内容信息。

3 实验

   预训练。与LayoutLM一致，LayoutLMv2模型在IIT-CDIP数据集（包含千万量级扫描文档）上进行预训练。文本内容与单词级别包围盒信息通过Microsoft Read API进行提取。研究人员训练了 BASE、LARGE 两种规模的模型，更多预训练细节参见论文3.2节。

   下游任务。前四个任务都是序列标记任务，需要为每个token预测BIO标记，包括表单理解（FUNSD数据集）、票据理解（CORD数据集、SROIE数据集）、合同文档理解（Kleister-NDA数据集）。第五个任务为文档级别分类任务（RVL-CDIP数据集），将[CLS] token特征、图像token经过Transformer前后的平均特征进行拼接，并输入分类层。第六个任务为文档VQA任务（Doc-VQA数据集）。

4 结论

  针对多模态文档理解任务提出的多模态预训练模型 LayoutLMv2，将文本、版面、图像信息融合到了统一的多模态预训练框架中，并且提出了新的预训练目标，以增强跨模态信息交互。LayoutLMv2在6个不同的视觉丰富文档理解任务中取得了大幅提升， FUNSD (0.7895 → 0.8420), CORD (0.9493 → 0.9601), SROIE (0.9524 → 0.9781), Kleister-NDA (0.834 → 0.852), RVL-CDIP (0.9443 → 0.9564), and DocVQA (0.7295 → 0.8672)，有助于真实场景的文档理解。

[1] Xu Y, Xu Y, Lv T, et al. LayoutLMv2: Multi-modal pre-training for visually-rich document understanding. ACL 2021.

[2] Xu Y, Li M, Cui L, et al. Layoutlm: Pre-training of text and layout for document image understanding. Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining. 2020: 1192-1200.

[3] https://bbs.huaweicloud.com/blogs/298660

[4] Vaswani A, Shazeer N, Parmar N, et al. Attention is all you need. Advances in neural information processing systems. 2017: 5998-6008.

[5] Xie S, Girshick R, Dollár P, et al. Aggregated residual transformations for deep neural networks. Proceedings of the IEEE conference on computer vision and pattern recognition. 2017: 1492-1500.

[6] Lin T Y, Dollár P, Girshick R, et al. Feature pyramid networks for object detection. Proceedings of the IEEE conference on computer vision and pattern recognition. 2017: 2117-2125.

想了解更多的AI技术干货，欢迎上华为云的AI专区，目前有AI编程Python等六大实战营供大家免费学习