End-to-End Information Extraction by
Character-Level Embedding and Multi-Stage Attentional U-Net

The document image, after being processed by a standalone OCR module to form the text lines, can help to construct the chargrid as follows:

• •

  Character separation: We first convert the text line boxes to character boxes by dividing them horizontally to the number of character reside in the text with the assumption that in each text line, the variance of width and height of each character are minimal. After this step, a list of character boxes $L_{c}=\{(c_{1},b_{1}),(c_{2},b_{2}),\ldots,(c_{N},b_{N})\}$ is obtained where $c_{i}$ is the character index and $b_{i}$ its respective bounding box.

• •

  Character encoding: Each character is indexed, a constant number of most frequent characters $N_{char}$ is chosen. The character boxes are then used to create a mask $CM\in\{0,1,\dots,N_{char}\}^{H\times W}$. For each character box in the list with the character index $i_{char}\in\{1,\dots,N_{char}\}$, the region of the mask covered by the box will be filled with the value $i_{char}$. $CM$ is then converted to one-hot format $CM_{one\_hot}\in\{0,1\}^{H\times W\times(N_{char}+1)}$

At this point, the problem of information extraction is now equivalent to the pixel-wise segmentation on the chargrid: Taking $CM_{one\_hot}$ as input, we wish to train the neural network that can output the mask $M\in\{0,\dots,K-1\}^{H\times W}$ masking the regions of $K$ target fields, with the supervision from labeled mask $LM\in\{0,\dots,K-1\}^{H\times W}$ for each input document.

The architecture of our model is depicted in Figure 1. We use the backbone as Coupled U-Net [Tang et al.(2018)Tang, Peng, Geng, Zhu, and Metaxas], with the inspiration behind this utilization is the benefit of multi-stage outputs. For example, in order to understand what kind of information a value in a table represents, a person have to look at the top header or the side description (see Figure 2(a)). This order of behavior can be satisfied by forcing the middle U-Net block to segment key and/or description while the final end outputs the full segmentation of all key/value collections. In addition, other advantages of the Coupled U-Net are also leveraged such as the avoidance of gradient vanishing while promoting feature reuse by the dense skip-connections between corresponding encoder and decoder blocks of each U-Net to another; and the use of auxiliary loss applied to the output of middle U-Net block.

The detail is further explained as follows: Firstly, for the sake of flexibility, each encoder/decoder block is constructed by chaining multiple ResBlocks, each of which is a ResNet-like bottleneck, with ($3\times 3$) convolutions, batch normalization and residual connection. The downsampling blocks are equipped with dilated convolution while the upsampling blocks use transposed convolution blocks for upscaling their resolutions. Secondly, from an intuitive observation that different information flows from table headers and descriptors are required to extract the right target value (illustrated in Figure 2(a)), we incorporate the Box-Convolution layer in [Burkov(2018)] (Figure 2(b)) to a variant of ResBlock, which helps modeling long range vertical and horizontal interaction. The box convolutions [Burkov(2018)] are sliding mean window with learnable offsets instead of weights over windows (see Figure 2(b)). Learning the offset patterns aids the model in aggregating information from distant regions in a straight forward way, making it a suitable alternative to conventional convolution blocks. Thirdly, to strengthen the Coupled U-Net design, expansion paths (also known as skip-connection) from the same level of downsampling blocks to upsampling blocks are added, with self-attention mechanism (non-local network) [Zhang et al.(2018)Zhang, Goodfellow, Metaxas, and Odena] (see Figure 3) to promote global information propagation between each encoder-decoder pair of an U-Net block.

In short, with the described building blocks, we have built an architecture that is light-weight and feasible to exploit both textual and spatial relations from the chargrid embedding.

We specify the additional detail about the training stage, including model hyper-parameters, the loss function and the data augmentation process. The hyper-parameters of the standard $msau$ model is chosen as follows: the number of most-frequent characters $N_{char}=256$, number of U-Net blocks $n_{blocks}=2$, the depth (number of convolution operations) of the ResBlock $res\_depth=2$, the number of channels of the first input feature map $C=16$, the number of downsampling blocks in an encoder $n_{downsampling}=4$. $msau\_big$ has the same configuration except $n_{downsampling}=5$. The same setups are also used for our baseline variants of U-Net, which contain a single encoder-decoder block. Our code, trained model and test samples will be tentatively released via this link¹¹1https://github.com/datvo06/MSAU to foster the reproducibility of this work.

Loss function. The multi-task loss used in the training process is a weighted sum of two smaller losses: the ending loss of the final output and the auxiliary loss of the intermediate output. We adopt the Focal Loss [Lin et al.(2017)Lin, Goyal, Girshick, He, and Dollar] to resolve the class imbalance problem during training. The auxiliary loss is computed from the ground-truth keys mask, after removing all the values. The corresponded weight are $\gamma=0.4$ for the auxiliary loss, and $1-\gamma$ is applied to the ending loss.

Data augmentation. In order to increase the robustness of the model against varied layouts, we perform various types of data augmentation. The text-boxes in the input document are re-scaled so that the median text height is 3 pixels in the char-grid. We use the following steps for augmentation: random character replacement to mimic OCR errors, random text box shifting, random affine transformation, random background padding.

Our in-house datasets comprise of challenging document images with a variety of layouts. The images are also corrupted by several distortions such as noises, blurs, missing strokes due to the scanning process or poor paper quality. These conditions often appear in a practical document processing pipeline, thus make our datasets a good benchmark for information extraction methods. Ground-truth for the data was annotated manually by human operators, which includes the location and content of text-boxes as well as the key information boxes that need to be extracted.

Japanese invoices dataset consists of 261 invoices in Japanese from several vendors. Each invoice contains 16 key-value pairs to extract. Example keys are date_issue, sender_name, receiver_name, total_amount, tax, item_name, item_amount, etc. A few sample images are shown in the supplement (Figure 5). This dataset features some notable challenges such as: mixed handwriting and printed text, random element placements, low resolution scanning.

Insurance medical receipts dataset consists of 200 medical receipts in Japanese with varied formats. Each document has 12 keys to be extracted: date_issue, billing_period, insurance_amount, patient_co-payment_ratio, surgery_fee, hospitalization_fee, etc. Examples of this dataset can be accessed via the supplement as well (see Figure 6). The challenges of this data include long-range key-value correlation, multiple keys relation, skewed images.

Both datasets are divided into 70% for training and 30% for testing.

The proposed network was implemented using Tensorflow library. We perform experiments on a server equipped by an Nvidia Quadro M4000 (8GB memory GPU). All models were trained for 200 epochs with a mini-batch size of 4. The RMSProp method [Tieleman and Hinton(2012)] is used for optimizing the model. The learning rate is set to $0.001$ at the beginning and decreases every 10 epochs with a polynomial decay of 0.9.

Baselines. In our experiment, we used two U-Net variants with dilated convolution [Papandreou et al.(2015)Papandreou, Kokkinos, and Savalle, Chen et al.(2016)Chen, Papandreou, Kokkinos, Murphy, and Yuille] and residual bottleneck in the downsampling path, which is similar to the encoder backbone of our architecture. The first variant has $n_{downsampling}=5$, $res\_depth=2$ and $C=16$. The second variant has $n_{downsampling}=6$, $res\_depth=3$. We abbreviate them as $unet\_small$ and $unet\_big$, respectively. These two networks are very strong baselines due to the proven effectiveness of their sub-components. Additionally, the same architecture was also presented in [Katti et al.(2018)Katti, Reisswig, Guder, Brarda, Bickel, Höhne, and Faddoul], which allows us to make a direct comparison to one of the current state-of-the-art methods. Table 2 summarized the number of parameters of our network and the baselines. As will be shown shortly, although our $msau\_small$ uses 40% less parameters, it is still able to outperform $unet\_big$ by a large margin in term of accuracy.

Metrics. To compare the performance across models, we use three metrics: mean Intersection-over-Union (mIOU), mean pixel accuracy (pix acc) and the box F1-score (F1-score). We note that we leveraged a modified metric, so-called box F1-score, to evaluate the correctness of the predicted bounding box for each key field. The bounding box of each field is constructed from the connected component analysis on the output segmentation mask. A predicted box is marked as a correct detection if its IoU ratio with one of the ground-truth boxes is larger than a certain threshold $0.8$. This metric is suitable for testing the model on practical OCR-ed document where the ground-truth mask may slightly misalign with the generated char-grid from the OCR text.

Overall results. We first compare the performance of our $msau$ models to the baseline architectures. Table 3 shows the results of all models that have been deployed on the Japanese invoices testset. As we can see, the MSAU-nets surpassed the baseline U-Net models in all metrics by a significant margin. The base $msau$ model improved the $mIOU$ score by $7\%$ compare to the $unet\_big$, with only 60% the number of parameters. The $msau\_big$ further increased this gap to 9%, leading to an remarkable 40% error reduction. This demonstrates the effectiveness and efficiency of our architectures. $F1$-$score$ is also largely improved, with a 9% increase from $unet\_big$ to $msau\_big$ and 7% increase from $unet\_big$ to $msau$.

Similar achievements are expected while testing on the Insurance medical receipts datasets (Table 3). The MSAU models continue to outperform the baseline U-Nets consistently. The $msau\_big$ model leads to an obvious 8% improvement in $mIOU$ from the $unet\_big$. The base $msau$ also demonstrated good parameter efficiency by delivering better performance compared to $unet\_big$ (5%) in both $mIOU$ and $F1$-$score$.

We further evaluate our models on the two datasets with output text from an OCR engine (Tesseract [Smith(2007)]). As can be seen in Table 2, the $msau\_big$ model shows its robustness to OCR errors with 5% and 4% drops in $F1$-$score$, while the baseline $unet\_big$ has 8% and 7% drops respectively. Furthermore, the gap between $unet\_big$ and $msau\_big$ is enlarged to 9$\sim$10% in the OCR-ed documents. This shows that our method is able to detect the errors produced by OCR step by a successful rate.

For a more comprehensive comparison, we show our model results across different levels of synthetic OCR errors (randomly replacing/removing character with probability $e$) in Figure 4(a). The impact of OCR errors is less severe on our MSAU-net as the proposed model still achieve more than 70% $F1$-$score$ in an extreme condition where character error rate is as high as 25%.

As we can see, the base CU-Net model (4) improves the mIOU score by 2% compared to the base U-Net (1). It also surpasses the UNet+NL (2) configuration. The Non-local (self-attention) block helps the base model gained 1% and 2% in $mIOU$, respectively, in cases of the U-Net (2) and CU-Net (6) baseline. The CU-Net configuration gained much greater improvement in comparison with the base U-Net when adding the box convolution (3% and 1% $mIOU$ in (3) and (8)), respectively. This demonstrates the effectiveness of the component, especially when combined with the CU-Net for modeling spatial relation in the char-grid. Notably, with the multi-stage training (5) plus the box convolution, configuration (7) improves the overall $mIOU$ by 4%. This is a significant achievement because in order to reach the same number, $unet\_small$ has to transform to $unet\_big$ with more 40% parameters. The effect of box convolution can be further analyzed in the Figure 5, in which we can see that the improved result come from the larger receptive field by expanding the box convolution filter dimension.