Tag, Copy or Predict: A Unified Weakly-Supervised Learning Framework for Visual Information Extraction using Sequences

1 Introduction

With the fast development of information interaction, document intelligent processing [Ferilli et al., 2011] has attracted considerable attention. As an important part of it, visual information extraction (VIE) technique has been integrated into many real-world applications.

The existing VIE methods usually first organized text blocks (text bounding boxes and strings, which were provided by the ground truth or parsed by an OCR system) into plain texts according to the reading order and utilized effective encoding structures such as [Katti et al., 2018; Liu et al., 2019a;

∗Corresponding author

Xu et al., 2020] to extract the most distinguishable representations for each input token from multi-sources. After this, a sequence tagging model like [Lample et al., 2016] was trained with token-level category supervision.

However, the token-level category supervision expends great annotation costs and may be exposed to label ambiguity. Given a document image as shown in Figure 1 (a), the most widely used annotation scheme is to label the bounding box and string of each utterance, and further point out which category does each token/box belongs to, as shown in Figure 1 (b). In this way, a heuristic label assignment procedure is needed to train the aforementioned tagging model, of which the core idea is matching the detected boxes and recognized transcriptions with the given annotations and then assign label to each token/box of OCR results. However, this procedure may encounter problems from mainly two aspects. First, wrong recognition results will bring troubles to the matching operation, especially for key information sequences. Second, the repeated contents will bring label ambiguities. As shown in Figure 1(a) and (b), three values with same content can be regarded as the answer of the key Total Amount. In most cases, it is hard to establish a uniform annotation specification to determine which one should be regarded as ground truth.

To address the aforementioned limitations, in this paper, we propose an end-to-end weakly-supervised learning framework, which can supervise the decoding process directly using the target key information sequences. The benefits it brings are mainly two-folds: first, it greatly saves the annotation costs, as shown in Figure 1 (c), and shortens the training process by skipping the matching between OCR results and the ground truth; second, our method solves the label ambiguity problem by automatically learning the alignment between OCR results and ground truth, which can adaptively distinguish the most likely one in the repeated contents. In addition, we also propose a flexible decoder, which is combined with our weakly-supervised training strategy and have two switchable modes - Copy or Predict Mode (TCPN-CP) and Tag Mode (TCPN-T), to balance the effectiveness and efficiency. In TCPN-CP, our decoder can generate key information sequences by copying a token from the input or predicting one in each time step, which can both retain novel contents in input and correct OCR errors. And in TCPN-T, our decoder can directly label each token's representations into a specific category in a single forward pass, which maintains the fast speed. It is notable that our decoder only needs to be trained once to work in different modes.

Besides, we propose an efficient encoder structure to simultaneously model the semantic and layout information in 2D OCR results. In our design, the semantic embeddings of different tokens in a document are re-organized into a vector matrix I ∈ RH×W×d (here H, W and d are the height dimension, the width dimension and the number of channels, respectively), which we called TextLattice, according to the center points of token-level bounding boxes. After this, we adopt a modified light-weight ResNet [He et al., 2016] combined with the U-Net [Ronneberger et al., 2015] architecture to generate the high-level representations for the subsequent decoding. It is notable that the most relevant work of our encoding method was CharGrid [Katti et al., 2018], which first used CNN to integrate semantic clues in a layout. However, it initialized an empty vector of the same size as the original document image, and then repeatedly filled each token's embedding at every pixel within its bounding box. This simple and direct approach may lead to the following limitations: 1) Larger tokens would be filled in more times, and it might result in the risk of category imbalance; 2) A pixel could completely represent a token, and repeated filling would waste extra cost; 3) The lack of concentration of information would make it difficult for the network to capture global clues. By comparison, our method greatly saves computing resources while maintains the space information of document. For example, for a receipt image in SROIE [Huang et al., 2019] benchmark whose size is generally more than 1600 × 800, the side length of I after our processing is less than 100. In this way, both holistic and local clues can be captured, the relative position relationship between utterances are retained, and distance can be perceived in a more intuitive way through the receptive field.

Experiments on the two public SROIE and EPHOIE [Wang et al., 2021] benchmarks demonstrate that our method shows competitive and even new state-of-the-art performance. Our main contributions can be summarized as follows:

• We propose an efficient 2D document representation called TextLattice, and the corresponding light-weight encoder structure.
• We propose a flexible decoder which has two inference modes - TCPN-CP for OCR error correction and TCPNT for fast inference speed.
• We propose a weakly-supervised learning framework which guides decoding process directly using the key information sequences. This greatly saves the annotation costs and avoids label ambiguity.
• Our method achieves competitive performances even compared with the setting of token-level full supervision, which totally proves its superior.

2 RelatedWork

Early works of visual information extraction mainly utilized rule-based [Muslea and others, 1999] or templatebased [Huffman, 1995] methods, which might tend to poor performance when the document layout was variable. With the development of deep learning, more advanced researches commonly extracted a feature sequence from the input plain text and used token-level supervision to train a sequence tagging model. [Lample et al., 2016] first used a bidirectional long short-term memory [Hochreiter and Schmidhuber, 1997] (BiLSTM) network to model sequential information and a conditional random field (CRF) layer to decode the optimal classification path. Most of the follow-up works mainly focused on the feature encoding part: [Liu et al., 2019a], GraphIE [Qian et al., 2019] and PICK [Yu et al., 2020] tried to use graph neural networks (GNNs) to extract node embeddings for better representation. LayoutLM [Xu et al., 2020] embedded multi-source information into a common feature space, and utilized a BERT-like [Devlin et al., 2019] model for feature fusion. TRIE [Zhang et al., 2020] and VIES [Wang et al., 2021] proposed the end-to-end VIE methods directly from image to key information, which introduced multi-head self-attention [Bahdanau et al., 2015] to integrate multimodal clues during encoding.

As for the decoding part, EATEN [Guo et al., 2019] first utilized sequence-level supervision to guide training. It generated feature maps directly from document image, and used several entity-aware attention-based decoders to iteratively parse the key information sequences. However, its efficiency could be significantly reduced as the number of entities increased, and it can only process simple and fixed layout since it had to overcome the difficulties of both OCR and VIE at the same time. When the given text blocks were accurate or directly the ground truth but the model still performed inference by step-by-step prediction, it might greatly slow down the speed and lead to the severe over-fitting problem of sequence generation due to the lack of corpus in VIE task.

3 Method

Here we provide the details of the proposed TCPN. First we describe the approach of generating TextLattice, and how to encode higher-level features. Next we introduce details of our switchable decoder and weakly-supervised training strategy. Finally, we explain when and how to inference in different modes. Figure 2 gives an overview of our approach.

3.1 Document Representation

In this section, we introduce how to re-organize the OCR results into our 2D document representation - TextLattice. The whole process can be summarized as: 1) We first normalize y coordinates of detected bounding boxes Bu , sort Bu from top to bottom and left to right, and utilize heuristic rules to modify y coordinates to divide Bu into multiple rows; 2) Then, Bu is divided into token-level Bt according to lengths of the recognized strings Su ; 3) Next, the x coordinates of Bt are also normalized and modified to avoid information loss caused by overlapping; 4) Finally, we initialize an all-zero matrix I ∈ RH×W×d where W and H are determined by the ranges of x and y coordinates of Bt , and fill in I according to the correspondence between token-level center points and d-dimensional token embeddings. The detailed procedure is shown in Appendix.

3.2 Feature Encoding

After acquiring I, we adopt ResNet [He et al., 2016] as CNN encoder to capture more holistic features. The U-Net [Ronneberger et al., 2015] structure is also combined to restore the down-sampled features to the same size as the input I and adaptively fuse both local and global clues extracted under diverse receptive fields. Vanilla ResNet adopts a 7 × 7 Conv2d as the conv1 layer to capture association between local pixels in an RGB image. However, it may not be applicable in our scenario since the features of adjacent tokens also need to be separable, instead of high fusion of features in the first few layers. To this end, we replace conv1 with a 3 × 3 Conv2d and remove the original maxpool layer. Thanks to the efficient design of TextLattice, both speed and superiority can be retained.

In order to further preserve the location clues, inspired by CoordConv [Liu et al., 2018], two extra channels are concatenated to the incoming representation I, which contain horizontal and vertical relative location information in the layout of range from -1 to 1. The whole procedure of feature encoding can be formulated as:

Here, ⊕ is the concatenation operator, I0 ∈ RH×W×2 is the extra coordinate vector. Since the output of CNNs has the same size as I, we add them together as a residual connection. Finally, the features at the center points of token-level bounding boxes Bt are retrieved to form F ∈ RN×d , where N is the number of tokens. We regard the rest pixels as useless and discard them directly for calculation efficiency.

3.3 Weakly-Supervised Training

As shown in Figure 2, the VIE task can be regarded as a setto-sequence problem after feature encoding, since F is orderindependent. We introduce the class embedding C ∈ Rd to control the category of information parsed by the decoder, which is taken from a pre-defined trainable look-up table. Given C, our attention-based decoder takes it into account at each time step and iteratively predicts target sequence. Such novel design avoids class-specific decoders, alleviates the shortage of isolated class corpus, and decouples the serial correlation between different categories in the traditional sequence tagging model into parallel.

When generating sequences, we need the model to be able to switch between copying tokens from input or directly predicting ones. The copying operation make the model be able to reproduce accurate information and retain novel words, while the predicting operation introduces the ability of correcting OCR errors. Inspired by [See et al., 2017], which implemented a similar architecture for abstractive text summarization, our model recurrently generates the hidden state st by referring to the current input tokens xt and the context vector in the previous step F ∗ t-1 :

Here, α is the attention score where the historical accumulated values are also referenced during calculation. All Ws and bs are learnable parameters.

Then, the probability distribution of tokens in a fixed dictionary P pred t is calculated and a copy score pcopy t is generated as a soft switch to choose between different operations in each time step t:

Pt(k∗ ) is the probability score of token k∗ in time step t, where k∗ is the current target token. LSt is the sequence alignment loss function. In this way, our method acquires the ability to produce out-of-vocabulary (OOV) tokens, and can adaptively perform optimal operations.

As of now, our method can be seen as a sequence generation model trained with sequence-level supervision. However, it is notable that since the class embedding C of entity category c is given, when the model decides to copy a token ki from the input at a step, ki's feature vector in F should be also classified as c by a linear classifier. More generally speaking, our method can first learn the alignment relationship, and then train a classifier using the matched tokens. This novel idea enables our approach the ability of supervising the sequence tagging model. We adopt a linear layer to model the entity probability distribution, which can be formulated as:

It is worth noting that, equation (11) - (13) do not train the tokens which do not belong to any key information sequences. The neglect of negative samples may lead to severe defect that all input tokens will be classified as positive. Thus we construct an extra auxiliary loss function LNt for negative sample suppression:

Here, Pt(c) indicates the sum of the probabilities belong to entity category c of all input tokens, and Length(Sc) is the length of current target sequence Sc. The main purpose of LNt is to limit the number of input tokens classified as c to be less than or equal to the actual number. This simple but effective design greatly improves the performance of the model in Tag Mode.

In summary, the final integrated loss function Lt is the weighted sum of multiple components mentioned above:

where λS, λC and λN are trade-off hyper-parameters.

3.4 Inference

In this section, we explain when and how to implement inference process in different modes. It is worth noting that, since class embeddings are sent into the decoder in the form of a batch, key information sequences of different categories of the same document can be generated under different modes according to the entity-specific semantic characteristics.

In most real-world scenarios, OCR results cannot be flawless. In this regard, users can switch our decoder to Copy or Predict Mode as described in equation (3) - (9) to supplement missing or wrong tokens. This mode is more suitable for sequences of categories with strong semantic relevance.

Thanks to the auto-alignment property of the proposed weakly-supervised training strategy, the decoder can also directly perform sequence tagging in a single forward pass in Tag Mode using equation (14). It prefers to extremely few OCR errors or categories of weak semantic correlation between adjacent contents.

4 Experiment

4.1 Implementation Details

We adopt ResNet-18[He et al., 2016] as backbone in feature encoding, and use BiLSTM in attention mechanism of decoder. The number of channels d is set to 256, the hyperparameters λS, λC and λN are all set to 1.0 in our experiments empirically. We set batch size as 4 and perform 450 training epochs. The learning rate is initialized as 1.0 with ADADELTA [Zeiler, 2012] optimization and decreased to a tenth for two times in 300 and 400 epochs.

4.2 Ablation Study

In this section, we evaluate the influences of components of the proposed TCPN on the public EPHOIE benchmark.

Comparison between Different Encoding Structures

We organize OCR ground truth into plain texts and use different encoding structures to generate representations. Then a sequence tagging model is trained utilizing the official tokenlevel category annotations. We mainly adopt the following models for comparison: BiLSTM: A bidirectional LSTM adopted in [Lample et al., 2016]; GAT: A graph attention network (GAT) adopted in [Liu et al., 2019a]; BERT-like: A BERT-like model similar to LayoutLM[Xu et al., 2020]. Since vanilla LayoutLM is trained on English corpus, the pretrained weights for Chinese in [Cui et al., 2020] are loaded for fair comparison; Chargrid: The document modeling method introduced by Chargrid[Katti et al., 2018].

The comparison results are given in Table 1. BiLSTM perceives sequential clues well, but it cannot effectively model location space in 1D form; GAT can adaptively fuse useful features using attention mechanism. However, the ability of capturing location clues highly depends on the way of feature embedding; BERT-like can perform forward calculation in parallel, and since the pretrained weights are loaded, it achieves satisfactory performance; Chargrid establishes input matrix using a more direct way, which means that both robustness and efficiency cannot be guaranteed. It is notable that TextLattice(Ours) achieves superior performance and maintains the fastest speed, which fully proves its efficiency. Our method has a more direct and sensitive perception of location clues than position embedding in GAT or BERT-like, and ensures a higher degree of information concentration than Chargrid. The fully parallel scheme also greatly contributes to the leading speed.

Effects of Components in TextLattice

We also conduct experiments to verify the effectiveness of different components in our encoding structure, such as CoordConv, the U-Net structure and the residual connection in equation (1). It can be seen in Table 2 that each design has a significant contribution to the final performance. Although

(a)

(b)

Table 3: Performance comparison on (a) EPHOIE and (b) SROIE under ground truth setting. 'Token-level Fully-Supervised' indicates the token-level category annotation, and 'Sequence-level WeaklySupervised' means the key information sequence annotation.

CNN can capture the relative position relationship, CoordConv can further provide the global position clues relative to the whole layout, which brings higher discernibility; we also try to use ResNet only where all stride and the U-Net structure are removed to perform feature encoding. However, the performance decreases obviously, which indicates the importance of semantic feature fusion under different receptive fields; Residual Connection gives model the chance to directly receive token-level semantic embedding, which further improves the performance.

4.3 Comparison with the State-of-the-Arts

We compare our method with several state-of-the-arts on the SROIE and EPHOIE benchmarks. The following Ground Truth Setting indicates that the OCR ground truth is adopted, while End-to-End Setting indicates the OCR engine result.

Results under Ground Truth Setting

As shown in Table 3, our method exhibits superior performance on both SROIE and EPHOIE in the case of tokenlevel full-supervision, which totally proves the effectiveness of our feature encoding method. Furthermore, the results under sequence-level weakly-supervised setting achieve competitive performance. This fully confirms the superiority of our learning strategy, which can model the correspondence between the input tokens and the output sequences.

Compared with SROIE, EPHOIE usually has less content and more token types, which reduces the difficulty of learning alignment. Relatively, since a receipt in SROIE often contains abundant tokens and the same token may appear repeatedly, which may lead to the alignment confusion, the gap between full- and weak-supervision is further widened.

Results under End-to-End Setting

We adopt BDN[Liu et al., 2019b] as text detector and CRNN[Shi et al., 2016] as text recognizer, and train them using the official annotations to get OCR results. It is worth noting that since there may inevitably exist errors in OCR results, all of our experiments under end-to-end setting are trained using our weakly-supervised manner, which avoids the matching process between OCR results and ground truth. The performances are shown in Table 4. Our method shows new state-of-the-art performance in every mode.

It can be inferred that an important basis for choosing TCPN-CP or TCPN-T mode is the richness of semantics and corresponding corpus. On SROIE, TCPN-CP obviously outperforms TCPN-T, which mainly benefits the ability to error correction; however, on EPHOIE, although both modes outperform counterparts, TCPN-T beats TCPN-CP by a large margin, where the main reason should be the diversity of Chinese tokens and the resulting lack of corpus.

Results on A Business License Dataset

In order to further explore the effectiveness of our framework in real-world applications, we collect an in-house dataset of business license. It contains 2331 photos taken by mobile phone or camera with real user needs, and most of images are inclined, distorted or the brightness changes dramatically. We randomly select 1863 images for training and 468 images for testing, and there are 13 types of entities to be extracted. Furthermore, the OCR results are generated by our off-theshelf engines, which definitely contains OCR errors due to the poor image quality.

The detailed results are shown in Table 5. Our endto-end weakly-supervised learning framework outperforms traditional rule-based matching method by a large margin, which can also greatly reduce the annotation cost. Compared with TCPN-T, the implicit semantic relevance learned by TCPN-CP can further boost the final performance by correcting OCR errors. Some qualitative results are shown in Appendix.

5 Conclusion

In this paper, we propose a unified weakly-supervised learning framework called TCPN for visual information extraction, which introduces an efficient encoder, a novel training strategy and a switchable decoder. Our method shows significant gain on EPHOIE dataset and competitive performance on SROIE dataset, which fully verifies its effectiveness.

Visual information extraction task is in the cross domain of natural language processing and computer vision, and our approach aims to alleviate the over-reliance on complete annotations and the negative effects caused by OCR errors. For future research, we will explore the potential of our framework through large-scale unsupervised data. In this way, the generalization of encoder, the alignment capability of decoder and the performance of our TCPN-CP can be further improved.

Acknowledgments

This research is supported in part by NSFC (Grant No.: 61936003, 61771199), GD-NSF (no.2017A030312006).

Appendix

A TextLattice Generation

Algorithm 1 shows the detailed procedure of our TextLattice Generation Algorithm introduced in Section 3.1.

B Qualitative Results

Some qualitative results are shown in Figure 3, 4 and 5.

Algorithm 1 TextLattice Generation Algorithm

Input: Utterance-level detected bounding boxes Bu and corresponding recognized strings Su

Parameter: Normalization threshold Rt and ratio Rr .

Output: TextLattice I

• 1: Calculate Rh = Average height of Bu .
• 2: Normalize y coordinates of Bu using Rh .
• 3: Sort Bu from top to bottom, and then from left to right.
• 4: Let Y = unique values of y coordinates of center points of Bu sorted from small to large, a list ˜ Y = [Y [1], ].
• 5: for i from 2 to Count(Y ) do
• 6: if (Y [i] - Y [i - 1]) ≤ Rt then
• 7: dY = 1.
• 8: else
• 9: dY = max(1, (Y [i] - Y [i - 1])/Rr ).
• 10: end if
• 11: Append ˜ Y [-1] +dY into ˜ Y .
• 12: end for
• 13: Modify y coordinates of center points of Bu corresponding to Y → ˜ Y .
• 14: Split utterance-level boxes Bu into token-level boxes Bt according to lengths of Su .
• 15: Calculate Rw = Average width of Bt .
• 16: Normalize x coordinates of Bt using Rw .
• 17: for ˆ y in Y do
• 18: Let ˆX = x coordinates of boxes Bt′ whose y coordinates of center points equal to ˆ y.
• 19: Sort ˆX from small to large, let ˜X = [ ˆX[1], ]
• 20: for i from 2 to Count( ˆX) do
• 21: if ( ˆX[i] - ˆX[i - 1]) ≤ Rt then
• 22: dX = 1.
• 23: else
• 24: dX = max(1, ( ˆX[i] - ˆX[i - 1])/Rr ).
• 25: end if
• 26: Append ˜X[-1] +dX into ˜X.
• 27: end for
• 28: Modify x coordinates of center points of Bt′ corresponding to ˆX → ˜X.
• 29: end for
• 30: Combine Su in the same order as Bu to get St .
• 31: Let xmin, xmax, ymin, ymax = minimal and max values of x and y coordinates of center points of Bt , initialize an all-zero matrix I ∈ R(ymax-ymin)×(xmax-xmin)×d . Subtract xmin, ymin from the coordinates of Bt .
• 32: Fill in I according to the correspondence between Bt and d-dimensional token embeddings of St .
• 33: return I.

Figure 5: Qualitative results of TCPN on the in-house Business License dataset. For the purpose of privacy protection, part of the correctly parsed information is masked.