Nougat 深度剖析 | 莫叶何竹🍀

type

status

date

slug

summary

0 背景

近日，MetaAI又放了大招，他们提出了一种全新的端到端的OCR模型，该模型基于自回归的方法，旨在实现给定图片后输出对应的Markdown标记。一个模型实现版面分析、文本检测、文本识别、公式识别等功能。笔者从论文、源码、测试对Nougat进行深度学习与理解。下面一起来看Nougat是如何做的吧！

1 方法大意

1.1 模型架构

该模型采用了常规的“编码器-解码器”（encoder-decoder）架构，下面对其进行详细说明：

编码器（Encoder）：

模型：使用了 SwinTransformer 模型2 作为编码器。

输入：接收来自PDF图像的输入，其分辨率为DPI96，并去除白边，然后将图像padding到指定的尺寸 (H, W)=(896, 672)。

输出：编码器的输出是图像patch的embedding序列。

解码器（Decoder）：

模型：采用了 mBART 模型的解码器部分3 作为解码器。

输入：图片patch的embedding序列

输出：token概率向量。（最长自回归解码长度为4096）

以上可见Nougat的encoder与decoder都用了较大transformer架构，整体pipeline得参数量达350M。

1.2 数据工程

Nougat将OCR的问题定义为：

核心关键是：如何用一种cheap的方法构造（图片，对应的markdown）pair。于我而言，这是这篇文章最有价值、最值得借鉴学习的地方。

1.2.1 数据源

目前并无大规模的pdf图片与对应markdown标记pair的数据集。Nougat从arXiv、PMC (PubMed Central)， IDL(Industry Documents Library)三个来源构建数据集。其中PMC与IDL的数据由于语义信息不充足仅用于预训练阶段，来使模型具备基础的ocr能力。arXiv数据有tex源码，能拿到所有需要的语义信息，用于预训练和微调阶段。

数据源	Pages	简介	使用阶段
arXiv	750w	有tex源码，信息最充足	预训练+微调
PMC (PubMed Central)	53.6w	有xml源码。但xml文件经常将公式、表格存为图片，导致语义信息不足。	预训练
IDL (Industry Documents Library)	44.6W	仅有text信息，缺失format信息，语义信息不充足。	预训练

1.2.2 图文对构建pipeline

1.2.2.1 思路介绍

图文对构造的整体pipeline由下图所示。从arXiv拿到的Tex源码出发拿到全篇文章的markdown标记，与pdf每页的图片与文本

Branch1:

Branch2:

TEX2HTML转化工具为 LaTeXML http://dlmf.nist.gov/LaTeXML/。（tex源码自定义性过强，转为HTML主要为了消歧）
HTML2markdown 转为代码见源码位置：nougat/dataset/parser/html2md.py

目前我们只能得到全文的markdown标记与pdf图片文本对。剩下需要做的就是根据pdf的page text信息将markdown进行划分得

1.2.2.2 markdown 划分

代码位置：nougat/nougat/dataset/split_md_to_pages/split_markdown

1.2.2.2.1 预处理

预处理1 去除PDF中的图片表格：

由于图片表格在PDF的位置和tex源码的位置可能有所差异，因此作者采取的办法是先用pdffigures2工具将PDF的图片和表格移除。当划分完markdown后再在markdown的末尾加入移除的信息。

pdffigures2提取的信息如

预处理2 将pdf的text格式转化：

去除PDF text中的尾注、页码等。

预处理3 将pdf的text格式转化为latex编码：

后续的markdown划分中会依据markdown的序列与pdf text的匹配度，为了更好的匹配，最好将pdf的text用pylatexenc工具转为为latex编码。

1.2.2.2.2 markdown page 划分

叙述核心逻辑，详细细节见源码

STEP1: HTML解析的全文markdown按段落划分, 得到doc_paragraphs_full，数据结构: List[str]，每一个元素是段落

STEP2: 用fitz以block拿到每页的text, 得到pdf_content，数据结构: List[List[str]]，例：pdf_content[0][0]为第一页的第一个block的文本信息。

STEP3: 基于pdf_content训练page分类器

分类器的数据: , N为页面数量，M_i为页面i中的文本block数量。随后将用结合了TF-IDF的Bag of Word (BoW)进行向量化。

分类器的标签：BoW向量对应的页面标签。

分类器的模型：SVM。

STEP4: 预测markdown的所有段落文本doc_paragraphs_full所属的页面标签。

STEP5: 根据Gini impurity来refine预测的doc_paragraphs_full页面标签，使其满足阶梯状分布。

为在a, b间预测为页面i的概率。t为最好的分割位置。

STEP6: 通过STEP5得到将markdown的段落以页划分。但这是段落级别的划分，字符始末位置与PDF任有差异，需要进一步修正。核心思想为(详细实现参考代码，此处省略了很多细节，仅用于理解)：

从PDF text中找到当前页的初始几个单词和末尾几个单词。

拿到当前页markdown的第一段，使用以编辑距离为指标的模糊匹配（库fuzzysearch）找到初始几个单词的位置，对其进行截断。

拿到当前页markdown的最后一段，找到末尾几个单词的位置，进行截断。

STEP7: 在每页markdown末尾添加预处理产出的表格、图片信息。

1.2.3 训练中数据增强

训练中的数据增强没太多可以介绍的，就是一些常规的方式。

2 小结

Nougat描绘了这么一个愿景，用一个端到端的方式来实现过去繁琐的数据加工pipeline。但目前尝试来看，并不适用实际场景。单纯从架构来看，主要有以下几点缺陷：

推理速度慢。虽然过去的pipeline设计多个模型，但每个模型都非常轻量化，组合起来的参数量甚至不到Nougat的1/10。

定制化难。

数据集构建成本高。Nougat需要构建带结构化信息的图文对，当构建完毕生产数据domain足够的图文对时，用过去的方法早已完成项目。
训练成本高。主要体现在机器成本，需要更多的GPU，更长的训练时间。
优化成本高。Nougat作为一种端到端的解决方法无法针对特定的badcase进行优化。比如在传统方案中，如果表格OCR这个模块效果较差单独优化即可，不会影响到其它模块。但用端到端的方案，当构建倾向表格的数据时，可能会导致其它场景出现新的badcase。

或许在未来以上问题不再是问题。

Reference

[1] Ali Furkan Biten, Rub` en Tito, Lluis Gomez, Ernest Valveny, and Dimosthenis Karatzas. OCR-IDL: OCR Annotations for Industry Document Library Dataset, February 2022.

[2] Liu, Ze, et al. "Swin transformer: Hierarchical vision transformer using shifted windows." Proceedings of the IEEE/CVF international conference on computer vision. 2021.

[3] Liu, Yinhan, et al. "Multilingual denoising pre-training for neural machine translation." Transactions of the Association for Computational Linguistics 8 (2020): 726-742.