ABM thesis translation

Address of thesis ： https://arxiv.org/pdf/2112.03603.pdf

Project address ：https://github.com/XH-B/ABM

One 、Abstract

This paper proposes an attention aggregation model based on two-way interactive learning （ABM）, This model consists of two parallel encoders with opposite directions （L2R and R2L） form . These two encoders distill each other , In the training of one-to-one information transmission at each step , Complementary information in both directions is fully utilized . in addition , In order to deal with mathematical symbols of different scales , This paper proposes an attention aggregation model （AAM）, This model can aggregate attention at different scales . It is worth noting that , In the reasoning stage , Considering that the model has learned knowledge from two directions , So just use L2R Some branches of reasoning , In this way, the size of the original parameters and the reasoning speed can be maintained .

Two 、Introduction

WAP First, two-dimensional attention is introduced , To solve the problem of insufficient space coverage , Here's the picture 1 The two-dimensional attention shown focuses on the sum of the past , Designed to track past alignment information , In this way, the attention model can be guided , Assign a higher probability of attention to the untranslated region of the image . However , The main limitation of this approach is , It only uses historical alignment information , Without considering future information （ Untranslated area ）. Such as , Many mathematical expressions are symmetrical structures , Left “{” And right “}” Parentheses always appear together , Sometimes far away . Some symbols in the equation are related , such as $f$ and $dx$. Most methods only use left to right attention to recognize the current symbol , While ignoring the future information from the right , This may cause attention drift . The dependent information between the symbol and the previous symbol becomes weaker with the increase of its distance . therefore , They do not take full advantage of long-distance correlations or grammatical norms of mathematical expressions .

BTTR Use with two directions transformer Decoder to solve the problem of attention drift , But there is no effective way to BTTR Learn the opposite direction supervision information , also BTTR There is no alignment of attention in the whole learning process , This makes it still limited in identifying long formulas .

DWAP-MSA By adding multi-scale features to the code, we can alleviate the problem of recognition difficulty or uncertainty caused by the change of character scale in mathematical expressions . However , They do not scale the local acceptance domain , Only zoom the feature map , This makes it impossible to pay attention to small characters accurately in the recognition process .

therefore , We proposed ABM frame , The framework contains three models ：（1） feature extraction . Use DenseNet The extracted features ;（2） Attention aggregation module （Attention Aggregation module）. We propose multiscale attention , Recognize characters of different sizes in mathematical expressions , Thus, the current recognition accuracy is improved , It alleviates the problem of error accumulation .（3） Two way learning module （ Bi-directional Mutual Learning module）. We propose a new decoder framework , There are two parallel decoder branches with opposite decoding directions (L2R and R2L), And use mutual distillation to learn from each other . Be careful , Although we use two decoders for training , But we only use one L2R Branch for reasoning .

chart 1 Typical HMER（Handwritten Mathematical Expression Recognition） Model architecture

3、 ... and 、Method

We propose a new end-to-end attention aggregation and two-way mutual learning (ABM) framework , Pictured 2 Shown . It mainly consists of three modules ：

• Feature extraction module （FEM), This module can extract feature information from a mathematical expression image .

• Note the aggregation module (AAM) Integrated multiscale coverage note , Align history note information , In the decoding stage, the different scale features of symbols of different sizes are effectively aggregated .

• Two way mutual learning module (BML) By two parallel decoders with opposite decoding directions (L2R and R2L) form , To complement each other . In the process of training , Each decoder branch can not only learn real latex Sequence , You can also learn the opposite latex Sequence , So as to improve the decoding ability .

chart 2 ABM Model architecture

3.1、 Feature extraction module （Feature Extraction Modul）

Use DenseNet The extracted features , Output $H\times W\times D$, We encode information transformation M dimension （$M=H\times W$）, The output vector is $a=(a_1,a_2,\dots ,a_M)$.

3.2、 Note the aggregation module （Attention Aggregation Module）

The attention mechanism can guide the encoder to pay more attention to the specific area of the input picture . Especially the attention mechanism based on the overall situation , It can better track the alignment information and guide the model to allocate a higher probability of attention to the translation area . Inspired by this , We proposed AAM Modules aggregate different receptive fields on global attention . Different from the traditional attention mechanism ,AAM Not only pay attention to local information , At the same time, we also pay attention to the overall information on the larger receptive field . therefore ,AAM Will produce finer alignment information , And help the model capture more accurate spatial relationships . differ DWAP-MSA The model passes dense Multi-scale branches of the encoder to generate low-level and high-level features ,AAM Propose a hidden state $h_t$、 Characteristics of figure $F$ And global attention ${\beta }_t$ Calculate the current attention weight ${\alpha }_t$, Then get the context vector $c_t$.
$$\begin{gathered} A_s=U_s{\beta }_t,A_l=U_l{\beta }_t \\ {\beta }_t=\sum _{l=1}^{t-1}{\alpha }_l \end{gathered}$$
$U_s$ and $U_l$ It means small nucleus and large nucleus respectively （ Such as 5、11） Convolution of ,${\beta }_t$ Represents the sum of all the attention probabilities in the past , Initialize to zero vector . among ,${\alpha }_l$ For the first time $l$ Step's attention score .

therefore , The current attention ${\alpha }_t$ The calculation process is as follows ：
$${\alpha }_t=v_a^{T}tanh(W_h{h}_t+U_{f}F+W_s{A}_s+W_l{A}_l)$$
The final context vector $c_t$ For characteristic information $a$ And attention t force ${\alpha }_t$ Weighted sum of , The calculation formula is as follows ：
$$c_t=\sum _{i=1}^M{\alpha }_{t,i}{a}_i$$

3.3、 Two way mutual learning module （Bi-directional Mutual Learning Module）

Given a mathematical formula, input the image , The traditional method is decoding from left to right (L2R), This method does not consider the problem of long-distance dependence . Therefore, we propose a bidirectional decoder to translate the input image into two opposite directions （L2R 、R2L） Of Latex Sequence , Then learn to decode information from each other . The two branches have the same architecture , Only in its decoding direction .

For two-way training , Let's add $<sos>$ and $<eos>$ As the beginning and end symbols of latex sequence . Specially , For length is T Of Latex Sequence $Y=(Y_1,Y_2,...,Y_T)$,

From left to right (L2R) Express ：$y=(<sos>,Y1,Y2,...,YT,<eos>)$

From right to left (R2L) Express ：$y=(<eos>,Y_T,Y_{T-1},...,Y_1,<eos>)$

L2R and R2L Branch in step t The probability of prediction at is calculated as follows ：
$$\begin{gathered} p(\vec{y}|{\vec{y}}_{y-1})=W_{o}max(W_{y}E{\vec{y}}_{t-1}+W_h{h}_t+W_t{c}_t) \\ p(\overleftarrow{y}|{\overleftarrow{y}}_{y-1})=W_o^{{}^{\prime }}max(W_y{E}^{{}^{\prime }}{\overleftarrow{y}}_{t-1}+W_h^{{}^{\prime }}{h}_t^{{}^{\prime }}+W_t^{{}^{\prime }}{c}_t^{{}^{\prime }}) \end{gathered}$$
among ,$h_t$、${\vec{y}}_t$ Express L2R Step in branch t Current state of and previous predicted output .${\ast }^{\prime }$ Express R2L Branch .$W_o\in R^{K\times d}$,$W_y\in R^{d\times n}$、$W_h\in R^{d\times n}$ and $W^{d\times D}$ Is a trainable matrix .d、K and n Respectively denote attention dimension 、 Number of all tag classes and GRU Dimension of .E Is an embedded matrix .Max Indicates the maximum activation function . Hide representation $\{h1、h2、...,ht\}$ from ：
$$\begin{gathered} {\hat{h}}_t=f_1(h_{t-1},E{\vec{y}}_{t-1}), \\ {h}_t=f_2({\hat{h}}_t,c_t) \end{gathered}$$
We define L2R The probability of branching is ${\vec{p}}_{l2r}=\{<sos>,{\vec{y}}_1,{\vec{y}}_2,...,{\vec{y}}_T,<eos>\}$,R2L The probability of branching is ${\vec{p}}_{r2l}=\{<eos>,{\overleftarrow{y}}_1,{\overleftarrow{y}}_2,...,{\overleftarrow{y}}_T,<sos>\}$.$y_i$ Is to execute the i Prediction probability of label symbol in step decoding . In order to learn the predicted distribution of the two branches from each other , We need to align by L2R and R2L Generated by decoder LaTeX Sequence . meanwhile , introduce kullback-leibler(KL) Loss to quantify the difference in the predicted distribution between them . In the process of training , We use the soft probability generated by the model to provide more information . therefore , about k Categories , come from L2R The soft probability of a branch is defined as ：
$$\sigma ({\vec{z}}_{i,k},S)=\frac{exp({\vec{z}}_{i,k})/S}{{\sum }_{j=1}^{K}exp({\vec{Z}}_{i,j}/S)}$$
among ,S Indicates the parameters for generating soft labels . The... Of the sequence calculated by the decoder network $i$ The logarithm of symbols is defined as $z_i=\{z_1,z_2,...,z_K\}$. Our goal is to minimize the distance between two branch probability distributions . therefore ,${\vec{p}}_{l2r}$ And ${\overleftarrow{P}}_{r2l}^{\ast }$ Between KL The distance is calculated as follows ：
$$L_{KL}=S^2\sum _{i=1}^T\sum _{j=1}^K\sigma ({\vec{Z}}_{i,j},S)log\frac{\sigma ({\vec{Z}}_{i,j},S)}{\sigma ({\overleftarrow{Z}}_{T+1-i,j},S)}$$

3.4、 Loss function （Loss Function）

Specially , For length is T Of Latex Sequence ${\vec{y}}_{l2r}=\{<sos>,Y_1,Y_2,...,Y_T,<eos>\}$, We will be the first to i Time steps correspond to one-hot The real label is expressed as $Y_i=\{x_1,x_2,...,x_K\}$. The first k It's symbolic softmax The probability is calculated as ：
$${\vec{y}}_{i,k}=\frac{exp({\vec{Z}}_{i,k})}{{\sum }_{j=1}^{K}exp({\vec{Z}}_{i,j})}$$
For multi classification , Target tag with two branches softmax The cross entropy loss between probabilities is defined as ：
$$\begin{gathered} L_{ce}^{l2r}=\sum _{i=1}^T\sum _{j=1}^K-Y_{i,j}log({\vec{y}}_{i,j}) \\ {L}_{ce}^{r2l}=\sum _{i=1}^T\sum _{j=1}^K-Y_{i,j}log({\overleftarrow{y}}_{T+1-i,j}) \end{gathered}$$
The global loss function is ：
$$L=L_{ce}^{l2r}+L_{ce}^{r2l}+\lambda L_{KL}$$

Four 、Experiments

chart 3 ABM Compared with the previous network

chart 4 Translate handwritten mathematical expressions into two directions (L2R and R2L) Of LaTeX The coverage of the sequence pays attention to the visualization process

chart 5 CROHME 2014 last year DWAP（ The baseline ） and ABM Of t-SNE visualization

chart 5 Accuracy of expressions with different lengths .