[论文解读] DiT: Self-supervised Pre-training for Document Image Transformer

Junlong Li, Yiheng Xu+4 more|arXiv (Cornell University)|Mar 4, 2022

Handwritten Text Recognition Techniques被引用 9

一句话总结

DiT 提出了一种自监督预训练的文档图像 Transformer 模型，通过利用大规模未标注文档图像学习通用视觉表征，无需人工标注。通过将掩码图像建模与面向文档的离散 VAE 相结合，DiT 在多个文档 AI 任务中实现了最先进性能，包括文档图像分类（92.69% 的 top-1 准确率）、版面分析（94.9% 的 F1）、表格检测（96.55% 的 wF1）以及文本检测（94.29% 的 F1）。

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 extbf{DiT}, a self-supervised pre-trained extbf{D}ocument extbf{I}mage extbf{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 $ ightarrow$ 92.69), document layout analysis (91.0 $ ightarrow$ 94.9), table detection (94.23 $ ightarrow$ 96.55) and text detection for OCR (93.07 $ ightarrow$ 94.29). The code and pre-trained models are publicly available at \url{https://aka.ms/msdit}.

研究动机与目标

• 解决文档 AI 中用于预训练视觉主干网络的大规模人工标注文档图像数据集缺乏的问题。
• 克服因模板和格式不匹配导致的真实世界文档 AI 应用中出现的领域偏移与性能下降问题。
• 开发一种面向文档图像的自监督预训练框架，使其在多样化文档布局与格式中具备泛化能力。
• 在不依赖监督标注的情况下，提升基于视觉的文档 AI 任务的下游性能。
• 建立一个支持计算机视觉与自然语言处理任务的统一 Transformer 架构文档 AI 基础模型。

提出的方法

• 在大规模未标注文档图像上使用掩码图像建模（MIM）预训练视觉 Transformer 主干网络。
• 在灰度文档图像上训练一个领域特定的离散 VAE（dVAE），以生成比通用模型（如 DALL-E）更相关的视觉 token。
• 通过掩码图像块的方式破坏输入文档图像，并训练 Transformer 从 dVAE 重建原始视觉 token。
• 使用标准图像 Transformer 架构，结合学习的位置嵌入和多头自注意力机制，建模文档布局中的长距离依赖关系。
• 在下游文档 AI 任务上使用标准分类头和检测头对预训练的 DiT 模型进行微调。
• 通过在额外 100 万张合成文档图像上进行微调，进一步提升性能，增强零样本泛化能力。

实验结果

研究问题

• RQ1 在大规模未标注文档图像上进行自监督预训练，是否能在无需人工标注的情况下提升文档 AI 任务的泛化能力？
• RQ2 面向文档的 dVAE 与通用图像 tokenizers（如 DALL-E）相比，在学习文档理解的有意义视觉表征方面表现如何？
• RQ3 与监督和自监督基线相比，DiT 在文档图像分类、版面分析、表格检测和文本检测等下游任务上的性能提升程度如何？
• RQ4 与在有限或领域特定数据集上预训练的模型相比，DiT 在多样化文档格式和布局上的泛化能力是否更强？
• RQ5 DiT 是否可作为多模态文档 AI 系统的强大基础模型，包括基于版面的模型（如 LayoutLM）？”

主要发现

• DiT 在 RVL-CDIP 文档图像分类任务上实现了 92.69% 的 top-1 准确率，创下新 SOTA，相比之前方法的 91.11% 显著提升。
• 在 PubLayNet 的文档版面分析任务中，DiT 达到 94.9% 的 F1，显著优于之前的 SOTA 方法（91.0%）。
• 在 ICDAR 2019 cTDaR 的表格检测任务中，DiT-L 达到 96.55% 的 wF1 分数，超过之前最佳结果（94.23%）超过 2.3 个百分点。
• 在 FUNSD 的文本检测任务中，DiT-L 达到 94.29% 的 F1（IoU@0.5），优于商业 OCR 引擎基线以及所有其他 ViT 和 CNN 主干网络。
• 通过在包含 100 万张文档图像的合成数据集上微调 DiT，性能进一步提升，在 FUNSD 上达到 94.29% F1，在表格检测任务上达到 96.55% wF1。
• 该模型在少样本泛化方面表现强劲，尤其在 cTDaR 的档案子集等低资源设置下，尽管 dVAE 存在颜色限制，仍优于大多数基线模型。