Robust Layout-aware IE for Visually Rich Documents with Pre-trained Language Models

Our work falls under the scope of Visually Rich Document information extraction which is a relatively new research topic. We roughly divide the current progress of the approaches on this problem into three categories. The first category is rule-based systems. (Rusinol et al., 2013) describes an invoice extraction system using a number of rules and empirical features, and (d’Andecy et al., 2018) builds a more stable system using tf-idf algorithm and a large number of human-designed features. Other document-level entity extraction systems combine rules and statistical models (Belaïd et al., 2001; Schuster et al., 2013). Although it is possible to craft high-precision rules in some closed-domain applications, rule-based systems are usually associated with extensive human effort and cannot be rapidly adapted to new domains.

The second category is graph-based statistical models, where graphs are used to model the relationships between layout components, such as text boxes. (Santosh and Belaïd, 2013) first performs graph mining in a document with a set of key-fields selected by clients, in order to learn the pattern to extract information in the absence of clients. More recently, graph neural networks are used to capture structural information in visually rich documents: (Liu et al., 2019a) applies graph modules to encode the visual information with deep neural networks and GraphIE (Qian et al., 2019) also assumes that the graph structure is ubiquitous in the text, and applies GCN between the BiLSTM encoder-decoder structure to model the layout information in the document. The limitation of these methods is that they do not have access to pre-trained language models such as BERT and have not explored rich visual information (e.g. font and weight of texts) beyond the position of texts.

The third category is approaches that exploit 2D grid information of characters or words. Chargrid (Katti et al., 2018) models the problem by encoding each document page as a two-dimensional grid of characters and used a fully convolutional encoder-decoder network to predict the class of all the characters and group them by the predicted item bounding boxes. LayoutLM (Xu et al., 2019) appends the language model embeddings with 2D grid information of words to jointly pre-train text and layout. These methods rely on the results of OCR and only model with information on the character and word levels, while valuable sequence-level information in digital-born documents is ignored.

Graph neural networks have attracted increasing attention. GCN (Kipf and Welling, 2017) models graph-structured data based on an efficient variant of convolutional neural networks, which has been proved effective in many NLP tasks. For example, GraphRel(Fu et al., 2019) uses a Bi-GCN to better extract relations by jointly learning named entities and relations. (Liu et al., 2018) models and matches long article through a graph convolutional network to identify the relationship between two articles. (Marcheggiani and Titov, 2017) encodes sentences with GCN for semantic role labeling. In our model, GCN architecture is applied to encode the positional and formatting relationship between the text nodes.

Attention mechanisms have been added to graph neural networks for more expressiveness. GAT (Velickovic et al., 2018) enables specifying different weights to different nodes and outperforms its convolutional counterparts on knowledge graph dataset. Then Graph Transformers(Koncel-Kedziorski et al., 2019) present an end-to-end system for graph-to-text generation from knowledge graphs with multi-head attention inspired by transformers.

Since BERT (Devlin et al., 2018) achieved SOTA performances on a range of NLP tasks, there has been much more attention on pre-training with large amounts of unlabeled data to produce powerful contextual representations for downstream tasks. After pre-training from massive text data with language modeling objectives (Howard and Ruder, 2018; Peters et al., 2018; Radford et al., 2018), the encoder could learn high-level dependencies between tokens and then feed the representations to downstream tasks to perform transfer learning. In addition to language model related objectives, many works apply other pre-training objectives which has also made significant progress. BERT introduces a next sentence prediction (NSP) loss considering the sentence level relationships. (Mayhew et al., 2019) comes up with a pre-training objective predicts casing in the text to address the robustness problem for named entity recognition (NER) systems and (Mehri et al., 2019) introduces four different pre-training objectives for dialog context representation learning.

In our work, we fine tune the proposed RoBERTa-GCN model with two objectives, the Masked Language Model (MLM) used by BERT and RoBERTa, as well as the novel sequence positional relationship classification loss (SPRC), which is designed to capture the relationships between text chunks given their position and content.

Training statistical models typically requires annotating a large amount of data, which can often be costly. As a remedy, the few shot learning task (Miller et al., 2000; Fei-Fei et al., 2006) is proposed to build models that can quickly learn from a small number of samples. To address this problem, our method should be robust enough to adapt to an unseen dataset based on very few samples. Many few-shot learning methods consider image domains e.g. (Vinyals et al., 2016). For NLP, (Yu et al., 2018) proposes an adaptive metric learning approach to automatically assign weights for newly seen few-shot task and (Geng et al., 2019) uses a few-shot learning method that leverages the dynamic routing algorithm in meta-learning for intention recognition in conversation.

Our approach of binding the GCN with pre-trained BERT-like model and fine-tuning the model with in-domain unlabeled data facilitates few-shot learning by optimizing both language features (through pre-trained transformer-based language models) and visual features (through layout-aware fine-tuning on unlabeled in-domain data) on unlabeled data as much as possible. Experiments show that our method outperforms strong text-based baseline by significant margins in zero- and few-shot settings.

We focus on digital-born documents, which are documents produced and received in digital form without printing and OCR. These documents preserve rich formatting information and constitute a large portion of business documents to be analyzed today. We model each document as a set of text boxes and apply sequence tagging on each text box. Each text box corresponds to a node in the graph. We create a graph for each page, if a document has multiple pages.

To convert a document page to a graph, we define the graph as G, where $G=(V,E)$, $V=\{v_{1},v_{2},\cdots,v_{N}\}$ is the set of text boxes in one page, and matrix $E\subset M\times V\times V$ is the undirected edges set, where M is the number of edge types. In our experiments, we define a basic undirected edge type $e_{i,j}=(v_{i},v_{j})$ that connects a text box to its closest vertical or horizontal neighbor. The model can also accommodate multiple edge types defined in different domains (cf. §3.3).

Our model combines text features and layout features to perform entity extraction.

Different from some existing works (Liu et al., 2019a; Qian et al., 2019) that encode words or sentences using the standard BiLSTM-CRF(Schuster and Paliwal, 1997), we take advantage of pre-trained transformer-based language models influenced by BERT (Devlin et al., 2018; Liu et al., 2019b) to perform text feature extraction as they are much more expressive.

The BERT model is based on the multi-layer bidirectional transformer (Vaswani et al., 2017) architecture that effectively encodes a sequence of tokens to produce final representations of the text. It is often trained on large unlabeled text data using a language model objective and then fine-tuned on in-domain data to effectively transfer knowledge from the large-scale unlabeled text corpus for a specific task.

The BERT model takes as input a set of tokens. For each token, its vector representation is computed by summing the corresponding token embedding, positional embedding, and segment embedding. A special classification embedding [CLS] is always added as the first token of the sequence, and a special end-of-sequence token [SEP] is added to the end. Then, these inputs are passed through a multi-layer bidirectional transformer encoder.

During pre-training, BERT uses two objectives: Masked Language Model (MLM) and Next Sentence Prediction (NSP). The MLM objective randomly masks tokens in a sequence and the language model is trained to recover the masked tokens. The NSP objective performs binary classification on whether two segments follow each other in the original text. As it is pre-trained on a large corpus with a large number of parameters, BERT-like language model has shown excellent transferability (Peng et al., 2019; Adhikari et al., 2019; Nogueira and Cho, 2019): downstream tasks can often achieve high accuracy even if the data size of the specific task is not large enough, when the models are initialized with pre-trained BERT weights.

In our work, we use the pre-trained RoBERTa model (Liu et al., 2019b), which is similar to BERT in architecture but tuned with different objectives on larger data, as the encoder backbone to obtain expressive representations of text. Given a tokenized text box $s_{i}=(\omega_{1}^{(i)},\omega_{2}^{(i)},\cdots,\omega_{k}^{(i)})$ of length $k$ as input, we pool the RoBERTa model output by concatenating the final hidden state corresponding to the first token [CLS] as the aggregate representation of the sentence. The sentence representation vector [CLS] is used as messages in the graph neural network to model page layout on the graph level. The other token embeddings $H_{1:k}^{(i)}$ will be concatenated with the outputs of the graph layers (cf. §3.3) as features for the entity tagger.

Empirically, visual information of non-standard format VRDs varies greatly in different scenarios. For example, information in invoices is often shown as lists, while information in resume is often organized in sections. Accordingly, we adapt the graph structure to model the sophisticated layouts and fonts to learn task specific visual features, rather than relying solely on the bounding box positions into the pre-trained model like (Xu et al., 2019).

With the layout encoding from GCN and the text encoding from the pre-trained language model, we sequentially tag each word in a sentence to extract entities. As illustrated in Figure3, the final layer output representation of the i-th text box $h_{i}^{L}$ from the GCN is concatenated to each token states $H_{1:k}^{(i)}$ from the RoBERTa encoder. Then, a sequence labeling layer is used to predict entity types for each token. Entity types are encoded at token level with the BIO schema.

Labeling data is expensive for information extraction applications, but unlabeled data is more likely abundant. So we try to make full use of the large amount of unlabeled data in our model through unsupervised in-domain fine-tuning.

Specifically, assuming that we have data $D$ for a given task. In addition to training all labeled document data $D_{labeled}$, we also make use of plenty of unlabeled data $D_{unlabeled}$ during the fine-tuning process to adapt the pre-trained model to a new task.

An extraction task can cover different domains. Invoice extraction, for example, may have to process both hotel invoice and appliance invoice for different users. Document from different domains often use different language and layout and it is traditionally expensive to adapt models trained on one domain to other domains. To better evaluate the performance of our system on different domains, we further divide $D_{labeled}$ into $D_{seen}$ and $D_{unseen}$, where $D_{unseen}$ is a set of testing documents from unseen domains. The model is first fine-tuned with $D_{unlabeled}$, and then trained to extract entities based on a portion of $D_{seen}$. It is finally tested on the test portion of $D_{seen}$ and $D_{unseen}$.

Inspired by existing pre-trained models, we propose to fine tune on unlabeled in-domain data with two training objectives before supervised training. Experiments show that in-domain fine-tuning helps extraction performance on both $D_{seen}$ and $D_{unseen}$.

Although we have collected a large number of invoices from different sources to build the dataset, the number of document layout types is still limited and it is impractical to prepare labeled invoice document from all domains, so it is necessary to evaluate our method in zero- and few-shot settings to verify the robustness of our model on new domains. Zero-shot learning is well-known as a problem where no in-domain labeled training data is available but the system has to make predictions. Few-shot learning is to test the ability of the neural networks on unseen dataset when provided only a very small number of training instances.

In zero-shot experiments, we select invoices from a specific department from labeled dataset as $D_{unseen}$ as test set and use $D_{seen}$ for training and validation. For few-shot learning, we select 70 invoices from the same department with $D_{unseen}$ from $D_{unlabeled}$ and manually labeled this set as $D_{few}$ to fine-tune our model. It is worth mentioning that $D_{few}\cap D_{unseen}=\emptyset$, $D_{few}\cap D_{labeled}=\emptyset$, which means that the $D_{few}$ is constructed to be used for few-shot experiments only.

Table 5 illustrates the results of these experiments. In zero-shot experiments, the performances of both methods are not as good as the supervised setting, where abundant training data is available. Nevertheless, our model still obtains a 3.9% improvement over the baseline on $D_{unseen}$.

We then focus on the few-shot setting where there are few labeled data fine-tuning the trained model, as a few shot experiment. In our experiment, $D_{few}$ is further splitted into 50 invoices as validation and 20 for fine-tuning. After 5 epochs fine-tuning process, out model outperforms the baseline by a significant margin at 94.62, which is comparable to the performance of fully supervised training. The few-shot results further prove the effectiveness and robustness of our method.

To understand how the system performs as the number of training instances increase, we manually label 500 invoices from the $D_{unseen}$ subset and select different number of training instances (1, 10, 50, 300, and 500) to fine-tune the model. The results on $D_{unseen}$ are illustrated in Figure 5. It can be seen from the figure that our model achieves $\sim 90\%$ F1-score with only about ten labeled instances. In contrast, to achieve the same level of performance as our model, the baseline model requires about 300 instances, which translates into 30x annotation cost reduction for the proposed model. We conjure that the improvement is facilitated by our models’ ability to effectively utilize both text and layout information in both supervised and unsupervised settings.

We test on another real-world dataset containing hundreds of international resumes which are all in different layouts. This dataset is different from the invoice dataset in both text content and layout and can be considered as a different task. Extensive experiments are conducted on this dataset to evaluate if the proposed model can generalize across different types of documents.