StrucTexTv2：端到端文档图像理解预训练框架

本文简要介绍ICLR 2023录用论文“StrucTexTv2: Masked Visual-Textual Prediction for Document Image Pre-training”的主要工作。针对当前主流多模态文档理解预训练模型需要同时输入文档图像和OCR结果，导致欠缺端到端的表达能力且推理效率偏低等问题，论文提出了一种全新的端到端文档图像多模态表征学习预训练框架StrucTexTv2。该框架设计了一种基于词粒度图像区域掩码、多模态自监督预训练任务（MIM+MLM），仅需要图像单模态输入，使得编码器网络能在大规模无标注文档图像上充分学习视觉和语言联合特征表达，并在多个下游任务的公开基准上取得SOTA效果。

一、研究背景

视觉富文档理解技术例如文档分类、版式分析、表单理解、OCR以及信息提取，逐渐成为文档智能领域一个热门研究课题。为了有效处理这些任务，前沿的方法大多利用视觉和文本线索，将图像、文本、布局等信息输入到参数网络，并基于大规模数据上的自监督预训练挖掘出文档的多模态特征。由于视觉和语言之间的模态差异较大，如图1所示，主流的文档理解预训练方法大致可分为两类：a）掩码语言建模（Masked Language Modeling）[9]，对输入的掩码文本Token进行语言建模，运行时文本的获取依赖于OCR引擎，整个系统的性能提升需要对OCR引擎和文档理解模型两个部件进行同步优化；b）掩码图像建模（Masked Image Modeling）[10]，对输入的掩码图像块区进行像素重建，此类方法倾向应用于图像分类和版式分析等任务上，对文档强语义理解能力欠佳。针对上述两种预训练方案呈现的瓶颈，本文提出了StrucTexTv2：c）统一图像重建与语言建模方式，在大规模文档图像上学习视觉和语言联合特征表达。

图1 主流文档图像理解预训练框架比较

二、方法原理简述

图2 整体框架图

图2描绘了StrucTexTv2的整体框架，主要包含编码器网络和预训练任务分支两部分。编码器网络，主要通过FPN结构串联CNN组件和Transformer组件构成；预训练分支则包含了掩码语言建模（MLM）和掩码图像建模（MIM）双预训练任务头。

2.1 编码器网络

StrucTexTv2采用CNN和Transformer的串联编码器来提取文档图像的视觉和语义特征。文档图像首先经过ResNet网络以获取1/4到1/32的四个不同尺度的特征图。随后采用一个标准的Transformer网络接收最小尺度的特征图并加上1D位置编码向量，提取出包含全局上下文的语义特征。该特征被重新转化为2D形态后，与CNN的其余三个尺度特征图通过FPN[6]融合成4倍下采样的特征图，作为整图的多模态特征表示。

2.2 预训练策略

为了统一建模MLM和MIM两种模态预训练方式，论文提出了一种基于词粒度图像区域掩码预测方式来学习视觉和语言联合特征表达。首先，随机筛选30%的词粒度OCR预测结果（仅在预训练阶段使用），根据OCR的位置信息直接在原图对应位置像素进行掩码操作（比如填充0值）。接着，掩码后的文档图像直接送入编码器网络去获得整图的多模态特征表示。最后，再次根据选中的OCR位置信息，采用ROIAlign[11]操作去获得每个掩码区域的多模态ROI特征。

掩码语言建模：借鉴于BERT[9]构建的掩码语言模型思路，语言建模分支使用一个2层的MLP将词区域的ROI特征映射到预定义的词表类别上，使用Cross Entropy Loss监督。同时为了避免使用词表对文本序列进行标记化时单个词组被拆分成多个子词导致的一对多匹配问题，论文使用分词后每个单词的首个子词作为分类标签。此设计带来的优势是：StrucTexTv2的语言建模无需文本作为输入。

掩码图像建模：考虑到基于图像Patch的掩码重建在文档预训练中展现出一定的潜力，但Patch粒度的特征表示难以恢复文本细节。因此，论文将词粒度掩码同时用作图像重建，即预测被掩码区域的原始像素值。词区域的ROI特征首先通过一个全局池化操作被压缩成特征向量。其次，为了提升图像重建的视觉效果，论文将通过语言建模后的概率特征与池化特征进行拼接，为图像建模引入“Content”信息，使得图像预训练专注于复原文本区域的“Style”部分。图像建模分支由3个全卷积 Block构成。每个Block包含一个Kernel=2×2，Stride=4的反卷积层，一个Kernel=1×1，以及两个Kernel=3×1卷积层。最后，每个单词的池化向量被映射成一个大小为64×64×3的图像，并逐像素与原本的图像区域做MSE Loss。

论文提供了Small和Large两种参数规格的模型，并在IIT-CDIP数据集上使用百度通用高精OCR的文字识别结果预训练编码网络。

三、实验结果

论文在四个基准数据集上测试模型对文档理解的能力，在五个下游任务上使用不同的Head进行Fine-tune并给出实验结论。表1给出模型在RVL-CDIP[13]验证文档图像分类的效果。同比基于图像单模态输入的方法DiT[4]，StrucTexTv2以更少的参数量取得了更优的分类精度。

表1 RVL-CDIP数据集上文档图像分类的实验结果

如表2和表3所示，论文结合预训练模型和Cascade R-CNN[1]框架fine-tune去检测文档中的版式元素以及表格结构，在PubLaynet[8]以及WWW[12]数据集上取得了当前的最好性能。

表2 PubLaynet数据集上版式分析的检测结果

表3 WWW数据集上表格结构识别的性能对比

在表4中，论文同时在FUNSD[3]数据集上进行了端到端OCR和信息提取两项实验，在基准测试中都取得了同期最优的效果。对比如StrucTexTv1[5]和LayoutLMv3[2]等OCR+文档理解的两阶段方法，证明了提出方法端到端优化的优越性。

表4 FUNSD数据集上端到端OCR以及信息抽取实验

接下来，论文对比了SwinTransformer[7]、ViT[10]以及StrucTexTv2的编码网络。从表5对比结果来看，论文提出CNN+Transformer的串联结构更有效地支持预训练任务。同时，论文给出了不同预训练配置的模型在文档图像分类和版式分析的性能增益，对两种模态预训练进行了有效性验证。

表5 预训练任务以及编码器结构的消融实验

同时，论文中评估了模型在预测时的耗时和显存开销。表6中给出了两种OCR引擎带来的开销以及并与现阶段最优的多模态方法LayoutLMv3进行了比较。

表6 与两阶段的方法LayoutLMv3的资源开销对比

最后，论文评估了表7所示在图像重建预训练中使用不同的掩码方式对下游任务的影响。在RVL-CDIP和PubLaynet两个数据集上，基于词粒度掩码的策略可以获取到更有效的视觉语义特征，确保更好的性能。

表7 预训练任务以及编码器结构的消融实验

总结及讨论

论文出的StructTexTv2模型用于端到端学习文档图像的视觉和语言联合特征表达，图像单模态输入条件下即可实现高效的文档理解。论文提出的预训练方法基于词粒度的图像掩码，能同时预测相应的视觉和文本内容，此外，所提出的编码器网络能够更有效地挖掘大规模文档图像信息。实验表明，StructTexTv2在模型大小和推理效率方面对比之前的方法都有显著提高。更多的方法原理介绍和实验细节请参考论文原文。

审核编辑 ：李倩