SlowGuess/NAACL_2022_OCR huggingface dataset:
Haoyu Cao $^{1}$ , Jiefeng Ma $^{2*}$ , Antai Guo $^{1}$ , Yiqing Hu $^{1}$ , Hao Liu $^{1}$ , Deqiang Jiang $^{1}$ , Yinsong Liu $^{1}$ , Bo Ren $^{1}$
Tencent YouTu Lab
{rechycao,ankerguo,hooverhu,ivanhliu,dqiangjiang}@tencent.com
jfma@mail.ustc.edu.cn {jasonysliu,timren}@tencent.com
Document Information Extraction (DIE) has attracted increasing attention due to its various advanced applications in the real world. Although recent literature has already achieved competitive results, these approaches usually fail when dealing with complex documents with noisy OCR results or mutative layouts. This paper proposes Generative Multi-modal Network (GMN) for real-world scenarios to address these problems, which is a robust multimodal generation method without predefined label categories. With the carefully designed spatial encoder and modal-aware mask module, GMN can deal with complex documents that are hard to serialized into sequential order. Moreover, GMN tolerates errors in OCR results and requires no character-level annotation, which is vital because fine-grained annotation of numerous documents is laborious and even requires annotators with specialized domain knowledge. Extensive experiments show that GMN achieves new state-of-the-art performance on several public DIE datasets and surpasses other methods by a large margin, especially in realistic scenes.
Document Information Extraction (DIE) aims to map each document to a structured form consistent with the target ontology (e.g., database schema), which has recently become an increasingly important task. Recent research (Xu et al., 2020, 2021; Wang et al., 2021a; Zhang et al., 2020; Li et al., 2021a) has achieved competitive results for information extraction in the idealized scenario with accurate OCR results, word-level annotations, and serialized document words in reading order. These methods regard DIE as a Sequence Labeling (SL) task. Given OCR results of a document image, the traditional sequence labeling method first serializes
Figure 1: Examples in public DIE benchmarks with practical problems. The three rows from top to bottom in each sub-figure are 1) input images, 2) raw intermediate tags and final results generated by the SL method, and 3) results of our GMN method, respectively. The error parts are marked in red, while the correct parts are in green. Best viewed in color.
words in reading order then classifies each input word into predefined categories.
As shown in Figure 1, multiple challenging problems for practical document understanding still exist in realistic scenes. 1) Document serialization requires pre-composition processing, which is difficult in real scenarios with ambiguous word orders. One entity may be incorrectly divided into multiple entities when the input sequences are sorted by coordinates. 2) OCR results are usually noisy because of inevitable recognition errors. 3) The volume of keys in practical scenarios is generally substantial and expanded frequently. Existing sequence labeling methods could not identify undefined keys. 4) While facing duplicated values, collecting word-level annotations is necessary for sequence labeling methods. However, this is difficult in practical scenarios since they are costly and labor-intensive.
To address the limitations mentioned above, we propose a robust information extraction method
*Work is done during an internship at Tencent YouTu Lab.
Human Language Technologies, pages 3768 - 3778
named Generative Multi-modal Network (GMN) for practical document understanding. Unlike sequence labeling methods that label each input word with a predefined category, we regard DIE as a translation task that translates source OCR results to a structured format (like key-value pairs in this paper). We use UniLM (Dong et al., 2019) as the basic model structure, which is a transformer-based pre-trained network that can handle both natural language understanding (NLU) and natural language generation (NLG) tasks simultaneously. Conditioned on a sequence of source words, GMN generates one word at each time step to compose a series of key-value pairs in them.
Regarding the sequence serialization problem, a novel two-dimensional position embedding method is proposed while the original one-dimensional positional embedding in the transformer is removed because all information in document understanding can be acquired from 2D layouts. In this manner, GMN bypasses the serialization problem. Furthermore, benefiting from the large-scale self-supervised pre-training processed on a vast document collection, GMN can correct OCR errors commonly encountered in practical scenarios. Moreover, using a weakly supervised training strategy that utilizes only key information sequences as supervision, GMN needs no word-level annotations that are indispensable in traditional sequence labeling methods like LayoutLM (Xu et al., 2020) and StructuralLM (Li et al., 2021a).
Experiments illustrate that the proposed GMN model outperforms several SOTA pre-trained models on benchmark datasets, including SROIE and CORD. The contributions of this paper are summarized as follows:
1) We present GMN tailored for the DIE task, which is more applicable for practical scenarios, including lack of word-level annotations, OCR errors as well as various layouts.
2) We propose a layout embedding method and multi-modal Transformer in a decoupling manner, which jointly models interactions between multiple modalities and avoids the reading order serialization.
3) Experiments on public DIE datasets demonstrate that the proposed method not only achieves a substantial performance improvement but also generalizes well to data under practical scenarios with unseen keys.
2 RELATED WORKS
Traditional methods (Esser et al., 2012; Schuster et al., 2013; Riloff, 1993) on DIE tasks rely heavily on predefined rules, templates, and handcrafted features, giving rise to difficulty in generalizing to unseen documents. With the development of deep learning technology, document information extraction methods have recently improved substantially in both performance and robustness. These deep learning-based methods can be classified into three categories: textual content-based methods, multi-modal-based methods, and pre-trained Transformer-based methods.
Textual content-based methods. Palm et al. (2017); Sage et al. (2019) adopt the idea from natural language processing and use recurrent neural networks (RNN) to extract entities of interest from documents. However, they discard the layout information during the text serialization, which is crucial for document understanding.
Multi-modal-based methods. Some works (Katti et al., 2018; Hwang et al., 2021) take the layout information into consideration and try to reconstruct character or word segmentation of the document. Katti et al. (2018) encode each document page as a two-dimensional grid of characters that represents text representation with a two-dimensional layout. Yu et al. (2020); Majumder et al. (2020); Zhang et al. (2020); Wang et al. (2021b) further integrate image embeddings for better feature extraction. Yu et al. (2020); Tang et al. (2021) represent documents by graphs, with nodes representing word segments and edges either connecting all the nodes or only spatially near neighbors. Convolutional or recurrent mechanisms are then applied to the graph for predicting the field type of each node. However, due to the lack of large-scale pre-training, the robustness and accuracy of the model are relatively limited.
Pre-trained Transformer-based methods. Recently, pre-trained models (Devlin et al., 2019; Liu et al., 2019) show effective knowledge transferability with large-scale training data and various self-supervised tasks. LayoutLM (Xu et al., 2020) first proposes a document-level pre-training framework that semantic and layout information are jointly learned. LayoutLM V2 (Xu et al., 2021) further improves the LayoutLM model by integrating the image information in the pre-training stage. Li et al. (2021a) propose the StructuralLM pre-training approach to exploit text block information. Methods
(a) The overall architecture of our generative multi-modal network, MD-Bert in encoder and decoder are with parameter sharing.
(b) Detail of MD-Bert module, with internal memory updated recursively.
Figure 2: Architecture of our generative multi-modal network. The MD-Bert module is composed of stacked multi-modal encoders, which fuses multi-modal features and iteratively generates structured results in a fixed order.
mentioned above all use one-dimensional position embeddings to model the word sequence, even with two-dimensional layout embeddings are involved, so that the reading order serialization in the document is required, which is challenging or even impossible due to the complex and diverse layout in the real world. What's more, they are all based on the classification of each input text segment to predefined labels, which means fine-grained annotations are indispensable and lack the ability to correct error OCR results.
On the contrary, the proposed GMN relies on a two-dimensional position embedding to bypass the serialization process and cross-modality encoders in a decoupling manner to model the layout information and the relative position of a word within a document simultaneously.
3 METHODOLOGY
In this section, we first introduce the overall architecture of GMN, followed by illustrating multimodal feature extraction, generative pre-training model with multi-modal decoupling in detail, respectively.
GMN aims at constructing an enhanced Transformer-based translator architecture for DIE for converting the document to structured,
machine-readable data. An overview of the architecture is as shown in Figure 2. It mainly consists of two parts: the multi-modal feature extraction module and stacked cross-modality module named MD-Bert (Modal Decoupling Bert), which simultaneously serves as encoder and decoder following the design of UniLM.
The whole process can be summarized as 1) Multi-modality embeddings of source inputs are extracted through an advanced OCR engine and a small CNN; 2) The extracted features from different modalities are fused as "multi-modal embeddings" through MD-Bert along with memory updating for each layer at each time step; 3) Next, MD-Bert output the encoding results by applying token prediction on multi-modal embeddings; 4) Finally, MD-Bert recursively generates structured results by taking multi-modal embeddings and accumulative memory as inputs until a terminator [SEP] is predicted.
3.2 Multi-Modal Feature Extraction
Based on the multi-modal information, including semantics, layout, and vision, we propose a unified layout embedding method named Cord2Vec which simultaneously encodes sequence information and spatial information to avoid complex reading order serialization.
3.2.1 Semantic Embedding
Intuitively, semantic contents are reliable signals to extract valuable information. The semantic content of each text fragment is acquired from the results of the OCR engine for practical application scenarios. After text fragments are acquired and tokenized, the start indicator tag [BEG] is added in front of the input token sequence, and the end indicator tag [SEP] is also appended to the end. Extra padding tag [PAD] is used to unify the length of sequence with predefined batch length  $L$ . In this way, we can get the input token sequence  $S$  as
$$
S = \left[ \left[ \mathrm {B E G} \right], t _ {1}, \dots , t _ {n}, \left[ \mathrm {S E P} \right], \left[ \mathrm {P A D} \right], \dots \right], | S | = L \tag {1}
$$
Here,  $t_i$  refers to  $i$ -th token in OCR texts. Moreover, though the sequence length of the input token sequence is fixed during training, GMN can handle variable lengths when making the inference due to the novel positional embedding method.
3.2.2 Layout Embedding
DIE task is a typical two-dimensional scene in which relative positions of words are essential evidence. While the reading order serialization is challenging, we propose Cord2Vec, a unified embedding method that fully utilizes spatial coordinates rather than one-dimensional sequence order information to bypass this problem.
As for the source input part, we normalize and discretize all coordinates to the integer in the range of  $[0, \alpha]$ , here  $\alpha$  is the max scale which is set to 1000 in our experiment. Then corner coordinates and edge lengths of each text fragment are gained using corresponding bounding boxes. In order to enhance the tokens' interaction in the same box, two tuples  $(x_0, x_1, w)$ ,  $(y_0, y_1, h)$  are used to represent the layout information. Here,  $(x_0, y_0)$  and  $(x_1, y_1)$  are the top-left and bottom-right coordinates of each token, and  $w$  is the average width of tokens in the same box while  $h$  representing box height. Such embedding represents both the layout and the word order information. As for target tokens generated by GMN which does not have the real coordinate, the Cord2Vec assumes each token is tied in the grid of  $[W_{grid}, H_{grid}]$  with row-first principle, and each token occupies a pixel with a width and height of 1. After the layout information is acquired, we use two embedding layers to embed x-axis features and y-axis features separately as stated in Equation 2.
$$
\begin{array}{l} \begin{array}{l} X _ {i} = \operatorname {P o s E m b 2 D} _ {\mathrm {x}} \left(x _ {0}, x _ {1}, w\right), \\ Y _ {i} = \operatorname {P o s E m b 2 D} _ {\mathrm {x}} \left(\dots , w\right) \end{array} \tag {2} \\ Y _ {i} = \operatorname {P o s E m b 2 D} _ {\mathrm {y}} \left(y _ {0}, y _ {1}, h\right) \\ \end{array}
$$
Here,  $PosEmb2D_{\mathrm{x}}$  and  $PosEmb2D_{\mathrm{y}}$  are the position embedding function which takes coordinate as input. Each input element is embedded separately and then added together with an element-wise function. Note that the placeholder such as [PAD] can be treated as some evenly divided grids, so their bounding box coordinates are easy to calculate. An empty bounding box  $X_{\mathrm{PAD}} = (0,0,0),Y_{\mathrm{PAD}} = (0,0,0)$  is attached to [PAD], and  $X_{\mathrm{SEP}} = (0,w,w),Y_{\mathrm{SEP}} = (0,h,h)$  is attached to other special tokens including [BEG] and [SEP].
3.2.3 Visual Embedding
We use ResNet-18 (He et al., 2016) as the backbone of the visual encoder. Given a document page image  $I$ , it is first resized to  $W * H$  then fed into the visual backbone. After that, the feature map is scaled to a fixed size by average-pooling with the width being  $W / n$  and height being  $H / n$ ,  $n$  is the scaling scale. Finally, RoI Align (He et al., 2017) is applied to extract each token's visual embedding with a fixed size. The visual embedding of the  $i$ -th token is denoted by  $v_{i} \in (v_{1}, v_{2}, v_{3}, \ldots, v_{L})$ . For source input, visual embedding can be represented as
$$
v _ {i} = R O I A l g n (C o n v N e t (I m a g e), P o s _ {i}) \tag {3}
$$
Here,  $Pos_{i}$  stands for the position of  $i$ -th token, and ConvNet is a convolutional neural network serving as feature extractor in terms of input image, and then ROAlign takes the image feature and location as input, and extracts the corresponding image features. Note that the [BEG] token represents the full image feature, and the other special tokens, as well as output tokens, are attached to the default null image feature.
3.3 Generative Pre-training Model
3.3.1 Model Structure
In order to learn more general features and make full use of the pre-training data, we propose a unified encoder-decoder module named MD-Bert which is composed of stacked hierarchical multimodal Transformer encoders. The context of input tokens is from OCR result during the encoding stage, while already decoded tokens are also included in the decoding stage. To solve this problem, inspired by UniLM, we use masking to control which part of context the token should attend to when computing its contextualized representation, as shown in Figure 3. The input features of semantics, layout and computer vision are mapped to hidden states by:  $\boldsymbol{F}_0^S = \{f_1^S,\dots ,f_N^S\} ,\boldsymbol{F}_0^X =$
Figure 3: Modal-aware Mask Module: modalities fusion in decoupling manner. The sub-module for processing X are with parameter sharing for each layer, the same for Y and Visual sub-module.
$$
\begin{array}{l} \left\{f _ {1} ^ {X}, \ldots , f _ {N} ^ {X} \right\}, F _ {0} ^ {Y} = \left\{f _ {1} ^ {Y}, \ldots , f _ {N} ^ {Y} \right\}, F _ {0} ^ {V} = \\ \left\{f _ {1} ^ {V}, \ldots , f _ {N} ^ {V} \right\}, \text {a l i n e a r m a p p i n g i s a s f o l l o w s :} \end{array}
$$
$$
\begin{array}{l} f _ {i} ^ {S} = \boldsymbol {W} _ {S} x _ {\mathrm {i}} ^ {S}, f _ {i} ^ {X} = \boldsymbol {W} _ {X} x _ {\mathrm {i}} ^ {X}, \\ f _ {i} ^ {Y} = \boldsymbol {W} _ {Y} x _ {\mathrm {i}} ^ {Y}, f _ {i} ^ {V} = \boldsymbol {W} _ {V} x _ {\mathrm {i}} ^ {V} \end{array} \tag {4}
$$
where matrices  $W_{S} \in \mathbb{R}^{d_{h} \times d_{\mathrm{S}}}$ ,  $W_{X} \in \mathbb{R}^{d_{h} \times d_{\mathrm{X}}}$ ,  $W_{Y} \in \mathbb{R}^{d_{h} \times d_{\mathrm{Y}}}$ ,  $W_{V} \in \mathbb{R}^{d_{h} \times d_{\mathrm{V}}}$  are used to project features into hidden-state in  $d_{h}$  dimensions. MD-Bert takes  $F^{\sim}$  and memory  $M^{\sim}$  of history state as input, generates output and update memory  $M^{\sim}$  step by step. where,  $F^{\sim} \in \{F^{S}, F^{X}, F^{Y}, F^{V}\}$ ,  $M^{\sim} \in \{M^{S}, M^{X}, M^{Y}, M^{V}\}$ ,  $M^{\sim}$  contains the history state of each layer and previous embeddings of the model. In the first timestep,  $M^{\sim}$  is initialized from scratch, input  $T_{0}$  means the full OCR result and MD-Bert acts as bi-directional encoder. For timestep  $t$ , where  $t \in [1, m]$ , the model takes the output of the previous timestep or [BEG] as input, and outputs the current result, MD-Bert acts as uni-directional decoder.
3.3.2 Cross-Modality Encoder
Traditional multi-modal models usually fuse different modal features by adding or concatenating them together, which inevitably introduces the unwanted correlations between each other during self attention procedure, e.g. word-to-position, image-to-word. However, these correlations are harmful to strengthening the model's capability as different data modalities are practically orthogonal, thus we need to design a customized pipeline for each one.
We propose MAMM (modal-aware mask module) encoder, a hierarchical structure multi-modal Transformer model in a decoupling manner that jointly models different modalities. As a consequence, the feature embedding decoupling is decomposed into three embeddings in GMN. The MAMM module follows the design of a basic module in BERT, but replaces the multi-head attention with modal-aware multi-head attention.
It also contains feed-forward (FF) layers, residual connections, and layer normalizations (LN), meanwhile, parameters of modals are not shared. When feeded with different modal content such as semantics, layout and computer vision, MAMM first calculates each modal's attention score separately, then added these attention scores together to get a fusion score, finally use this fusion score to apply masking and following operations on semantic content. As shown in Figure 3, let  $F_{l}^{\sim} = \{f_{1}^{\sim},\dots,f_{N}^{\sim}\}$  be the encoded feature in the  $l$ -th layer.  $F_{0}^{\sim}$  is the vector of the input features as mentioned in Equation 4. Features output by the next layer  $F_{l + 1}^{\sim}$  can be obtained via:
$$
\boldsymbol {F} _ {l - a t t} ^ {\sim} = L N \left(f _ {M A M M} \left(\boldsymbol {F} _ {l} ^ {\sim}\right) + \boldsymbol {F} _ {l} ^ {\sim}\right) \tag {5}
$$
$$
\boldsymbol {F} _ {\tilde {l} + 1} ^ {\sim} = L N \left(f _ {\mathrm {F F}} \left(\boldsymbol {F} _ {\tilde {l} - a t t} ^ {\sim}\right) + \boldsymbol {F} _ {\tilde {l} - a t t} ^ {\sim}\right) \tag {6}
$$
where  $f_{MAMM}(\cdot)$  is the modal-aware mask function defined as
$$
f _ {\mathrm {M A M M}} \left(\boldsymbol {F} _ {l} ^ {\sim}\right) = \operatorname {s o f t m a x} \left(\operatorname {M a s k P r o d} \left(\boldsymbol {F} _ {l} ^ {\sim}\right)\right) v \left(\boldsymbol {F} _ {l} ^ {\sim}\right) \tag {7}
$$
$$
\operatorname {M a s k P r o d} \left(\boldsymbol {F} _ {l} ^ {\sim}\right) = \frac {q \left(\boldsymbol {F} _ {l} ^ {\sim}\right) k \left(\boldsymbol {F} _ {l} ^ {\sim}\right) ^ {\top}}{\sqrt [ 4 ]{d _ {k}}} + f _ {\text {M a s k o p t}} \tag {8}
$$
where  $q(\cdot), k(\cdot), v(\cdot)$  are linear transformation layers applied to the proposals' feature, which represent the query, key and value in attention mechanism accordingly. Benefited from the parameters' sharing among layers with regard to X, Y and Vision, GMN has comparable weights as Bert. Symbol  $d_{k}$  is the number of attention headers for normalization, and the  $f_{Maskopt}$  is the Mask operation which controls the attention between each token. In GMN, we apply full attention on all OCR input tokens, and input tokens of the model for the output structural sequence can attend to the whole inputs as well as tokens that have been decoded which are like auto-regressive encoder. Finally,  $\pmb{F}_{l + 1}^{\sim}$  can be obtained by  $\pmb{F}_{l - att}^{\sim}$  via a feed-forward sub-layer composed of two fully-att connected layers of function  $f_{FF}(\cdot)$ . Hierarchically stacked layers form the multi-modal encoder.
3.3.3 Pre-training Method
Similar to UniLM, three cloze tasks including Unidirectional LM, Bidirectional LM and Sequence-to-Sequence LM are used in the GMN. Meanwhile, we propose NER-LM for better entity correlation extraction. The whole loss function is defined as,
$$
\mathcal {L} = \mathcal {L} _ {\text {u n i - L M}} + \mathcal {L} _ {\text {b i - L M}} + \mathcal {L} _ {\text {s 2 s - L M}} + \mathcal {L} _ {\text {N E R - L M}} \tag {9}
$$
In a cloze task, we randomly choose some WordPiece (Wu et al., 2016) tokens in the input, and replace them with the special token [MASK]. Then, we feed their corresponding output vectors computed by the Transformer network into a softmax classifier to predict the masked token. The parameters of GMN are learned to minimize cross-entropy loss, which is computed using the predicted tokens and the original tokens.
NER-LM is an extension of sequence-to-sequence LM for better integrity constraints on the entity. Given the source segment which includes entity values  $s_1, s_2$ , and the corresponding entity types  $n_1, n_2$  as well as some background sentence e.g.  $b_1, b_2$ , we form the input format as A “[BEG]  $s_1b_1s_2b_2$ [SEP]” and B “[BEG]  $n_1s_1n_2s_2$ [SEP]”. Each token in A can access all others of A, while each token in B can access all tokens of A as well as the preceded tokens in B. The target entity in B is masked for prediction during training.
4 EXPERIMENTS
4.1.1 Pre-training Dataset
Our model is pre-trained on the IIT-CDIP Test Collection 1.0 (Lewis et al., 2006), which contains more than 6 million documents, with more than 11 million scanned document images. Moreover, each document has its corresponding text and metadata stored in XML files which describe the properties of the document, such as the unique identity and document labels. And the NER-LM is pre-trained on the Enron Email Dataset (Klimt and Yang, 2004), which contains 0.5 million emails generated by employees of the Enron Corporation. We follow the organization of the letter and generate the content on the image. The structured information in the letters acts as an entity, such as subject, date, etc.
4.1.2 Fine-tuning Datasets
We conduct experiments on three real-world public datasets, FUNSD-R, CORD and SROIE.
The FUNSD-R Dataset. FUNSD (Jaume et al., 2019) is a public dataset of 199 fully annotated forms, which is composed of 4 entity types (i.e. Question, Answer, Header and Other). The original dataset has both semantic entity extraction (EE) and semantic entity linking (EL) tasks. It's noteworthy that the linking between different entities are complicated, one header entity may have linking to several question entities with more answer entities linked.