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

本文简要介绍ACM MM 2022录用论文“DiT: Self-supervised Pre-training for Document Image Transformer”[1]的主要工作。该论文是2022年微软亚研院发表的LayoutLM V3[2]的前身工作，主要解决了文档领域中标注数据稀少和以视觉为中心的文档智能任务骨干网络的预训练问题。

一、研究背景

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

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

二、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 模型微调

Fig 3 Illustration of applying DiT as the backbone net-work in different detection frameworks.

模型预训练完成后，需针对下游任务进行微小的结构改动，针对分类任务，输入经过平均池化和线性层进行分类。针对检测任务，如Fig 3所示，在ViT的特定层进行下采样或上采样，然后输入到FPN和后续的检测框架中。

三、主要实验结果及可视化效果

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

从TABLE 1、TABLE 2、TABLE 3、TABLE 4来看，文章所提方法在各种下游任务中取得了state-of-the-art的结果，验证了该方法在文档领域的有效性。Fig 4中展示了重新训练的离散变分编码器的可视化输出，结果显示本文中的离散变分编码器效果更好。

四、总结及讨论

本文设计了一个利用大量无标签文档图像预训练ViT的自监督方法，该方法的核心是利用离散变分编码器对图像Patch进行Token化，再使用NLP领域的掩码重建任务进行预训练。从实验结果可以看出，该方法在多个下游任务的有效性，探索了自监督任务在文档领域的可能性。

五、相关资源

DiT论文地址：https://arxiv.org/abs/2203.02378

代码开源地址：https://aka.ms/msdit

参考文献

[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 16x16 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

原文作者: Junlong Li, Yiheng Xu, Tengchao Lv, Lei Cui, Cha Zhang, Furu Wei

撰稿：简晟

编排：高学

审校：殷飞

发布：金连文

免责声明：（1）本文仅代表撰稿者观点，撰稿者不一定是原文作者，其个人理解及总结不一定准确及全面，论文完整思想及论点应以原论文为准。（2）本文观点不代表本公众号立场。

往期精彩内容回顾