FreeTex背后的算法：UniMERNet算法解读

概述

之前在写 FreeTex，有不少读者问道其背后的算法问题。

本文将根据UniMERNet发布的论文进行进一步研究。

论文标题：UniMERNet: A Universal Network for Real-World Mathematical Expression Recognition

论文地址：https://arxiv.org/abs/2404.15254

看完本文，将会获得以下问题的答案：

• UniMERNet 具体是什么算法？
• 数据集是如何收集的？
• 数据集规模有多大？
• 模型训练需要什么样的资源？

主要贡献

这篇工作聚焦数学表达式识别(MER)这个领域，旨在将基于图像的数学表达式转换为相应的标记语言，如 LaTeX 或 Markdown。

UniMERNet算法和Pix2tex、Texify算法在简单打印表达式(SPE)、复杂打印表达式(CPE)、屏幕捕获表达式(SCE)和手写表达式（HWE)四种数据类型上的评估结果如下图所示。

本文主要有以下两点主要贡献：

• 1.构建了一个通用的 MER 数据集 UniMER，涵盖多场景下的数学表达式。

• 2.提出了新的网络结构UniMERNet，用于公式识别任务。

数据集构建

1. 数据集规模

UniMER 数据集分成训练集和测试集两部分:

• 训练集UniMER-1M：包含1061791个LaTeX图像对。
• 测试集UniMER-Test：包含23789个LaTeX图像对。

UniMER 数据集和同类数据集规模如下表所示，

2. 数据集类型

UniMER 数据集包含四种类型的数据集：SPE、CPE、SCE和HWE。

对于打印渲染的公式(SPE、CPE): 以公开数据集 Pix2tex 作为基础，从 Arxiv、Wikipedia 和 StackExchange 平台中获取额外的latex公式代码进行拓展。

对于屏幕捕获表达式(SCE)：收集了1000页中英文PDF文档，涵盖书籍、论文、教科书、杂志、报纸。通过 Mathpix 进行公式识别，进行人工校正整理。

对于手写表达式(HWE)：直接将公开数据集 CROHME 和 HME100K 进行合并。

最终，进一步考虑数据集长度和复杂性的平衡分布，进行处理采样，得到 UniMER 数据集。

UniMERNet网络

网络整体结构如下图所示。

UniMERNet 整体网络架构是基于 Swin-Transformer 编码器和 mBART 解码器。

UniMER-Encoder 集成了细粒度嵌入(FGE)、卷积增强(CE)和移除滑动窗口(RSW)来提升识别能力。

UniMER-Decoder 采用挤压注意力(SA)来加速推理速度。

1. 细粒度嵌入(FGE)

原始的Swin-Transformer使用一个4×4单层卷积核，步长为4，用于图像块嵌入，导致输入分辨率降低四倍。然而，这些非重叠的卷积核经常提取碎片化的特征，导致公式中的字符被分离，从而影响识别精度。

为此，UniMERNet 实现了一个重叠的细卷积层，该层由两个卷积层、一个层归一化(LN)层和一个GELU激活层组成。这两个卷积层都使用3的内核大小、2的步长和1的填充。这种重叠的细卷积扩展了感受野，同时减轻了碎片化，从而提高了模型的整体性能。

2. 卷积增强(FGE)

Transformer架构很好地捕捉全局信息，但辨别局部细节至关重要。为此，引入了一个名为ConvEnhance的局部感知模块，该模块增强了模型识别上标和下标的能力。

ConvEnhance位于每个注意力和MLP模块之前，由3x3深度卷积和GELU激活函数组成。卷积提供局部感知，而GELU激活函数实现门控阈值滤波。这一改进使UniMERNet能够在局部和全局信息之间切换，显著提高了其性能。

3. 移除滑动窗口(RSW)

滑动窗口操作旨在通过层间重叠窗口来增加模型的感受野，类似于卷积核。

然而，随着FGE和CE的引入，与原始Swin架构相比，感受野已经显著增大。这使得滑动窗口操作变得多余。移除它不仅提高了模型的性能，还加快了模型的运行速度。

4. 压缩注意力(SA)

为了增加模型的输出速度，引入了压缩注意力。

具体来说，压缩注意力将查询和键映射到低维空间，同时尽量减少信息损失，从
而加速注意力的计算。

实验

1. 评估指标

采用2个模型性能评估指标：

• BLEU：用于量化候选句子和参考句子之间的n-gram匹配。
• EditDis：测量将一个字符串转换为另一个字符串所需的最少字符更改。

2. 实验细节

使用8张A100(80GB)显卡进行训练，迭代次数(iteration)为300000。

3. 性能比较

UniMERNet 模型根据参数量从小到大共有T、S、B三个型号，和其它SOTA方法的性能比较如下表所示：

因为UniMERNet是在自己扩充的数据集测试的，因此比较有参考价值的就看Texify*和不同型号的模型比较。

从数值上看，相同参数规模(>300M)下，每类数据集的BLEU基本都有1-2个点的提升，说明模型还是比较有效的。

当然，数据集的作用更为明显，Texify 在使用新的数据集训练后，在 CPE、SCR 和 HWE 三类数据上的性能具有显著提升。

下面是三组示例的可视化对比情况：

总结

这篇文章写得是非常清晰，UniMERNet效果好的主要得益于自建的数据集和网络架构的局部调优，前者的贡献度可能更高一些。

材料征集

之前在浏览PPOCR的官网时，发现其也支持公式识别，并放了一张对比表，包括UniMERNet。

数据地址：https://paddlepaddle.github.io/PaddleOCR/latest/version3.x/module_usage/formula_recognition.html

从它给的数据表来看，百度自研的PP-FormulaNet_plus算法，似乎比UniMERNet更胜一筹。

因此，想向使用过 FreeTex 的用户征集一些困难样本，即一些识别效果差的图片，作为测试案例，以便实际评估算法性能。

如有材料，欢迎添加我的微信zstar1003，发送整理好的图片信息，一起促进产品朝更好的方向进行迭代。