DiT: Self-supervised Pre-training for Document Image Transformer

TL;DR

DiT tackles the lack of large labeled document datasets by performing self-supervised pre-training on 42M unlabeled document images using a document-specific dVAE tokenizer and Masked Image Modeling. The model adopts a ViT-based Document Image Transformer backbone, pre-trained on IIT-CDIP, and fine-tuned for document image classification, layout analysis, table detection, and text detection. It achieves new state-of-the-art results across RVL-CDIP, PubLayNet, ICDAR cTDaR, and FUNSD, outperforming both supervised and other self-supervised baselines. The work demonstrates the value of domain-specific self-supervised pre-training for document understanding and provides open-source models and code for the community.

Abstract

Image Transformer has recently achieved significant progress for natural image understanding, either using supervised (ViT, DeiT, etc.) or self-supervised (BEiT, MAE, etc.) pre-training techniques. In this paper, we propose \textbf{DiT}, a self-supervised pre-trained \textbf{D}ocument \textbf{I}mage \textbf{T}ransformer model using large-scale unlabeled text images for Document AI tasks, which is essential since no supervised counterparts ever exist due to the lack of human-labeled document images. We leverage DiT as the backbone network in a variety of vision-based Document AI tasks, including document image classification, document layout analysis, table detection as well as text detection for OCR. Experiment results have illustrated that the self-supervised pre-trained DiT model achieves new state-of-the-art results on these downstream tasks, e.g. document image classification (91.11 $\rightarrow$ 92.69), document layout analysis (91.0 $\rightarrow$ 94.9), table detection (94.23 $\rightarrow$ 96.55) and text detection for OCR (93.07 $\rightarrow$ 94.29). The code and pre-trained models are publicly available at \url{https://aka.ms/msdit}.

Figures (5)

• Figure 1: Visually-rich business documents with different layouts and formats for pre-training DiT.
• Figure 2: The model architecture of DiT with MIM pre-training.
• Figure 3: Illustration of applying DiT as the backbone network in different detection frameworks.
• Figure 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.
• Figure 5: An example of pre-processing with adaptive image binarization on the ICDAR 2019 cTDaR archival subset.