# DiT: Self-supervised Pre-training for Document Image Transformer > Research note on DiT, a self-supervised pre-trained Document Image Transformer using BEiT-style masked image modeling on 42M unlabeled document images from IIT-CDIP. - Authors: Hunter Heidenreich - Published: 2026-03-13 - Paper: Preprint (arXiv) - Paper: Published (ACM MM 2022) - Code: DiT Code & Models - Model: dit-base (pre-trained) - Model: dit-large (pre-trained) - Model: dit-base-finetuned-rvlcdip - Model: dit-large-finetuned-rvlcdip - Source: Bevaya Labs, --- ## TL;DR DiT applies BEiT-style masked image modeling (MIM) to document images, pre-training a ViT backbone on 42 million unlabeled document pages from IIT-CDIP. A domain-specific dVAE tokenizer replaces the natural-image DALL-E tokenizer. The resulting model achieves top scores on document image classification, layout analysis, table detection, and text detection benchmarks at the time of publication. ## What kind of paper is this? **Dominant:** $\Psi\_{\text{Method}}$. The headline contribution is the self-supervised pre-training strategy (MIM with a document-domain dVAE) applied to a ViT backbone for document images. The paper's bulk consists of architectural choices, training recipes, and ablation-style comparisons across four downstream tasks. **Secondary:** $\Psi\_{\text{Resource}}$. The pre-trained DiT-Base and DiT-Large checkpoints are publicly released, and the domain-specific dVAE tokenizer trained on IIT-CDIP is a reusable asset. ## What is the motivation? Prior Document AI pre-training (LayoutLM, DocFormer, etc.) combines text, layout, and image modalities but still relies on supervised CNN backbones (typically ImageNet-pretrained) for the vision component. No large-scale human-labeled dataset analogous to ImageNet exists for document images, making supervised vision pre-training impractical for this domain. Weakly supervised datasets like PubLayNet and DocBank are drawn almost entirely from academic papers, creating a domain mismatch with real-world business documents (forms, invoices, receipts, reports). DiT addresses this by bringing self-supervised pre-training to the document vision backbone, using large-scale unlabeled document images that span diverse templates and formats. ## What is the novelty? DiT adapts the BEiT framework to document images with two key modifications: 1. **Domain-specific dVAE tokenizer.** Rather than using the DALL-E tokenizer (trained on 400M natural images), the authors retrain a discrete VAE on 42M document images from IIT-CDIP. The document tokenizer produces sharper reconstructions of text boundaries and line structures compared to the DALL-E tokenizer. The tokenizer uses a codebook of size 8,192 and a downsampling factor of 8, producing a $14 \times 14$ token map from a $112 \times 112$ input crop (aligned with the $14 \times 14$ patch grid from $224 \times 224$ input at $16 \times 16$ patch size). 2. **Masked Image Modeling (MIM) on document images.** The pre-training objective masks 40% of image patches (using blockwise masking as in BEiT) and requires the model to predict the corresponding discrete visual tokens from the dVAE. The training loss is cross-entropy over the 8,192-entry codebook: $$ \mathcal{L}\_{\text{MIM}} = -\sum\_{i \in \mathcal{M}} \log p(z\_i \mid \tilde{x}) $$ where $\mathcal{M}$ is the set of masked positions, $z\_i$ is the visual token from the dVAE, and $\tilde{x}$ is the corrupted input. The architecture itself is a standard ViT (no modifications to the Transformer blocks). For object detection tasks, DiT feeds into Mask R-CNN or Cascade R-CNN via a multi-scale FPN adapter that extracts features at four resolution levels from intermediate Transformer blocks (1/3, 1/2, 2/3, and the final block), using transposed convolutions for upsampling and max pooling for downsampling. ## What experiments were performed? DiT is evaluated on four Document AI benchmarks covering classification and detection: ### Document Image Classification (RVL-CDIP) 400K grayscale images across 16 document categories. Single-model accuracy comparison against CNN (ResNeXt-101) and ViT baselines (DeiT-B, BEiT-B, MAE-B), all at $224 \times 224$ resolution. ### Document Layout Analysis (PubLayNet) 360K document images with 5 element types (text, title, list, table, figure). Evaluated with mAP @ IoU [0.50:0.95] using both Mask R-CNN and Cascade R-CNN detection heads. ### Table Detection (ICDAR 2019 cTDaR) Combined archival and modern document subsets. Evaluated with weighted F1 across IoU thresholds (0.6, 0.7, 0.8, 0.9): $$ wF1 = \frac{0.6 \cdot F1\_{0.6} + 0.7 \cdot F1\_{0.7} + 0.8 \cdot F1\_{0.8} + 0.9 \cdot F1\_{0.9}}{0.6 + 0.7 + 0.8 + 0.9} $$ Notable: the authors apply adaptive binarization (OpenCV) to archival images before fine-tuning, improving performance on historical documents with color backgrounds. ### Text Detection (FUNSD) 199 scanned forms (31,485 words) with word-level bounding boxes. Evaluated with precision, recall, and F1 at IoU@0.5. Additionally tested with 1M synthetic document images for further pre-training. **Baselines:** ResNeXt-101-32x8d (CNN), DeiT-B, BEiT-B, MAE-B (all ImageNet-1K pre-trained ViTs). All baselines were re-run by the authors for fair comparison. ## What are the outcomes/conclusions? DiT reports the highest scores across all four benchmarks at time of publication: | Task | Previous Best | DiT Best | Config | | :--- | :--- | :--- | :--- | | RVL-CDIP (accuracy) | 91.11% (single) | **92.69%** | DiT-L | | PubLayNet (mAP) | 91.0 | **94.9** | DiT-L + Cascade R-CNN | | ICDAR 2019 cTDaR (wF1) | 94.23 | **96.55** | DiT-L + Cascade R-CNN | | FUNSD (F1 @ IoU 0.5) | 93.07 | **94.29** | DiT-L + synthetic data | Key observations: - The domain-specific dVAE tokenizer is important: visual reconstructions show much sharper text boundaries compared to the DALL-E tokenizer, particularly for fine-grained text and table borders. - DiT benefits more from Cascade R-CNN than baselines do, suggesting the richer representations from document-specific pre-training compound with stronger detection heads. - On the archival subset of ICDAR 2019, BEiT-B slightly outperforms DiT-B at lower IoU thresholds (0.6, 0.7). The authors attribute this to the DALL-E tokenizer being trained on color images (400M), while DiT's tokenizer was trained on grayscale IIT-CDIP images, disadvantaging it on colorful historical documents. - DiT-L consistently outperforms DiT-B, indicating that the self-supervised pre-training scales well with model capacity. **Limitations not discussed by the authors:** - All evaluations are on English-centric or Latin-script benchmarks. Performance on CJK, Arabic, or Indic scripts is unknown. - The $224 \times 224$ resolution for classification is low for dense document pages. Detection tasks use higher resolution but the pre-training is still at $224 \times 224$. - No ablation on pre-training data size or diversity; all experiments use the full 42M IIT-CDIP corpus. - The paper does not compare against concurrent multimodal approaches (LayoutLMv2, LayoutLMv3) that use text + layout + image jointly, though this is acknowledged as a different paradigm (DiT is vision-only by design). ## Reproducibility ### Models - **DiT-Base:** 12-layer ViT, 768 hidden size, 12 attention heads, 3072 FFN intermediate size. 87M parameters. - **DiT-Large:** 24-layer ViT, 1024 hidden size, 16 attention heads, 4096 FFN intermediate size. 304M parameters. - Pre-trained checkpoints are publicly available on HuggingFace (`microsoft/dit-base`, `microsoft/dit-large`). Fine-tuned checkpoints for RVL-CDIP classification are also on HuggingFace. Fine-tuned detection checkpoints (PubLayNet, ICDAR 2019 cTDaR) are available via the GitHub repo README. - Detection framework: Detectron2-based Mask R-CNN and Cascade R-CNN with multi-scale FPN adapter. ### Algorithms - **Pre-training:** 500K steps, batch size 2048, learning rate 1e-3, 10K warmup steps, weight decay 0.05. Adam optimizer ($\beta\_1 = 0.9$, $\beta\_2 = 0.999$). Stochastic depth rate 0.1, no dropout. Blockwise masking at 40%. - **dVAE training:** Learning rate 5e-4, minimum temperature 1e-10, 3 epochs on full IIT-CDIP. Codebook size 8192, encoder with three layers (each: 2D conv stride 2 + ResNet block), downsampling factor 8. Trained with MSE reconstruction loss + perplexity loss. Implementation based on the open-source [DALL-E PyTorch](https://github.com/lucidrains/DALLE-pytorch) codebase. - **Fine-tuning (RVL-CDIP):** 90 epochs, batch size 128, learning rate 1e-3, $224 \times 224$ with RandomResizedCrop. - **Fine-tuning (PubLayNet):** Batch size 16, learning rate 4e-4 (base) / 1e-4 (large). DETR-style data augmentation (random crop + multi-scale resize, shortest side 480-800px, longest side 1333px max). - **Fine-tuning (ICDAR 2019 cTDaR):** Batch size 16, learning rate 1e-4 (archival) / 5e-5 (modern). Same DETR augmentation. Adaptive binarization applied to archival images. - **Fine-tuning (FUNSD):** Batch size 16, learning rate 1e-4 (base) / 5e-5 (large). Anchor sizes [4, 8, 16, 32, 64] for word-level detection (vs. [32, 64, 128, 256, 512] for paragraph-level tasks). ### Data - **Pre-training corpus:** IIT-CDIP Test Collection 1.0 (42M document images after splitting multi-page documents). Grayscale, diverse business documents. - **dVAE training data:** Same IIT-CDIP corpus. - **Downstream datasets:** RVL-CDIP (400K), PubLayNet (360K), ICDAR 2019 cTDaR (1,200 train / 439 test), FUNSD (199 forms: 150 train / 49 test, 31,485 words total). - For FUNSD, an additional 1M synthetic document images were used in extended experiments. ### Evaluation - **RVL-CDIP:** Overall classification accuracy. - **PubLayNet:** Category-wise and overall mAP @ IoU [0.50:0.95]. - **ICDAR 2019 cTDaR:** Weighted F1 across IoU thresholds {0.6, 0.7, 0.8, 0.9}. - **FUNSD:** Precision, recall, F1 at IoU@0.5. - No error bars, confidence intervals, or multi-run statistics reported. - All baselines were re-run by the authors under identical conditions. ### Hardware - Not explicitly reported. Given 500K steps at batch size 2048, the pre-training likely required multiple V100 or A100 GPUs over several days, but the paper does not provide GPU count, total GPU-hours, or training time. ## BibTeX ```bibtex @inproceedings{li2022dit, title={DiT: Self-supervised Pre-training for Document Image Transformer}, author={Li, Junlong and Xu, Yiheng and Lv, Tengchao and Cui, Lei and Zhang, Cha and Wei, Furu}, booktitle={Proceedings of the 30th ACM International Conference on Multimedia}, pages={3530--3539}, year={2022}, doi={10.1145/3503161.3547911} } ```