Source-linked AI summary
DiT: Self-supervised Pre-training for Document Image Transformer
Junlong Li, Yiheng Xu, Tengchao Lv, Lei Cui, Cha Zhang, Furu Wei
TL;DR
Document AI lacks large-scale human-labeled document images for supervised pre-training, limiting the direct transfer of standard image-Transformer practices. DiT addresses this gap with self-supervised masked-image pre-training on unlabeled document images and a document-domain tokenizer. Across document classification, layout analysis, table detection, and OCR text detection, DiT outperforms strong baselines and achieves new state-of-the-art performance.
Problem
Large-scale supervised pre-training is impractical for document images because commonly used large human-labeled benchmarks like ImageNet do not exist in this domain.
Method
DiT is a self-supervised Document Image Transformer pre-trained on unlabeled document images using masked image modeling and a dVAE trained on 42 million document images.
Results
DiT outperforms strong baselines across document image classification, layout analysis, table detection, and text detection, achieving new state-of-the-art performance.
Takeaways & Limitations
The publicly available DiT models are intended to facilitate Document AI research across downstream vision-based tasks.
Abstract
from arXiv · showhide
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}.
1 INTRODUCTION
DiT addresses the lack of large-scale human-labeled document-image data by using self-supervised pre-training on unlabeled document images. The resulting backbone supports multiple Document AI tasks and achieves new state-of-the-art results.
- Document AI vision models commonly rely on supervised backbones trained with human-labeled samples, creating a challenge for real-world domain shifts and template variation.
- Large-scale supervised pre-training is impractical for document images because no commonly used human-labeled benchmark comparable to ImageNet exists.
- DiT uses self-supervised pre-training on large-scale unlabeled document images and adapts the visual tokenizer to the document domain.Images are resized to 224 × 224, split into 16 × 16 patches, and tokenized with a dVAE retrained on document images.
- DiT achieves new state-of-the-art results on the evaluated downstream tasks.
- DiT can serve as a backbone for document image classification, layout analysis, table detection, and text detection for OCR.
2 RELATED WORK
Related work establishes Image Transformers and self-supervised objectives as foundations for visual understanding, while Document AI remains constrained by limited human-labeled data. DiT extends this direction to vision-based document tasks.
- Image Transformers have advanced classification, detection, and segmentation through both supervised and self-supervised pre-training.
- Self-supervised image Transformer methods include autoregressive pixel prediction, self-distillation, Siamese learning, and masked visual-token prediction.
- Vision-based Document AI covers tasks such as OCR, document layout analysis, and document image classification.
- Document AI lacks large-scale human-labeled datasets, so existing approaches commonly use ConvNets pre-trained on ImageNet or COCO before task-specific training.
- DiT is presented as the first large-scale self-supervised pre-trained model for vision-based Document AI and can also support multimodal pre-training.
3 DOCUMENT IMAGE TRANSFORMER
DiT combines a vanilla Vision Transformer backbone with masked image modeling tailored to document images. The pre-trained representations are fine-tuned for classification and object-detection-based Document AI tasks.
- Architecture: DiT uses a vanilla Transformer backbone that converts non-overlapping document-image patches into position-embedded patch representations.
- Pre-training: Masked Image Modeling trains DiT to predict visual tokens from representations of masked image patches.
- Pre-training: DiT trains a document-domain dVAE on 42 million IIT-CDIP document images to avoid the domain mismatch of a natural-image tokenizer.
- Downstream tasks: The model is fine-tuned on RVL-CDIP, PubLayNet, ICDAR 2019 cTDaR, and FUNSD for image classification, layout analysis, table detection, and text detection.
- Downstream tasks: Image classification aggregates patch representations with average pooling before applying a linear classifier.
- Downstream tasks: Object detection uses Mask R-CNN or Cascade R-CNN with ViT-based backbones and resolution-modifying modules that adapt single-scale features to multi-scale FPNs.
4 EXPERIMENTS
The experiments evaluate DiT across document classification, layout analysis, table detection, and text detection, using established datasets and task-specific preprocessing or detection settings. DiT achieves strong results across these downstream tasks, with benefits reported for document-specific pre-training, model scale, and detection framework choice.
- 4.2 Settings: 42 million document images from IIT-CDIP are used for DiT pre-training, alongside a document-specific dVAE tokenizer trained on the same collection.The pre-training setup also uses blockwise masking with 40% of patches masked.
- 4.3 Results: RVL-CDIP evaluates image-only document classification at 224×224 resolution, and DiT-B outperforms selected single-model baselines.DiT-L obtains a score comparable to the previous state-of-the-art ensemble under the single-model setting.
- 4.3 Results: DiT-B is stronger than DeiT-B, BEiT-B, and MAE-B for PubLayNet layout analysis, while DiT-L achieves a higher mAP score.The reported improvement mainly comes from the List and Figure categories, and Cascade R-CNN further raises performance.
- 4.3 Results: DiT surpasses most baselines on ICDAR 2019 cTDaR table detection, while its archival performance trails BEiT because its tokenizer uses grayscale training images.Adaptive image binarization is applied to archival images after raw-image fine-tuning produced suboptimal performance.
- 4.3 Results: On FUNSD text detection, DiT achieves new state-of-the-art results, reaching an F1 of 0.9429 with DiT-L trained using one million synthetic document images.The task evaluates text bounding-box detection at IoU@0.5 using the Mask R-CNN framework with different backbones.
5 CONCLUSION AND FUTURE WORK
DiT is presented as a self-supervised foundation model for vision-based Document AI, pretrained on large-scale unlabeled document images and evaluated across several benchmarks. The authors report broad gains over strong baselines and plan larger-scale pretraining and multimodal integration.
- DiT is pretrained on large-scale unlabeled document images spanning varied templates and formats for general Document AI tasks.
- The model is evaluated on table detection, document layout analysis, document image classification, and text detection.
- DiT outperforms several strong baselines across the evaluated vision-based Document AI benchmarks and achieves new state-of-the-art performance.
- Future work will pretrain DiT on a much larger dataset and integrate it into multimodal pretraining for visually rich document understanding.