摘要

图像transformer（Image Transformer）最近在自然图像理解方面取得了显著进展， 无论是使用监督（ViT、DeiT等）还是自监督（BEiT、MAE等）预训练技术。在本文中，我们提出了DiT，一种自监督预训练的文档图像transformer模型，利用大规模未标记的文本图像进行文档AI任务的训练。由于缺乏人工标注的文档图像，迄今为止并不存在任何监督的对应模型，因此该方法至关重要。我们将DiT作为骨干网络，应用于各种基于视觉的文档AI任务，包括文档图像分类、文档布局分析、表格检测以及OCR文本检测。实验结果表明，自监督预训练的DiT模型在这些下游任务中取得了新的最先进的结果，例如文档图像分类（91.11 → 92.69）、文档布局分析（91.0 → 94.9）、表格检测（94.23 → 96.55）和OCR文本检测（93.07 → 94.29）。代码和预训练模型可以在https://aka.ms/msdit获取。

CCS概念

• 计算方法 → 计算机视觉

关键词

文档图像transformer、自监督预训练、文档图像分类、文档布局分析、表格检测、文本检测、OCR

以下是第一部分“Introduction”的翻译：

1 引言

自监督预训练技术在过去几年已经成为文档AI的事实标准做法[10]，其中图像、文本和布局信息通常使用统一的transformer架构[2, 19, 21, 25, 28, 32, 33, 38, 41–44, 48]联合训练。在所有这些方法中，预训练文档AI模型的典型流程通常从基于视觉的理解开始，例如光学字符识别（OCR）或文档布局分析，这仍然在很大程度上依赖于具有人工标注训练样本的监督计算机视觉骨干模型。尽管在基准数据集上取得了良好的结果，但这些视觉模型在实际应用中往往面临由于领域转移和训练数据中的模板/格式不匹配而导致的性能差距。这种准确性回退[26, 46]对预训练模型及下游任务也产生了重要影响。因此，研究如何利用自监督预训练来增强文档图像理解的骨干网络是不可避免的，这能够更好地促进不同领域的通用文档AI模型。

图像transformer[3, 9, 12, 13, 17, 31, 36, 47]最近在自然图像理解方面取得了巨大成功，包括分类、检测和分割任务，无论是在ImageNet上的监督预训练，还是自监督预训练。与CNN基于预训练的模型相比，预训练的图像transformer模型在相似的参数规模下能够实现可比甚至更好的性能。然而，对于文档图像理解，目前没有像ImageNet这样的常用大规模人工标注基准，这使得大规模监督预训练变得不切实际。尽管弱监督方法已被用于创建文档AI基准数据集[26, 27, 45, 46]，这些数据集的领域通常来自共享相似模板和格式的学术论文，这些与实际世界中的文档（如表格、发票/收据、报告等）存在差异，见图1所示。这可能导致通用文档AI问题的结果不尽如人意。因此，使用来自通用领域的大规模未标记数据对文档图像骨干网络进行预训练至关重要，这能够支持多种文档AI任务。

为此，我们提出了DiT，一种自监督预训练的文档图像transformer模型，用于通用文档AI任务，且不依赖于任何人工标注的文档图像。受最近提出的BEiT模型[3]的启发，我们采用了类似的预训练策略，使用文档图像作为输入。输入的文本图像首先被调整为224 × 224的大小，然后图像被切分成一系列16 × 16的图块，这些图块作为输入传递给图像transformer。与BEiT模型不同，后者的视觉token来源于DALL-E中的离散VAE[34]，我们重新训练了离散VAE（dVAE）模型，使用大规模的文档图像，以便生成的视觉token更加与文档AI任务相关。预训练目标是基于腐蚀的输入文档图像，使用BEiT中的掩蔽图像建模（Masked Image Modeling, MIM）来恢复dVAE的视觉token。通过这种方式，DiT模型不依赖于任何人工标注的文档图像，而是仅利用大规模未标记数据来学习文档图像内各个图块之间的全局关系。

我们在四个公开的文档AI基准数据集上评估了预训练的DiT模型，包括用于文档图像分类的RVL-CDIP数据集[16]、用于文档布局分析的PubLayNet数据集[46]、用于表格检测的ICDAR 2019 cTDaR数据集[15]，以及用于OCR文本检测的FUNSD数据集[22]。实验结果表明，预训练的DiT模型超越了现有的监督和自监督预训练模型，并在这些任务上达到了新的最先进水平。

本文的贡献总结如下：

1. 我们提出了DiT，一种自监督预训练的文档图像transformer模型，可以利用大规模未标记的文档图像进行预训练。
2. 我们将预训练的DiT模型作为骨干网络，应用于多种文档AI任务，包括文档图像分类、文档布局分析、表格检测以及OCR文本检测，并取得了新的最先进结果。
3. 代码和预训练模型公开可用，网址为 https://aka.ms/msdit。

2 相关工作

图像transformer（Image Transformer）最近在计算机视觉问题上取得了显著进展，包括分类、目标检测和分割等。[12]首次将标准transformer直接应用于图像，修改最少。他们将图像切分为16 × 16的图块，并将这些图块的线性嵌入序列作为输入提供给名为ViT的transformer。ViT模型在图像分类任务中进行了监督训练，并超越了ResNet基准。[36]提出了数据高效的图像transformer和通过注意力蒸馏的DeiT，它仅依赖于ImageNet数据集进行监督预训练，并且与ViT相比，取得了最先进的结果。[31]提出了一种层次化的transformer，其表示通过移动窗口计算。移动窗口方案通过将自注意力计算限制在不重叠的局部窗口中，带来了效率，同时也允许跨窗口连接。

除了监督预训练的模型，[8]训练了一种序列transformer，称为iGPT，用于自回归地预测像素，而不考虑2D输入结构的知识，这是图像transformer自监督预训练的首次尝试。此后，图像transformer的自监督预训练成为计算机视觉领域的热门话题。[7]提出了DINO，它使用自蒸馏技术对图像transformer进行预训练，并且无需标签。[9]提出了基于Siamese网络的MoCov3，用于自监督学习。最近，[3]采用了类似BERT的预训练策略，首先将原始图像转换为视觉token，然后随机掩蔽一些图像块并将它们输入骨干transformer。类似于掩蔽语言建模，他们提出了一个掩蔽图像建模任务作为预训练目标，并取得了最先进的性能。[47]提出了一个自监督框架iBOT，它可以通过在线tokenizer执行掩蔽预测。在线tokenizer可以与MIM目标共同学习，并且不需要预先训练tokenizer的多阶段流水线。

基于视觉的文档AI 通常指利用计算机视觉模型进行文档分析的任务，例如OCR（光学字符识别）、文档布局分析和文档图像分类。由于该领域缺乏大规模人工标注的数据集，现有的方法通常基于使用ImageNet/COCO数据集预训练的ConvNets模型。然后，这些模型会继续用任务特定的标注样本进行训练。据我们所知，预训练的DiT模型是第一个用于基于视觉的文档AI任务的大规模自监督预训练模型。同时，它也可以进一步用于文档AI的多模态预训练。

3 文档图像transformer（DiT）

在本节中，我们首先介绍DiT的架构和预训练过程。接着，我们描述DiT模型在不同下游任务中的应用。

3.1 模型架构

沿用ViT [12]，我们使用标准的Transformer架构 [37] 作为DiT的骨干网络。我们将文档图像分割成不重叠的图块，并获得一系列图块嵌入。加上1D位置嵌入后，这些图像图块被传入一系列多头注意力机制的Transformer块。最后，我们将Transformer编码器的输出作为图像图块的表示，如图2所示。

3.2 预训练

受到BEiT [3] 的启发，我们使用掩蔽图像建模（MIM）作为我们的预训练目标。在此过程中，图像分别作为图像图块和视觉tokens进行表示。在预训练过程中，DiT接受图像图块作为输入，并通过输出表示来预测视觉tokens。

像自然语言中的文本token一样，图像也可以通过图像tokenizer表示为一系列离散的token。BEiT使用来自DALL-E [34] 的离散变分自编码器（dVAE）作为图像tokenizer，该模型在包括4亿张图像的大型数据集上进行训练。然而，天然图像与文档图像之间存在领域不匹配，这使得DALL-E的tokenizer不适用于文档图像。因此，为了获得更适合文档图像领域的离散视觉token，我们在IIT-CDIP [24] 数据集上训练了一个dVAE，该数据集包含4200万张文档图像。

为了有效地预训练DiT模型，我们在给定图像图块序列的情况下，随机掩蔽输入的子集，并用特殊的token [MASK]表示。DiT编码器通过线性投影嵌入掩蔽图块序列，并添加位置嵌入，接着通过一系列Transformer块进行上下文化处理。模型需要预测来自掩蔽位置的视觉token索引。与预测原始像素不同，掩蔽图像建模任务要求模型预测通过图像tokenizer获得的离散视觉token。

3.3 微调

我们在四个文档AI基准测试上微调了模型，包括用于文档图像分类的RVL-CDIP数据集、用于文档布局分析的PubLayNet数据集、用于表格检测的ICDAR 2019 cTDaR数据集，以及用于文本检测的FUNSD数据集。这些基准数据集可以形式化为两种常见任务：图像分类和目标检测。

图像分类：
对于图像分类，我们使用平均池化来聚合图像图块的表示。接下来，我们将全局表示传入一个简单的线性分类器。

目标检测：
对于目标检测，如图3所示，我们利用Mask R-CNN [18] 和Cascade R-CNN [5] 作为检测框架，并使用基于ViT的模型作为骨干网络。我们的代码基于Detectron2 [39] 实现。根据[14, 29]，我们在四个不同的Transformer块上使用分辨率修改模块，将单尺度ViT适应为多尺度FPN。设𝑑为总块数，1𝑑/3块通过2×2转置卷积模块进行4倍上采样，步幅为2。对于1𝑑/2块的输出，我们使用一个单步幅2的2×2转置卷积进行2倍上采样。2𝑑/3块的输出不做额外操作。最后，3𝑑/3块的输出通过步幅为2的2×2最大池化进行2倍下采样。

4 实验

4.1 任务

我们在此简要介绍第3.3节中提到的数据集。

RVL-CDIP：RVL-CDIP [16] 数据集包含400,000张灰度图像，分为16个类别，每个类别有25,000张图像。数据集包括320,000张训练图像，40,000张验证图像和40,000张测试图像。16个类别包括{信件、表格、电子邮件、手写、广告、科学报告、科学出版物、规格、文件夹、新闻文章、预算、发票、演示文稿、问卷、简历、备忘录}。评估指标是整体分类准确率。

PubLayNet：PubLayNet [46] 是一个大规模的文档布局分析数据集。该数据集通过自动解析PubMed XML文件构建了超过360,000张文档图像。生成的注释涵盖了典型的文档布局元素，如文本、标题、列表、图形和表格。模型需要检测分配元素的区域。我们使用按类别的和总体的平均精度（MAP）@交并比（IoU）[0.50:0.95]作为评估指标。

ICDAR 2019 cTDaR：cTDaR数据集 [15] 包含两个任务，包括表格检测和表格结构识别。在本文中，我们专注于Track A，其中提供了带有一个或多个表格注释的文档图像。该数据集包含两个子集，一个用于档案文档，另一个用于现代文档。档案子集包括600张训练图像和199张测试图像，展示了各种各样的表格，包括手绘的会计账本、证券交易清单、火车时刻表、生产普查等。现代子集包含600张训练图像和240张测试图像，包含不同种类的PDF文件，如科学期刊、表单、财务报表等。数据集包含中英文文档，格式多样，包括扫描文档图像和数字出生格式。评估该任务的指标是精度、召回率和F1分数，这些值是根据模型排名输出相对于不同交并比（IoU）阈值计算的。我们分别使用IoU阈值为0.6、0.7、0.8和0.9计算这些值，并将它们合并为最终的加权F1分数：
$$wF1=\frac{0.6F{1}_{0.6}+0.7F{1}_{0.7}+0.8F{1}_{0.8}+0.9F{1}_{0.9}}{0.6+0.7+0.8+0.9}$$
这个任务进一步要求模型将现代子集和档案子集合并为一个整体，以获得最终的评估结果。

FUNSD。FUNSD [22] 是一个带有噪声的扫描文档数据集，标注了三个任务：文本检测、光学字符识别（OCR）文本识别和表单理解。本文聚焦于 FUNSD 的任务 #1，旨在检测扫描表单文档中的文本边界框。FUNSD 包含 199 个完全标注的表单，包含 31,485 个词，其中训练集包含 150 个表单，测试集包含 49 个表单。评估指标为在 IoU@0.5 下的精度、召回率和 F1 分数。

4.2 设置

预训练设置。我们在 IIT-CDIP 测试集 1.0 [24] 上对 DiT 进行了预训练。我们通过将多页文档拆分为单页文档来预处理数据集，获得了 4200 万个文档图像。我们还在训练过程中引入了随机裁剪（random resized cropping）来增强训练数据。我们使用与 ViT base 相同的架构训练 DiT-B 模型：一个 12 层的 Transformer，隐藏层大小为 768，注意力头数为 12。前馈网络的中间大小为 3072。我们还训练了一个更大的版本 DiT-L，包含 24 层，隐藏层大小为 1024，注意力头数为 16，前馈网络的中间大小为 4096。

dVAE 分词器。BEiT 借用了由 DALL-E 训练的图像分词器，但该分词器与文档图像数据并不匹配。因此，我们充分利用了 IIT-CDIP 数据集中的 4200 万个文档图像，训练了一个文档 dVAE 图像分词器，以获得视觉标记。与 DALL-E 图像分词器类似，文档图像分词器具有 8192 的字典维度，图像编码器包含三层。每一层由一个步长为 2 的二维卷积和一个 ResNet 块组成。因此，分词器最终的降采样因子为 8。在这种情况下，给定一个 112×112 的图像，最终得到一个与 14×14 输入补丁对齐的 14×14 离散标记图。

我们从开源的 DALL-E 实现中实现了我们的 dVAE 代码库，并使用包含 4200 万文档图像的整个 IIT-CDIP 数据集训练了 dVAE 模型。新的 dVAE 分词器通过结合 MSE 损失（用于重建输入图像）和困惑度损失（用于增加量化代码本表示的使用）进行训练。输入图像的大小为 224×224，我们以 5e-4 的学习率和 1e-10 的最小温度训练分词器，训练 3 个周期。我们通过从下游任务中重建文档图像样本，比较我们的 dVAE 分词器与原始的 DALL-E 分词器，结果如图 4 所示。我们从文档布局分析数据集 PubLayNet 和表格检测数据集 ICDAR 2019 cTDaR 中采样图像。经过 DALL-E 和我们分词器的重建后，DALL-E 的图像分词器难以区分线条和标记的边界，但我们 dVAE 的图像分词器更接近原始图像，边界更加锐利和清晰。我们确认，较好的分词器能产生更准确的标记，能够更好地描述原始图像。

在获得预训练数据和图像分词器后，我们以 2048 的批次大小、1e-3 的学习率、10K 的 warmup 步数和 0.05 的权重衰减训练 DiT 500K 步。Adam [23] 优化器的 𝛽1 和 𝛽2 分别为 0.9 和 0.999。我们在 DiT 的预训练中采用了随机深度 [20]，深度率为 0.1，并禁用了 dropout，和 BEiT 预训练相同。我们还在 DiT 的预训练中应用了块级掩码，40% 的补丁被掩盖，像 BEiT 一样。

在 RVL-CDIP 上的微调
我们在 RVL-CDIP 上评估了预训练的 DiT 模型和其他图像骨干网络，用于文档图像分类。我们对图像 Transformer 进行了 90 个周期的微调，批次大小为 128，学习率为 1e-3。在所有设置中，我们使用 RandomResizedCrop 操作将原始图像调整为 224 × 224。

在 ICDAR 2019 cTDaR 上的微调
我们在 ICDAR 2019 数据集上评估了预训练的 DiT 模型和其他图像骨干网络，用于表格检测。由于物体检测任务的图像分辨率要比分类任务大得多，因此我们将批次大小限制为 16。对于档案和现代子集，学习率分别设置为 1e-4 和 5e-5。在初步实验中，我们发现直接使用档案子集中的原始图像进行微调 DiT 会导致次优性能，因此我们应用了 OpenCV [4] 实现的自适应图像二值化算法来二值化图像。图 5 中展示了该预处理的示例。在训练过程中，我们应用了 DEFT [6] 中使用的数据增强方法，作为多尺度训练策略。具体来说，输入图像以 0.5 的概率裁剪成一个随机矩形补丁，然后再次调整大小，使得最短边至少为 480 像素，最多为 800 像素，最长边不超过 1333 像素。

在 PubLayNet 上的微调
我们在 PubLayNet 数据集上评估了预训练的 DiT 模型和其他图像骨干网络，用于文档布局分析。与 ICDAR 2019 cTDaR 数据集类似，批次大小为 16，基础版的学习率为 4e-4，大版本的学习率为 1e-4。我们还使用了 DETR [6] 中的数据增强方法。

在 FUNSD 上的微调
我们使用相同的目标检测框架对预训练的 DiT 模型和其他骨干网络进行微调，处理 FUNSD 中的文本检测任务。在文档布局分析和表格检测中，我们使用锚框大小 [32, 64, 128, 256, 512]，因为检测区域通常是段落级别的。与文档布局分析不同，文本检测的目标是定位文档图像中的较小物体，通常是单词级别。因此，在检测过程中，我们使用锚框大小 [4, 8, 16, 32, 64]。批次大小设置为 16，基础版的学习率为 1e-4，大版本的学习率为 5e-5。

我们选择的图像骨干模型作为基线，其参数量与我们的 DiT-B 相当。包括两种类型：CNN 和图像 Transformer。对于基于 CNN 的模型，我们选择 ResNext101-32×8d [40]。对于图像 Transformer，我们选择 DeiT [36]、BEiT [3] 和 MAE [17] 的基础版，这些模型是在 ImageNet-1K 数据集上预训练的，输入尺寸为 224×224。我们重新运行了所有基线模型的微调。

4.3 结果

RVL-CDIP
表格 1 显示了在 RVL-CDIP 数据集上的文档图像分类结果。为了进行公平比较，表中方法仅使用数据集中的图像信息。DiT-B 显著优于所有选定的单模型基线。由于 DiT 与其他图像 Transformer 基线共享相同的模型结构，因此较高的得分表明我们的文档特定预训练策略的有效性。更大的版本 DiT-L 在单模型设置下与之前的 SOTA 集成模型得分相当，这进一步突显了它在文档图像建模上的能力。

PubLayNet
表格 2 显示了在 PubLayNet 数据集上的文档布局分析结果。由于该任务具有大量的训练和测试样本，并且需要对常见文档元素进行综合分析，因此清楚地展示了不同图像 Transformer 模型的学习能力。可以观察到，DeiT-B、BEiT-B 和 MAE-B 明显优于 ResNeXt-101，而 DiT-B 更强于这些强大的图像 Transformer 基线。根据结果，改进主要来自列表和图形类别，并且在 DiT-B 的基础上，DiT-L 提供了更高的 mAP 分数。我们还调查了不同目标检测算法的影响，结果表明更先进的检测算法（在我们的案例中为 Cascade R-CNN）可以将模型性能提升到更高的水平。我们还将 Cascade R-CNN 应用于 ResNeXt-101-32×8d 基线，结果表明 DiT 在基础版和大版设置下，分别超越了其 1% 和 1.4% 的绝对得分，表明 DiT 在不同检测框架下的优势。

ICDAR 2019 cTDaR
表格 3 显示了在 ICDAR 2019 cTDaR 数据集上的表格检测结果。该数据集相对较小，旨在评估低资源场景下模型的少样本学习能力。我们首先分别分析模型在档案子集和现代子集上的性能。在表格 3b 中，DiT 超越了除 BEiT 外的所有基线。这是因为在 BEiT 的预训练中，它直接使用了 DALL-E 的 dVAE，后者是在包含 400M 图像的庞大数据集上训练的，并且包括多种颜色。而 DiT 的图像分词器是用灰度图像训练的，可能对于带有颜色的历史文档图像不够充分。在将 Mask R-CNN 替换为 Cascade R-CNN 后，我们也观察到类似于 PubLayNet 设置的改进，DiT 仍然显著超越其他基线。类似的结论也适用于表格 3c 中的现代子集。我们进一步将两个子集的预测结合成一个单一集合。结果如表格 3a 所示，DiT-L 在所有 Mask R-CNN 方法中取得了最高的 wF1 分数，展示了 DiT 在不同文档类别中的多样性。值得注意的是，IoU@0.9 的指标明显更好，这意味着 DiT 具有更好的细粒度目标检测能力。在所有三个设置下，我们通过最佳模型和 Cascade R-CNN 检测算法，将 SOTA 结果提升了超过 2%（从 94.23 提升到 96.55）的绝对 wF1 分数。

FUNSD（文本检测）
表格 4 显示了在 FUNSD 数据集上的文本检测结果。由于 OCR 的文本检测一直是长期存在的实际问题，我们从一个流行的商业 OCR 引擎中获取了按字检测的结果，以设置一个高水平的基准。此外，DBNet [30] 是一个广泛使用的文本检测模型，用于在线 OCR 引擎，我们还使用 FUNSD 训练数据对预训练的 DBNet 模型进行了微调，并评估了其准确性。两者在 IoU@0.5 上都达到了约 0.85 的 F1 分数。接下来，我们使用 Mask R-CNN 框架来比较不同的骨干网络（CNN 和 ViT），包括 ResNeXt-101、DeiT、BEiT、MAE 和 DiT。结果显示，由于先进的模型设计和更多的参数，基于 CNN 和 ViT 的文本检测模型显著优于基线模型。我们还观察到，DiT 模型相比其他模型取得了新的 SOTA 结果。最后，我们进一步用包含 100 万个文档图像的合成数据集训练 DiT 模型，最终 DiT-L 模型取得了 0.9429 的 F1 分数。

5 结论与未来工作

本文提出了 DiT，一个自监督的基础模型，用于通用的文档 AI 任务。DiT 模型通过大规模未标注的文档图像进行预训练，这些图像涵盖了各种模板和格式，非常适合不同领域下游的文档 AI 任务。我们在多个基于视觉的文档 AI 基准数据集上评估了预训练的 DiT，包括表格检测、文档布局分析、文档图像分类和文本检测。实验结果表明，DiT 在各方面超越了几种强大的基线模型，并取得了新的 SOTA 性能。我们将公开预训练的 DiT 模型，以促进文档 AI 研究。

对于未来的研究，我们将使用更大规模的数据集对 DiT 进行预训练，进一步推动文档 AI 中的 SOTA 结果。同时，我们还将 DiT 集成到多模态预训练中，作为视觉丰富的文档理解的基础模型，例如下一个基于布局的模型如 LayoutLM，其中统一的基于 Transformer 的架构可能足以同时处理 CV 和 NLP 应用中的文档 AI。