多模态：Nougat详解

---

前言

科学知识主要存储在书籍和科学期刊中，通常以PDF的形式。然而PDF格式会导致语义信息的损失，特别是对于数学表达式。文章提出Nougat，一种视觉transformer模型，它执行OCR任务，用于将科学文档处理成标记语言。

paper：https://arxiv.org/pdf/2308.13418
github：https://github.com/facebookresearch/nougat/tree/main

一、模型结构

模型是一个encoder-decoder模型，允许端到端的训练。

1. encoder

视觉encoder首先接受一张文档图像，裁剪边距并调整图像大小成固定的尺寸(H,W)；如果图像小于矩形，那么增加额外的填充以确保每个图像具有相同的维度。encoder使用了Swin Transformer，将图像分为不重叠的固定大小的窗口，然后应用一系列的自注意力层来聚集跨窗口的信息。该模型输出一个embedding patch（d*N），其中d是隐层维度，N是patch的数目。

2. decoder

使用带有cross-attention的mBART解码器解码（d*N），生成一系列tokens，tokens以自回归方式生成，使用自注意力和交叉注意力分别关注输入序列和编码器输出的不同部分，最后tokens被投影到vocabulary的大小，产生logits。

3. set

以 96 DPI 的分辨率渲染文档图像。由于 Swin Transformer 的限制性可能输入维度，选择输入大小 (H, W ) = (896, 672)。使用预训练的权重初始化模型。Transformer decoder的最大序列长度为 S = 4096。这种相对较大的尺寸是由于学术研究论文的文本可能是密集的，尤其是表格的语法是token密集型的。BART 解码器是一个具有 10 层的仅解码器Transformer。整个架构共有 350M 参数。此外还使用较小的模型 (250M 参数) 进行实验，序列长度略小 S = 3584，只有 4 个解码器层。

训练：使用AdamW优化器训练3个epoch，batch_size是192；由于训练的不稳定性，选择lr init = 5·10−5的学习率，每15次更新减少0.9996倍，直到达到lrend = 7.5·10−6。

二、数据增强

在图像识别任务中，使用数据增强来提高泛化性是有效的。文章应用一系列的transformation来模拟扫描文档的缺陷和可变性。这些变换包括：腐蚀，膨胀，高斯噪声，高斯模糊，位图转换，图像压缩，网格失真和弹性变换（这些变换可以借鉴一下，几乎包括了所有图像领域常用的变换）。每个都有一个固定的概率来应用给给定图像（这一步在训练前完成）。转换效果如下：

在训练过程中，会用随机替换token的方式给groud truth增加扰动

三、数据

目前没有pdf页面和其对应的source code的成对数据集。
根据arxiv上的开源文章，建立了自己的数据集。对于layout多样性，我们引入了PMC开源非商业数据集的子集。在预训练过程中，也引入了一部分行业文档库数据。

ARXIV
我们从arxiv上收集了174w+的pape，收集其源代码并编译pdf。为了保证格式的一致性，首先用latex2html处理源文件，并将他们转为html文件。然后解析html文件，并将他们转换为轻量级标记语言，支持标题，粗体和斜体文本、公式，表等各种元素。这样，我们能保证源代码格式是正确的，方便后续处理。整个过程如图

PMC
我们还处理了来自PMC的文章，其中除了PDF文件之外，还可以获得具有语义信息的XML文件。我们将这些文件解析为与arxiv文章相同的标记语言格式，我们选择使用PMC少得多的文章，因为XML文件并不总是具有丰富的语义信息。通常，方程和表格存储为图像，这些情况检测起来并非易事，这导致我们决定将PMC文字的使用限制在预训练阶段。

IDL
IDL是行业产生的文档集合。这个仅用在预训练阶段，用于教模型基本的OCR；

splitting the pages

我们根据pdf的页中断来分割markdown标记，然后将每个pdf页面转为图像，来获得图像-标记pair。在编译过程中，Latex会自动确定pdf的页面中断。由于我们没有重新编译每篇论文的Latex源，我们必须启发式地将源文件拆分为对应不同页面的部分。为了实现这一点，我们使用PDF页面上的嵌入文本和源文本进行匹配。

然而，PDF中的图像和表格可能不对应他们在源代码中的位置。为了解决这个问题，我们在预处理阶段去掉了这些元素。然后将识别的标题和XML文件中的标题进行比较，并根据他们的Levenshtein距离进行匹配。一旦源文档被分成单个页面，删除的图形和表格就会在每个页面的末尾重新插入。

四、实验评估

在一篇科研文章中，存在三种不同类型的文本：1) 纯文本，占文档的大部分；2) 数学表达式；3) 表格。在评估过程中，检查这三个部分中的每一个都很重要。这是因为在LaTeX中，表达同一个数学表达式有多种方式。尽管在LaTeXML预处理步骤中消除了一些可变性，但仍然存在大量歧义，这种歧义降低了数学和纯文本的得分。数学表达式的预期得分低于纯文本。

repetitions during inference

我们观察到模型会退化为一遍又一遍地重复相同的句子。模型无法自行从这种状态中恢复过来。在最简单的情况下，最后一个句子或段落被不断地重复（在大模型中也经常出现）。我们在测试集中的1.5%页面观察到了这种行为，但领域外文档的频率会增加。陷入重复循环是使用贪婪解码采样时基于Transformer模型的一个已知问题。模型也可能在两个句子之间交替，但有时会改变一些单词，因此严格的重复检测是不够的。作者提出在训练过程中引入随机扰动以增强模型对错误预测标记的处理能力，以及在推理时检测和处理重复的方法，但是仍然无法根本解决这类问题，这也让模型的实用性受到限制。

个人认为有两种方法可以降低重复和幻觉：

1. 输出添加位置信息的预测，比如TextMonkey和mPlug-DocOwl1.5都把text grounding任务做为下游任务进行微调，在消融实验里也提到了增加文本位置预测可以显著降低模型输出的幻觉；

2. 降低图像里文本的数量，在图片中文本数量较少时，一般不容易出现重复和幻觉，所以可以按上篇文章提到的办法，将图片划分为多个block，再对单个block训练对应的端到端识别模型，这样也可以解决模型输出的重复问题

五、代码

1. 环境安装

按照github教程安装即可，但注意transformer库的版本, 如果transformer版本可能会报错：
TypeError: BARTDecoder.prepare_inputs_for_inference() got an unexpected keyword argument ‘cache_position’

由于Nougat自己写了“prepare_inputs_for_generation”，而与最新版本的transformers实现不同。因此应该在Nougat中安装旧版本的transformers，例如transformers==4.38.2是没有问题的。

2. Dataset（dataset.py）

主要函数：NougatDataset —— return input_tensor, input_ids, attention_mask

input_tensor: 图像经过预处理后的输入编码
input_ids: 文本信息经过tokenizer处理后的id
attention_mask: input_ids对应的mask

依赖函数：SciPDFDataset——return {“image”: img, “ground_truth”: data.pop(“markdown”), “meta”: data}

数据格式：

简单来说，SciPDFDataset输出一个样本的基本信息，NougatDatase则输出模型可直接训练的信息。

3. Model（model.py）

Swintransformer基于timm构建，Bart基于transformer构建。有个小技巧，如果训练显存太大，可以尝试在编码器解码器中添加一个参数

编码器中添加 —— use_checkpoint=True

解码器中添加 —— gradient_checkpointing=True

---

总结

Nougat尝试用一个端到端的方式来实现过去无数小模型+策略配合的结果。确实效果很惊艳，但其缺点也十分明显：

推理速度慢。虽然过去的pipeline设计多个模型，但每个模型都非常轻量化，组合起来的参数量甚至不到Nougat的1/10。
定制化难。
数据集构建成本高。（但是nougat的数据工程确实也很惊艳，非常值得学习！！！）
训练成本高。主要体现在机器成本，需要更多的GPU，更长的训练时间。
优化成本高。Nougat作为一种端到端的解决方法无法针对特定的badcase进行优化。比如在传统方案中，如果表格OCR这个模块效果较差单独优化即可，不会影响到其它模块。但用端到端的方案，当构建倾向表格的数据时，可能会导致其它场景出现新的badcase。