Handwritten Mathematical Expression Recognition with Bidirectionally Trained Transformer

In the last decades, many approaches [4, 32, 37, 27, 13, 29, 28] related to HMER have been proposed. These approaches can be divided into two categories: grammar-based and encoder-decoder based. In this section, we will briefly review the related work in both categories.

Transformer [22] is a neural network architecture based solely on attention mechanisms. Its internal self-attention mechanism makes transformer a breakthrough compared to RNN in two aspects. Firstly, transformer does not need to depend on the state of the previous step as RNN does. Well-designed parallelization allows transformer to save a lot of time in the training phase. Secondly, tokens in the same sequence establish direct one-to-one connections through the self-attention mechanism. Such a mechanism fundamentally solves the gradient vanishing problem of RNN [2], making transformer more suitable than RNN on long sequences. In recent years, RNN is replaced by transformer in various tasks in computer vision and natural language processing [20, 8].

Recently, transformer has been used in offline handwritten text recognition. Kang et al. [14] first adopted transformer networks for the handwritten text recognition task and achieved state-of-the-art performance. For the task of mathematical expression, “Univ. Linz” method in CROHME 2019 [17] used a Faster R-CNN detector and a transformer decoder. Two subsystems were trained separately.

To solve the problem that traditional autoregressive models can only perform left-to-right(L2R) language modeling, many studies have attempted right-to-left(R2L) language modeling. Liu et al. [15] trained a R2L model separately. During the inference phase, hypotheses from L2R and R2L models are re-ranked to produce the best candidate. Furthermore, Zhang et al. [38] used R2L model to regularize L2R model in the training phase to obtain a better L2R model. Zhou et al. [39] proposed SB-NMT that utilized a single decoder to generate sequences bidirectionally. However, all of these methods increase model complexity. On the other hand, our approach achieves bidirectional language modeling on a single decoder while keeping the model concise.

In the encoder part, DenseNet is used as the feature extractor for HME images. The main idea of DenseNet is to increase information flow between layers by introducing direct connections between each layer and all its subsequent layers. In this way, given output features $\mathbf{x}_{0},\mathbf{x}_{1},\ldots,\mathbf{x}_{l-1}$ from $0^{th}$ to $(l-1)^{th}$ layer, the output feature of $l^{th}$ layer can be computed by:

where $\left[\mathbf{x}_{0};\mathbf{x}_{1};\ldots;\mathbf{x}_{\ell-1}\right]$ denotes the concatenation operation of all the output features, and $H_{\ell}(\cdot)$ denotes a composite function of three consecutive layers: a batch normalization (BN) [12] layer, followed by a ReLU [9] layer and a $3\times 3$ convolution (Conv) layer.

Through concatenation operation in the channel dimension, DenseNet enables better propagation of gradient. The paper [11] states that by this dense connection, DenseNet can achieve better performance with fewer parameters compared to ResNet [10].

In addition to DenseNet, we also add a $1\times 1$ convolution layer in the encoder part to adjust the image feature dimension to the size of embedding dimension $d_{\text{model}}$ for subsequent processing.

The positional information of image features and word vectors can effectively help the model to identify regions that need to attend. In the previous studies [36, 34], although the RNN-based model inherently takes the order of word vectors into account, it neglects the positional information of image features.

In this paper, since the transformer model itself doesn’t have any sense of position for each input vector, we use two types of positional encodings to address this information. In detail, we use image positional encodings and word positional encodings to represent the image feature position and word vector position, respectively.

We refer to image features and word vector features as content-based, and the two types of positional encodings as position-based. As illustrated in Fig. 1, content-based and position-based features are summed up, as the input to transformer decoder.

For the decoder part, we use the standard transformer model [22]. Each basic transformer decoder layer module consists of four essential parts. In the following, we will describe the implementation details of these components.

First, two special symbols “$\langle\text{SOS}\rangle$” and “$\langle\text{EOS}\rangle$” are introduced in the dictionary to denote the start and end of the sequence. For the target LaTeX sequence $y=\{y_{1},\ldots,y_{T}\}$, we denote the target sequence from left to right (L2R) as $\overrightarrow{y}=\{\langle\text{SOS}\rangle,y_{1},\ldots,y_{T},\langle\text{EOS}\rangle\}$, and the right-to-left (R2L) target sequence as $\overleftarrow{y}=\{\langle\text{EOS}\rangle,y_{T},\ldots,y_{1},\langle\text{SOS}\rangle\}$.

Conditioned on image $x$ and model parameter $\theta$, the traditional autoregressive model need to compute the probability distribution:

where $j$ is the index in target sequence.

In this paper, since the transformer model itself does not actually care about the order of input symbols, we can use a single transformer decoder for bi-directional language modeling. Modeling both Eq. (10) and Eq. (11) at the same time.

To achieve this goal, a simple yet effective bidirectional training strategy is proposed, in which for each training sample, we generate two target sequences, L2R and R2L, from the target LaTeX sequence, and compute the training loss in the same batch (details in Section 4.2). Compared with unidirectional language modeling, our approach trains a model to perform bidirectional language modeling without sacrificing model conciseness. Experiment results in Section 5.3 verify the effectiveness of our bidirectional training strategy.

In the encoder part, to make a fair comparison with the previous state-of-the-art method, we use the same DenseNet feature extractor as the DenseWAP model [34]. Specifically, three bottleneck layers are used in the backbone network and transition layers are added in between to reduce the number of feature-maps. In each bottleneck layer, we set the growth rate to $k=24$, the depth of each block to $D=16$, and the compression hyperpatameter of the transition layer to $\theta=0.5$.

In the decoder part, we use the standard transformer model. We set the embedded dimension and model dimension to $d_{\text{model}}=256$, the number of heads in the multi-head attention module to $H=8$, the dimension of intermediate layers in the FFN to $d_{ff}=1024$, and the number of transformer decoder layer to $N=3$. The 0.3 dropout rate is used to prevent overfitting.

Our training objective is to maximize the predicted probability of the ground truth symbols in Eq. (10) and Eq. (11), so we use the standard cross-entropy loss function to calculate the loss between the predicted probabilities w.r.t. the ground truth at each decoding position. Given the training sample $\left\{x^{(z)},y^{(z)}\right\}_{z=1}^{Z}$, the objective function for optimization is shown as follows:

The model is trained from scratch using the Adadelta algorithm [33] with a weight decay of $10^{-4}$ , $\rho=0.9$, and $\epsilon=10^{-6}$. PyTorch framework is used to implement our model. The model is trained on four NVIDIA 1080Ti GPUs with $11\times 4$ GB memory.

In the inference phase, we aim to generate the most likely LaTeX sequence conditioned on the input image. Which can be fomulated as follows:

where $\mathbf{x}$ denotes the input image and $\theta$ denotes the model parameter.

Unlike the training phase where a lower triangular mask matrix is used to generate the prediction for all time steps simultaneously. Since we have no ground truth of the previously predicted symbol, we can only predict symbols one by one until the “End” symbol appears or the predefined maximum length is reached.

Obviously, we cannot search for all possible sequences, thus a heuristic beam search is proposed to balance the computational cost with the quality of decoding. Further, taking advantage of the fact that our decoder is capable of bidirectional language modeling, approximate joint search [15] is used to improve the performance. The basic idea consists of three steps: (1) Firstly, a beam search is performed on L2R and R2L directions to obtain two k-best hypotheses.(2) Then, we reverse L2R hypotheses to R2L direction and R2L hypotheses to L2R direction and treat these hypotheses as ground truth to compute the loss values for each of them as in the training phase.(3) Finally, those loss values are added to their original hypothesis scores to obtain the final scores, which is then used to find the best candidate. In practice, we set beam size $k=10$, the maximum length to 200, and length penalty $\alpha=1.0$.

We use the Competition on Recognition of Online Handwritten Mathematical Expressions (CROHME) benchmark, which currently is the largest dataset for handwritten mathematical expression to validate the proposed model. We use the same training dataset but different test datasets. The training set contains 8836 handwritten mathematical expressions, while the test sets of CROHME 2014/2016/2019 contain 986/1147/1199 expressions respectively.

In the CROHME dataset, each handwritten math expression is saved in a InkML file, which contains handwritten stroke trajectory information and ground truth in both MathML and LaTeX formats. We transform the handwritten stroke trajectory information in the InkML files to offline images in bitmap format for training and testing. With official evaluation tools provided by the CROHME 2019 [17] organizers and the ground truth in symLG format, we convert the predicted LaTeX sequences to symLG format and evaluate the performance.

The results of some models on the CROHME 2014/2016/2019 datasets are shown in Table 1. To ensure the fairness of performance comparison, the methods we show all use only the officially provided 8836 training samples. Neither we nor the methods we compare use data augmentation.

We first provide results of three traditional handwritten mathematical formula recognition methods based on tree grammar in CROHME 2014 as the baseline, denoted as I, VI, and VII. For CROHME 2016, we provide the best performing “TOKYO” method as the baseline, which only used official training samples. For CROHME 2019 official methods, we provide “Univ. Linz” [17] method as the baseline. For Image-to-LaTeX methods, we use the previous state-of-the-art “WYGIWYS” [7], “PAL-v2” [23], “WAP” [36], “Weakly supervised WAP”(WS WAP) [21], “DenseWAP” [34] as well as the tree decoder-based “DenseWAP-TD” [35] method. The “Ours-Uni” and “Ours-Bi” methods denote the model trained using the vanilla transformer decoder and the model trained with the bidirectional training strategy on top of it.

In Table 1, by comparing “DenseWAP” with “Ours-Uni”, both are unidirectional string decoder based models, we can obtain 4.29% average performance improvement in ExpRate by simply replacing the RNN-based decoder with the vanilla transformer decoder.

Compared with other methods, our proposed BTTR model outperforms the previous state-of-the-art methods in nearly all metrics and is about 3.4% ahead of the “DenseWAP-TD” method in ExpRate, which explicitly encodes our prior knowledge about the LaTeX grammar through a tree decoder.

In Table 2, the first “IPE” column denotes whether to use image positional encoding or not. Secondly, the “Bi-Trained” column shows whether the bidirectional training strategy is used. On the “AJS” column, ✓ indicates the use of approximate joint search [15], while ✗ represents L2R search. The last “Ensemble” column denotes whether the ensemble method is used.

First, we can see that whether to use image positional encoding or not makes huge difference on the CROHME 2019 test set. This shows that image positional encoding improves the generalization ability of our model in different scales.

Comparing the \nth2 and the \nth3 rows in each dataset, we can see that the model trained using the bidirectional training strategy still outperforms the unidirectionally trained model by about 2.52% in ExpRate, though both of them using L2R search. This shows that while training a bidirectional language model, the bidirectional training strategy also helps the whole model to extract information from the images more comprehensively.

Further, using the properties of the bidirectional language model, we evaluate the decoding results of both L2R and R2L directions using approximate joint search, resulting in an improvement of about 4.68% in ExpRate.

Finally we report the results using the ensemble method, showing that this can significantly improve the overall recognition performance by about 3.35% in ExpRate. Specifically, We train five models initialized with different random seeds and average their prediction probabilities at each decoding step.

As can be seen in Fig. 2, we give several case studies for the “DenseWAP”, “Ours-Uni” and “Ours-Bi” models. These three models use the same DenseNet encoder. The difference between these three models is that “DenseWAP” uses an RNN-based decoder, “Ours-Uni” adapts a vanilla transformer decoder, and “Ours-Bi” uses the bidirectional training strategy to train transformer decoder.

Firstly, comparing the prediction results between “DenseWAP” and “Ours-Uni”, we can see that for input images with complex structure, the “DenseWAP” model cannot predict all the symbols completely. Moreover, for input image with discontinuous structure, the right half is unnecessarily predicted twice. Problems mentioned above reflecting under-parsing and over-parsing phenomenon. However, the “Ours-Uni” and “Ours-Bi” models employed with positional encodings are able to identify all the symbols in these images accurately.

Secondly, by comparing the prediction results of “Ours-Uni” and “Ours-Bi”, we find that “Ours-Bi” gives more accurate predictions. Owing to the re-score mechanism in approximate joint search procedure, “Ours-Bi” avoids the asymmetry of “{” and “}” in the prediction results given by “Ours-Uni”.