[ACM MM 2022] DiT: 基于Transformer的文档图像自监督预训练方法

近年来自监督预训练技术已在文档智能领域进行了许多的实践，大多数技术是将图片、文本、布局结构信息一起输入统一的Transformer架构中。在这些技术中，经典的流程是先经过一个视觉模型提取额外文档图片信息，例如OCR引擎或版面分析模型，这些模型通常依赖于有标注数据训练的视觉骨干网络。已有的工作已经证明一些视觉模型在实际应用中的性能经常受到域迁移、数据分布不一致等问题的影响。而且现有的文档有标注数据集稀少、样式单一，训练出来的骨干网络并非最适用于文档任务。因此，有必要研究如何利用自监督预训练技术训练一个专用于文档智能领域的骨干网络。本文针对上述问题，利用离散变分编码器和NLP领域的常用预训练方式实现了文档图像的预训练。

Fig 1 Visually-rich business documents with different layouts and formats for pre-training DiT

2.1总体结构

Fig 2 Overall Architecture of DiT

Fig 2是DiT的整体结构。DiT使用ViT[3]作为预训练的骨干网络，模型的输入是图像Patch化后的Embedding特征向量，Patch的数量和离散变分编码器的下采样比例有关。输入经过ViT后输出到线性层进行图像分类，分类层的大小是8192。预训练任务和NLP领域的完型填空任务一致，先对输入的Patch随机掩膜，在模型输出处预测被遮盖的Patch对应的Token，Token由Fig 2 中左侧的离散变分编码器生成，作为每个Patch的Label，预训练过程使用CE Loss监督。

2.2 离散变分编码器dVAE

离散变分编码器作为Image Tokenizer，将输入的Patch Token化，来源于论文DALL-E[4]，在预训练任务开始前需要额外训练。本文使用数据集IIT-CDIP[5]重新训练了DALL-E中的离散变分编码器以适用于文档任务。在预训练任务中只使用到编码器的部分，解码器不参与预训练，编码器将输入图片下采样到原来的1/8，例如输入尺度为112*112，那编码后的Token Map为14*14，此时的Map大小，应与ViT输入Patch数保持一致。

2.3 模型微调

Table 1.  Document Image Classification accuracy (%) on RVL-CDIP, where all the models use the pure image information(w/o text information) with the 224×224 resolution.

Table 2. Document Layout Analysis mAP @ IOU [0.50:0.95] on PubLayNet validation set. ResNext-101-32×8d is shortened as ResNext and Cascade as C.

Table 3. Table detection accuracy (F1) on ICDAR 2019 cTDaR

Table 4. Text detection accuracy ([email protected]) on FUNSD Task#1, where Mask R-CNN is used with different backbones(ResNeXt, DeiT, BEiT, MAE and DiT). “+syn” denotes that DiT is trained with a synthetic dataset including 1M document images, then fine-tuned with the FUNSD training data.

Fig 4 Document image reconstruction with different tokenizers. From left to right: the original document image, image reconstruction using the self-trained dVAE tokenizer, image reconstruction using the DALL-E tokenizer

[1] Junlong Li, Yiheng Xu, Tengchao Lv, Lei Cui, Cha Zhang, and Furu Wei. 2022. DiT: Self-supervised Pre-training for Document Image Transformer. In Proceedings of the 30th ACM International Conference on Multimedia (MM ’22). Association for Computing Machinery, New York, NY, USA, 3530–3539. https://doi.org/10.1145/3503161.3547911

[2] Yupan Huang, Tengchao Lv, Lei Cui, Yutong Lu, and Furu Wei. 2022. LayoutLMv3: Pre-training for Document AI with Unified Text and Image Masking. In Proceedings of the 30th ACM International Conference on Multimedia (MM ’22). Association for Computing Machinery, New York, NY, USA, 4083–4091. https://doi.org/10.1145/3503161.3548112.

[3] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. 2021. An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale. ICLR (2021).

[4] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Rad-ford, Mark Chen, and Ilya Sutskever. 2021. Zero-Shot Text-to-Image Generation.arXiv:2102.12092 [cs.CV]

[5] D. Lewis, G. Agam, S. Argamon, O. Frieder, D. Grossman, and J. Heard. 2006.Building a Test Collection for Complex Document Information Processing. InProceedings of the 29th Annual International ACM SIGIR Conference on Researchand Development in Information Retrieval (Seattle, Washington, USA) (SIGIR ’06).ACM, New York, NY, USA, 665–666. https://doi.org/10.1145/1148170.114830

• [IJCAI 2022] C3-STISR: 基于三重线索引导的场景文本图像超分辨率方法（有源码）
  

• [CVPR 2022]基于语法感知网络的手写数学公式识别（已开源）
  

• [ACM MM 2022] 解耦检测与识别：单阶段自依赖场景文本识别器
  

• [ECCV 2022] CoMER: 基于Transformer与覆盖注意力机制建模的手写数学公式识别（已开源）
  

• [ECCV 2022] 场景文字端到端识别中的全局到局部注意
  

• [ECCV2022] MGP-STR:一种基于视觉Transformer的多粒度文字识别方法(已开源)
  

• [IEEE TMM 2022] |手写汉字纠错的树结构分析网络
  

• [SIGGRAPH 2022] 利用真实数据来提升文档图像矫正性能（有源码）
  

• [IEEE TIP 2022] | 基于EM算法的混合监督场景文本检测
  

• [ACM 2022] 基于判别式和生成式的自监督文本图像识别方法

欢迎加入中国图象图形学学会!（附入会攻略）

扫码关注，获取最新OCR资讯

原文始发于微信公众号（CSIG文档图像分析与识别专委会）：[ACM MM 2022] DiT: 基于Transformer的文档图像自监督预训练方法