UDOP：Unifying Vision, Text, and Layout for Universal Document Processing ---- 论文阅读

CVPR2023
论文地址：https://arxiv.org/abs/2212.02623
代码地址：https://github.com/microsoft/i-Code/tree/main/i-Code-Doc

简介

文档人工智能具有独特的挑战，使其有别于其他视觉语言领域。例如，文本和视觉模态之间的跨模态交互在文档中比在常规视觉语言数据中强得多，因为文本模态在视觉上位于图像中。此外，下游任务在领域和范例上是多样化的，例如，文档问答、布局检测、文档分类、信息提取等。这增加了两个挑战

1. 如何利用图像、文本和布局模式之间的强相关性，并将它们统一对整个文档进行建模？
2. 模型如何有效地有效地学习不同领域的不同视觉、文本和布局任务？

UDOP（Universal Document Processing）是一个文档 AI 模型，它将文本、图像和布局模态与各种任务结合在一起，包括文档理解和生成。UDOP利用文本内容和文档图像之间的空间相关性，用一种统一的表示对图像、文本和布局模态进行建模。 UDOP 使用一种 VTL Transformer （Vision-Text-Layout Transformer），将预训练和多领域的下游任务统一成基于 prompt 的序列生成方案。UDOP在 1100 万个公共无标注文档和11个监督数据集的180万数据上进行了预训练。UDOP也可以通过掩码图像从文本和布局模态生成文档图像。这是第一次在文档AI领域，一个模型同时实现了高质量的文档编辑和内容定制。该方法针对金融报告、学术报告和网站数据等不同领域数据的 8 个文档 AI 任务（例如文档理解和QA）达到了领先水平，在文档理解基准排行榜上达到了第一。

UDOP通过 VTL Transformer 和统一的生成性预训练任务（包括视觉任务、文本任务、布局任务和混合任务）将视觉、文本和布局统一起来。上图展示了所有自监督示例（联合文本-布局重构、视觉文本识别、布局建模和掩码自动编码）和两个监督示例（问答和布局分析）的任务 prompt (左) 和任务目标(右)

本文方法

通用文档处理

编码器

从一张文档图像中使用OCR技术可以获得所有文本块以及每个文本块的坐标信息（每个文本块的布局信息）。

在输入阶段，使用一个统一的 Transformer 编码器融合视觉、文本和布局模态。在文档图像中，文本是嵌入在图像中的，文本和图像像素有一一对应的关系，为了利用这种对应关系，本文使用一种 VTL Transformer 来动态的融合和统一基于布局信息的图像像素和文本，接下来详细说明：

首先将图像分成 N 个 patches，把每个 patch 编码成一个 D 维的向量，并组成一个向量序列。同时使用词汇表把文本信息也编码成 D 维的向量。接下来，使用视觉、文本和布局构建一个统一的表示，如下图所示，

通过定义一个图像 patch 和文本向量的布局指示函数，把文本和图像 patch 进行匹配：如果文本框的中心在哪个 patch 中，所对应的指示函数就是1，否则就是0，

然后把每个文本把匹配到的图像 patch 和自身相加得到了文本和图像的联合表示 。对于那些没有匹配到任何文本的图像 patch，联合表示就是图像本身。

然后把得到的联合表示输入到 VTL Transformer 中。这些联合表示通过显式的利用视觉、文本和布局之间的空间关系，大大增强了模型输入阶段三者之间的交互作用。

为了进一步统一布局和文本表示，受生成对象检测的启发，本文将布局模态（即连续坐标文本边界框）离散化为布局 token。假设我们有归一化到 [0, 1] 的边界框 $(x_i^1,x_i^1,y_i^2,y_i^2)$。生成的布局 token 将是每个坐标乘以词汇量，然后四舍五入到最接近的整数。例如，如果我们有布局词汇量为 500 的边界框 (0.1, 0.2, 0.5, 0.6)，布局 token 将是<50><100><250><300>，布局 token 可以方便地插入到文本上下文中，并优雅地用于布局生成任务（例如，位置检测）。之后会用到。

与 TILT¹ 相同将 2D 文本位置编码为 2D 相对注意力偏置，类似于 T5 中使用的相对注意力偏置。没有使用 1D 的位置嵌入，因为联合嵌入和 2D位置偏置已经包含了输入文档的布局结构。

解码器

解码器包含文本-布局解码器和视觉解码器。文本-布局解码器是一个单向的 Transformer 解码器，以序列到序列的方式生成文本和布局 token。视觉解码器采用的是 MAE²，直接生成带有文本和布局信息的图像。

统一生成式预训练

为了统一不同的训练目标和数据集，本文创建了一个具有 task prompt 的通用任务格式。在无标注和有标注的大规模文档上进行预训练，下表总结了文任务的 prompt 和 target：

以下通过例子 “Ship Date ti Retail: Week of March 14, 1994” 来说明：

自监督预训练任务

联合文本布局重建 让模型重建文档中丢失的文本和位置。通过 mask 部分文本，然后问模型这些文本以及文本框，例如假设 mask 文本 “Date” 和 “of”，输入序列和目标序列可以写成：

其中 <text_layout_0> 和 <text_layout_1> 表示文本布局提示标记。<100><350><118><372> 和 <100><370><118><382> 表示文本 “Date” 和 “of” 的布局 token。类似于 MLM（Masked Lauguage Modeling），文中使用 15% 的掩码率，这个任务可以解释为 Masked Text-Layout Modeling。
布局建模 给定文档图像和文本，让模型去预测文本的位置。例如，预测 “Ship Date” 和 “of” 的位置，输入序列和目标序列可以写成：

这个预训练任务有不同的提示词 <layout_sent_0>，使用比较大的掩码率 75%，小的掩码率太简单了。
视觉文本识别 给定文档图像的位置，让模型预测文本。例如，识别 <100><350><118><372> 和 <100><370><118><382> 这两个位置的文本，输入和目标可以写成：

这个预训练任务的提示词是 <text_0>，使用 50% 的掩码率。这个任务有助于模型通过理解视觉-文本对应关系来学习视觉-文本联合嵌入。
用文本和布局重建掩码图像

本文使用了 MAE² 目标进行视觉自监督学习。MAE 掩码一定比例的图像 patches，并将非掩码的 patches 输入到视觉编码器中。然后将编码器的输出输入到视觉解码器中来重建掩码的 patches。MAE使用了均方误差，并且只在掩码 patches 上使用。本文对 MAE 做了以下修改来做文档图像生成和统一框架：

1. 带有字符嵌入的交叉注意力：在文档中，文本内容主要由字母、数字和标点组成。文本的字符级组合应该有助于视觉生成。本文在视觉解码器中添加了交叉注意力，既关注文本编码特征，也关注字符嵌入特征。字符嵌入是训练参数，不通过编码器编码。这种字符交叉注意只增加了线性计算的复杂度，但大大提高了图像生成的质量。
2. 图像解码：本文没有直接把统一编码器的输出输入到视觉解码器，因为联合视觉文本嵌入只包含了非掩码图像 patches，并且图像 patches 和文本融合在一起。因此，本文提出了用一系列可训练的占位符嵌入作为视觉解码器的输入。本文使用了两种类型的占位符嵌入来指示图像 patches 是否在输入文档图像中被 mask。视觉解码器通过交叉注意力编码视觉文本输出和字符嵌入。

监督预训练任务

本文包括了以下的监督任务：文档分类、布局分析、信息提取、文档和文档自然语言推理。下游任务的验证或测试集不用于监督预训练。
文档分类 这个任务是去预测文档图像的类别。任务提示词是 “Document Classification on (Dataset Name)”， 比如 “Document Classification on RVLCDIP”，然后跟上文本。目标是文档的类别。数据集是RVLCDIP， 一共有 16 类。
布局分析 这个任务是预测文档图像中一个实体的位置，比如标题、段落等等。任务提示词是 “Layout Analysis on (Dataset Name)”，然后跟上实体的名称。目标是给定实体的文本框。数据集是 PubLayNet。
信息提取 这个任务是去预测文本查询的实体类型和位置（例如摘要段落）。提示词是 “Information Extraction on (Dataset Name) (Text Query)”，目标是要查询的每个 token 的实体标签和位置。使用数据集 DocBank, Kleister Charity(KLC)， PWC 和 DeepForm。
问答 这个任务是回答一个与文档相关的问题。提示词是 “Question Answering on (Dataset Name)”，然后跟上问题和文档所有的词。目标是这个问题的答案。使用数据集 WebSRC，VisualMRC，DocVQA， InfographicsVQA 和 WTQ (WikiTableQuestions)。
文档自然语言推理 这个任务是预测文档中两个句子之间的蕴含关系。提示词是 “Document Natural Language Inference on (Dataset Name)”，然后跟上句子对。目标是 “Entailment” 或者 “Not Entailment”。使用数据集 TabFact。

实验

模型预训练

模型配置
在UDOP中，统一的编码器和文本布局解码器遵循T5-large³的编码器-解码器架构。视觉解码器为 MAE-Large ²。总体UDOP有 794M 个可训练参数。对于 tokenizer，本文使用 T5 tokenizer 和 Hugging Face 中的 Transformers 嵌入。本文还扩展了词汇表用来适应特殊的 token 。
数据
自监督学习使用 IIT-CDIP Test Collection 1.0，这是常用的大规模文档数据集。包含 1100 万份扫描文档，其中包含OCR结果（文本和坐标）。监督数据集如上所述。
课程学习
本文在最终设置中使用大型图像分辨率 1024，因为低分辨率使得文档文本对于检测和生成都无法识别。这将导致 $(1024/16{)}^2=4096$ 个图像 patch 序列长度，其训练时间比小图像分辨率长，例如 224。因此，本文使用课程学习从相对较小的分辨率开始，逐渐扩展到 1024 分辨率。实际训练时，本文在预训练期间使用具有 3 个分辨率的等级 224 → 512 → 1024 。
训练
使用 Adam 优化器，学习率 5e-5, 1000 warmup steps，batch size 512，1e-2 weight decay， ${\beta }_1=0.9,{\beta }_2=0.98$。每一个课程学习阶段，训练 1 epoch。

下游评估

在数据集 FUNSD、CORD、RVL-CDIP 和 DocVQA 上的结果

在数据集 DUE-Benchmark 上的结果

在 DUE-Benchmark 上不同分辨率的结果

可视化结果

掩码图像重建

文档生成和编辑

布局定制

---

1. Powalski, Rafał, et al. “Going Full-TILT Boogie on Document Understanding with Text-Image-Layout Transformer.” arXiv: Computation and Language, Feb. 2021. ↩︎

2. He, Kaiming, et al. “Masked Autoencoders Are Scalable Vision Learners.” arXiv: Computer Vision and Pattern Recognition, Nov. 2021. ↩︎ ↩︎ ↩︎

3. Raffel, Colin, et al. “Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer.” arXiv: Learning, Oct. 2019. ↩︎