SS-TDNet论文阅读

本文分析了目前Tree Decoder的缺陷，并提出了一些解决方案

摘要

Offline HMER任务近期使用TreeDecoder进行。尽管Tree Decoder将数学表达式视为树结构，并且将2D空间结构解析到树节点上。

但目前的方法存在以下两个问题：

目前使用的方式受制于树节点预测的错误，性能仍然较差。
它们缺乏语法规则来规范表达式的输出。

本文提出了一个新的模型SS-TD，即Spatial Attention and Syntax Rule Enhanced Tree Decoder，使用了一个特殊的attention machanism来缓解树结构的预测误差，并使用语法掩码（有语法规则的变换得到）来约束不合法的数学表达式的出现

引言

最近使用Tree Decoder的方法有很多：

SAN
A Tree-Structured Decoder for Image-to-Markup Generation
A tree structure based decoder for online handwritten mathematical expression recognition.（DenseWAP-TD）

上述文章重点介绍了HMER结构领域正在兴起。

基于树解码器的方法本质上将表达式表示为树结构，这对于HMER更自然

DenseWAP-TD这篇论文中提出，使用Tree Decoder来预测解码过程中树的父子关系。树解码器虽然显式地建模了HME的树结构，但由于只利用节点符号信息预测节点结构，结构预测精度可能会降低。

本文受到DenseWAP-TD模型的启发，提出了SS-TD模型，用于整合空间注意力和语法规则。

首先提出空间注意力机制来预测父节点，这一工作有效地提高了结构准确度。

此外，为了减少不合语法的数学表达式的产生，我们将表达式符号之间关系的语法转化为语法掩码，并将其引入关系预测过程。

本文的主要贡献如下：

提出了新的Tree Decoder（SS-TD）
提出空间注意力机制来强化节点信息预测，提出语法掩码来限制非法公式的产生
该方案在CROHME14/16/19数据集上效果不错

结论

本文论证了结合空间注意力和语法掩码的新TreeDecoder的有效性，以及该结构能够减少数学表达式的显性语法错误。

工作

本文以HME图片作为输入，输出一个三元组：

$(y_t^c, y_t^p, y_t^{rel})$ 分别代表子节点、父节点，以及父子节点之间的关系。

其中父节点是之前的节点 $y_1^c,...,y_{t-1}^c$ 中与当前子节点 $y_t^p$ 在语义空间信息上最匹配的节点

$y_t^{rel}$ 包含六个预测值：After、Below、Sup、Sub、Inside、Right

本文与Zhang提出的Dense-WAP使用同样的DenseNet作为Encoder

而本文的解码器包含三个预测模块：

子节点预测模块
父节点预测模块，用来预测与当前节点匹配最好的前一个子节点
关系预测模块用来预测父子节点之间的关系

预测模块最终得到一个三元组作为输出结果，最后遍历树结构得到LATEX字符串。

子节点预测模块

与之前基于树解码器的方法DenseWAP - TD一样，本文也采用父节点解码器和子节点解码器来依次识别子节点符号，如上图所示。这两个解码器生成一个包含符号及其识别顺序的元组，定义为子节点信息。

父解码器由两个GRU单元和一个注意力模块组成。

以前一个预测得到的子节点 $y_{t-1}^c$ ，隐状态 $s_{t-1}^c$ 作为输入，输出父上下文向量 $c_t^p$ 及其隐状态 $s_t^p$ ，这也是子解码器和另外两个模块ParentN.P和Relation.P的输入。

$ c_t^p = \sum_{i=1}^L a_{ti}^ph_i$

$ s_t^p = GRU^p_2(c_tp, \hat{s}_t^p)$

其中 $h_i$ 是特征图E的第i个元素， $L = H' \times W'$

$s_t^p$ 是当前父隐藏状态的预测值

$a_t^p = \{a_{ti}^p\}$ 是当前父解码器计算得到的注意力权重，计算发放如下：

$\hat{s}_t^p = GRU_1^p(y_{t-1}^c,s_{t-1}^p)$

$a_t^p = f_{att}^p(\hat{s}^p_t,a_{1,2,...,t-1}^p)$

其中 $f_{att}^p$ 表示一个注意力函数

子解码器的结构与父解码器一致，子解码器以上一个父节点 $y_{t-1}^p$ 和它的隐藏状态 $s_t^p$ 作为输入，输出子上下文向量 $c_t^c$ 和它的隐状态 $s_t^c$

最后使用一个全连接层和一个softmax，输入 $y_{t-1}^p,s_t^c,c_t^c$ (上一个父节点，当前子节点隐状态，当前子节点特征向量)来预测当前节点的值。

$p(y_t^c) = softmax(W_{out}^c(W_ey_{t-1}^p + W_h s_t^c + W_c c_t^c))$

子节点预测模块的loss使用如下方式计算：

$L_c = -\sum^T_{t=1}log(p(y_t^c))$

基于空间注意力的父节点预测模块

本文将父节点的预测任务视为寻找与当前节点最匹配的先前生成的节点。

本文定义先前子节点的顺序为父节点的位置 $\hat{y}^{pos}$

首先使用了一个MLP结构来获取语义信息因子 $e_{ti}^{mem}$ ，即图中 $MLP_1$ 所示的功能：

$e_{ti}^{mem} = v_{mem}^Ttanh(W_{mem}s_t^p + U_{mem}s_{1,2,...,t-1})^c$

其中 $e_{ti}^{mem}$ 表示第t步的第i个语义信息因子

$s_t^p$ 是父编码器的隐状态， $s_{1,2,...,t-1}^c$ 是历史子节点的隐状态

然而，仅仅使用语义信息来预测父节点可能会导致不可避免的误差，尤其是当面对相同的表达符号时。因此，本文还引入空间注意力机制来缓解父节点的预测误差。

父节点和前一个子节点的空间信息分别表示为注意力分布 $a_t^p$ 和 $a^c_{1，2，..，t - 1}$ 。空间能量因子的计算如下：

$e^{alpha}_{ti} = v^T_{alpha}tanh(W_{alpha}p(a^p_t ) + U_{alpha}p(a^c_{1,2...,t−1}))$

该值同样表示第t步的第i个空间信息因子

$p(.)$ 表示自适应池化,其输出向量的大小为4 × 32。

本文进一步采用门机制来控制更新的空间位置信息的输入：

$gt = σ(W_{yg}y^c_{t−1} + U_{sg}s^c_{t−1})$

$e^{position}_{ti} = e^{mem}_{ti} + g_t \odot e^{alpha}_{ti}$

$G^{position}_{ti} = σe^{position}_{ti}$

式中 $\odot$ 为元素乘积。

该模块的Loss使用交叉熵算是函数，公式如下：

$L_{pos} = -\sum_{t-1}^T \sum_{i=1}^L[\bar{G}_{ti}^{position}log(G_{ti}^{position}) + (1 - \bar{G}_{ti}^{position})log(1 - G_{ti}^{position})]$

其中 $\bar{G}_{ti}^{position}$ 表示父节点的真标签，如果第i个历史子节点是当前节点的父节点，那么 $\bar{G}_{ti}^{position}$ 值为1，否则为0

预测时则直接取最大值：

$\hat{y}^{pos} = \underset{g}{ \arg \max \, G^{position}}$

随后使用表示子节点序列的向量 $\hat{y}^{pos}$ 得到父节点的预测 $\hat{y}^p$

随后本文将父节点信息 $\hat{y}^{pos}$ 输入子解码器来预测子节点。

基于语法规则的关系预测

本文在关系预测模块中引入了语法规则，规则如下：

不同的数学符号有不同的连接关系限制，比如 $\sum$ 不能与其他符号有’ Inside '关系。
每个数学符号不能同时存在重复的连接关系，就像一个符号不能与另外两个符号存在"Right"的关系一样

对于第一条规则，本文为了将语法规则结合到训练过程中，提出了用于表示数学连接关系的语法掩码。即如果两个数学符号能够以Right、Sup、Sub、Above、Below、Inside的关系表示，则对应位置掩码值为1，否则为0.

静态语法掩码如下：

这些掩码会被作为先验知识输入。作者将字符根据掩码分类来减少先验知识的数量，同时，对于具有多种可能的字符，可以直接使用逻辑或操作将每种可能的掩码进行处理。例如字符e，当它为小写字母变时，掩码为111000，而当它作为自然底数的时候，掩码为110000，那么字符e最后的掩码为 $111000||110000 = 111000$

本文使用一个掩码矩阵 $R_{static} \in R^{C \times 6}$ 来表示，父节点的静态掩码矩阵即为 $m_t^s \in R^6$ ：

$m_t^s = y_t^p R_{static}$

$y_t^p$ 是当前父节点的独热向量

对于第条规则，模型需要获取到父子节点之间关系的历史信息。本文提出来一个动态掩码矩阵 $R_{dynamic}^t$ ，它被初始化为一个全零矩阵来存储t步中的历史关系。例如，当父节点及其关系被预测为’ \ sqrt ‘和’ Inside ‘时， $R_{dynamic}^t$ 中’ \ sqrt ‘对应行和’ Inside ‘对应列的值应更新为’ 1 '。父节点的动态语法掩码表示为 $m_t^d \in R^6$ :

$m_t^d = y_t^p R_{dynamic}^{t-1}$

$m_t = m_t^s \otimes m_t^d$

其中 $\otimes$ 表示逻辑异或运算

动态掩码矩阵的更新计算如下：

$R_{dynamic}^t = R_{dynamic}^{t-1} + y_t^{p^T}y_t^{rel}$

其中 $y_t^{p^T}$ 表示父节点 $y_t^{p}$ d的转置

最终预测位置关系的公式如下：

$p(y_r^{rel}) = softmax(m_t \odot (W_{out}^{rel}(W_{cp}c_t^p + W_{cc}C_t^c)))$

其中 $\odot$ 代表逐元素乘

分类误差计算如下所示：

$L_{rel} = -\sum_{t=1}^T log(p(y_t^{rel}))$

预测阶段的任务我们可以用以下方程表示：

$\hat{y}^{rel} = \underset{y^{rel}}{\arg \max\, p(y_t^{rel})}$

损失计算

本文使用一个注意力自正则化损失来加速网络收敛。具体来说，采用KL散度来衡量父子解码器产生的注意力分布差异：

$L_{alpha} = \sum_{t=1}^{T} \hat{a}_t^p log(\frac{\hat{a}_t^p}{a_t^p})$

因此最终得到的目标损失函数如下：

$L = \lambda _1 L_c + \lambda _2 L_{pos} + \lambda _3 L_{rel} + \lambda_4 L_{alpha}$

其中 $\lambda$ 是每个项目的权重

实验

实验配置

对于训练Loss，权重 $\lambda_1 = \lambda_2 = \lambda_3 = 1, \lambda_4 = 0.1$

本文使用的Encoder为DenseNet，包含三个Dense块和两个Transition层，每一个Dense块均包含48个卷积层，卷积核大小为11*11

模型在一台两张3090的设备上进行

使用的数据集是CROHME14/16/19

本文对DenseWAP-TD模型进行了简化，将子注意力和父注意力以及记忆注意力的embedding维度由512更改为了128

消融实验

消融实验仅在CROHME14上进行

本文对空间信息、静态语义掩码、动态语义掩码分别进行了消融实验，第一项表示树结构的表达式准确率，第二项表示表达式准确率，第三项表示父节点的预测错误率，最后一项表示父子节点关系的预测错误率。

上表论证了影响识别正确率的是空间注意力机制而不是注意力权重。

本文也比较了不适用注意和使用注意力的情况，使用注意力时计算

$\frac{precision(\%)}{Params(M)}$

得到的值为10.79，比没有使用注意力得到的值多了0.12

如下图所示

当仅使用语义信息预测时得到了错误的结果1

在整合了空间信息之后，能够正确地预测到父节点是n

最后本模型与其余模型对比结果如下：

本文提出的模型仍然低于一些字符串编码器，例如

Truong, T., Ung, H.Q., Nguyen, H.T., Nguyen, C.T., Nakagawa, M.: Relationbased representation for handwritten mathematical expression recognition. In: ICDAR 2021: 16th International Conference Document Analysis and Recognition. pp. 7–19 (2021), https://doi.org/10.1007/978-3-030-86198-8_1

中使用一个额外的弱监督分支达到了53.35%

Ding, H., Chen, K., Huo, Q.: An encoder-decoder approach to handwritten mathematical expression recognition with multi-head attention and stacked decoder. In: ICDAR 2021: 16th International Conference Document Analysis and Recognition. pp. 602–616 (2021), https://doi.org/10.1007/978-3-030-86331-9_39

中使用数据增强和Beam Search解码的大规模输入图像在CROHME 16/19两个数据集上分别达到了57.72%和61.38%，但这一方法的迭代时间更长。

作者认为本方法过于追求表达式的空间结构分辨率二损失了部分符号的识别性能。

总结

本文提出来一个SS-TD模型，并说明了该模型的良好新能，未来作者将进一步提升SS-TD的泛化性。