LayoutLMv2: Multi-modal Pre-training for Visually-Rich Document Understanding

Introduction

Visually-rich Document Understanding (VrDU) aims to analyze scanned/digital-born business documents (images of invoices, forms in PDF format, etc.) where structured information can be automatically extracted and organized for many business applications. Distinct from conventional information extraction tasks, the VrDU task relies on not only textual information but also visual and layout information that is vital for visually-rich documents. Different types of documents indicate that the text fields of interest located at different positions within the document, which is often determined by the style and format of each type as well as the document content. Therefore, to accurately recognize the text fields of interest, it is inevitable to take advantage of the cross-modality nature of visually-rich documents, where the textual, visual, and layout information should be jointly modeled and learned end-to-end in a single framework.

The recent progress of VrDU lies primarily in two directions. The first direction is usually built on the shallow fusion between textual and visual/layout/style information (Yang et al., 2017; Liu et al., 2019; Sarkhel and Nandi, 2019; Yu et al., 2020; Majumder et al., 2020; Wei et al., 2020; Zhang et al., 2020). These approaches leverage the pre-trained NLP and CV models individually and combine the information from multiple modalities for supervised learning. Although good performance has been achieved, the domain knowledge of one document type cannot be easily transferred into another, so that these models often need to be re-trained once the document type is changed. Thereby the local invariance in general document layout (key-value pairs in a left-right layout, tables in a grid layout, etc.) cannot be fully exploited. To this end, the second direction relies on the deep fusion among textual, visual, and layout information from a great number of unlabeled documents in different domains, where pre-training techniques play an important role in learning the cross-modality interaction in an end-to-end fashion (Lockard et al., 2020; Xu et al., 2020). In this way, the pre-trained models absorb cross-modal knowledge from different document types, where the local invariance among these layouts and styles is preserved. Furthermore, when the model needs to be transferred into another domain with different document formats, only a few labeled samples would be sufficient to fine-tune the generic model in order to achieve state-of-the-art accuracy. Therefore, the proposed model in this paper follows the second direction, and we explore how to further improve the pre-training strategies for the VrDU tasks.

In this paper, we present an improved version of LayoutLM (Xu et al., 2020), aka LayoutLMv2. Different from the vanilla LayoutLM model where visual embeddings are combined in the fine-tuning stage, we integrate the visual information in the pre-training stage in LayoutLMv2 by taking advantage of the Transformer architecture to learn the cross-modality interaction between visual and textual information. In addition, inspired by the 1-D relative position representations (Shaw et al., 2018; Raffel et al., 2020; Bao et al., 2020), we propose the spatial-aware self-attention mechanism for LayoutLMv2, which involves a 2-D relative position representation for token pairs. Different from the absolute 2-D position embeddings that LayoutLM uses to model the page layout, the relative position embeddings explicitly provide a broader view for the contextual spatial modeling. For the pre-training strategies, we use two new training objectives for LayoutLMv2 in addition to the masked visual-language modeling. The first is the proposed text-image alignment strategy, which aligns the text lines and the corresponding image regions. The second is the text-image matching strategy popular in previous vision-language pre-training models (Tan and Bansal, 2019; Lu et al., 2019; Su et al., 2020; Chen et al., 2020; Sun et al., 2019), where the model learns whether the document image and textual content are correlated.

We select six publicly available benchmark datasets as the downstream tasks to evaluate the performance of the pre-trained LayoutLMv2 model, which are the FUNSD dataset (Jaume et al., 2019) for form understanding, the CORD dataset (Park et al., 2019) and the SROIE dataset (Huang et al., 2019) for receipt understanding, the Kleister-NDA dataset (Graliński et al., 2020) for long document understanding with a complex layout, the RVL-CDIP dataset (Harley et al., 2015) for document image classification, and the DocVQA dataset (Mathew et al., 2021) for visual question answering on document images. Experiment results show that the LayoutLMv2 model significantly outperforms strong baselines, including the vanilla LayoutLM, and achieves new state-of-the-art results in all of these tasks.

The contributions of this paper are summarized as follows:

We propose a multi-modal Transformer model to integrate the document text, layout, and visual information in the pre-training stage, which learns the cross-modal interaction end-to-end in a single framework. Meanwhile, a spatial-aware self-attention mechanism is integrated into the Transformer architecture.

In addition to the masked visual-language model, we add text-image alignment and text-image matching as the new pre-training strategies to enforce the alignment among different modalities.

LayoutLMv2 significantly outperforms and achieves new SOTA results not only on the conventional VrDU tasks but also on the VQA task for document images, which demonstrates the great potential for the multi-modal pre-training for VrDU.

Approach

In this section, we will introduce the model architecture and the multi-modal pre-training tasks of LayoutLMv2, which is illustrated in Figure 1.

We build a multi-modal Transformer architecture as the backbone of LayoutLMv2, which takes text, visual, and layout information as input to establish deep cross-modal interactions. We also introduce a spatial-aware self-attention mechanism to the model architecture for better modeling the document layout. Detailed descriptions of the model are as follows.

Following the common practice, we use WordPiece (Wu et al., 2016) to tokenize the OCR text sequence and assign each token to a certain segment $s_{i}\in\{\texttt{[A]},\texttt{[B]}\}$. Then, we add [CLS] at the beginning of the sequence and [SEP] at the end of each text segment. Extra [PAD] tokens are appended to the end so that the final sequence’s length is exactly the maximum sequence length $L$. The final text embedding is the sum of three embeddings. Token embedding represents the token itself, 1D positional embedding represents the token index, and segment embedding is used to distinguish different text segments. Formally, we have the $i$-th ($0\leq i<L$) text embedding

Visual Embedding

Although all information we need is contained in the page image, the model has difficulty capturing detailed features in a single information-rich representation of the entire page. Therefore, we leverage the output feature map of a CNN-based visual encoder, which converts the page image to a fixed-length sequence. We use ResNeXt-FPN (Xie et al., 2017; Lin et al., 2017) architecture as the backbone of the visual encoder, whose parameters can be updated through backpropagation.

Layout Embedding

Multi-modal Encoder with Spatial-Aware Self-Attention Mechanism

The encoder concatenates visual embeddings $\{\mathbf{v}_{0},...,\mathbf{v}_{WH-1}\}$ and text embeddings $\{\mathbf{t}_{0},...,\mathbf{t}_{L-1}\}$ to a unified sequence and fuses spatial information by adding the layout embeddings to get the $i$-th ($0\leq i<WH+L$) first layer input

Following the architecture of Transformer, we build our multi-modal encoder with a stack of multi-head self-attention layers followed by a feed-forward network. However, the original self-attention mechanism can only implicitly capture the relationship between the input tokens with the absolute position hints. In order to efficiently model local invariance in the document layout, it is necessary to insert relative position information explicitly. Therefore, we introduce the spatial-aware self-attention mechanism into the self-attention layers. For simplicity, the following description is for a single head in a single self-attention layer with hidden size of $d_{head}$ and projection matrics $\mathbf{W}^{Q}$, $\mathbf{W}^{K}$, $\mathbf{W}^{V}$. The original self-attention mechanism captures the correlation between query $\mathbf{x}_{i}$ and key $\mathbf{x}_{j}$ by projecting the two vectors and calculating the attention score

Finally, the output vectors are represented as the weighted average of all the projected value vectors with respect to normalized spatial-aware attention scores

2 Pre-training Tasks

Similar to the vanilla LayoutLM, we use the Masked Visual-Language Modeling (MVLM) to make the model learn better in the language side with the cross-modality clues. We randomly mask some text tokens and ask the model to recover the masked tokens. Meanwhile, the layout information remains unchanged, which means the model knows each masked token’s location on the page. The output representations of masked tokens from the encoder are fed into a classifier over the whole vocabulary, driven by a cross-entropy loss. To avoid visual clue leakage, we mask image regions corresponding to masked tokens on the raw page image input before feeding it into the visual encoder.

Text-Image Alignment

To help the model learn the spatial location correspondence between image and coordinates of bounding boxes, we propose the Text-Image Alignment (TIA) as a fine-grained cross-modality alignment task. In the TIA task, some tokens lines are randomly selected, and their image regions are covered on the document image. We call this operation covering to avoid confusion with the masking operation in MVLM. During pre-training, a classification layer is built above the encoder outputs. This layer predicts a label for each text token depending on whether it is covered, i.e., [Covered] or [Not Covered], and computes the binary cross-entropy loss. Considering the input image’s resolution is limited, and some document elements like signs and bars in a figure may look like covered text regions, the task of finding a word-sized covered image region can be noisy. Thus, the covering operation is performed at the line-level. When MVLM and TIA are performed simultaneously, TIA losses of the tokens masked in MVLM are not taken into account. This prevents the model from learning the useless but straightforward correspondence from [MASK] to [Covered].

Text-Image Matching

Furthermore, a coarse-grained cross-modality alignment task, Text-Image Matching (TIM) is applied to help the model learn the correspondence between document image and textual content. We feed the output representation at [CLS] into a classifier to predict whether the image and text are from the same document page. Regular inputs are positive samples. To construct a negative sample, an image is either replaced by a page image from another document or dropped. To prevent the model from cheating by finding task features, we perform the same masking and covering operations to images in negative samples. The TIA target labels are all set to [Covered] in negative samples. We apply the binary cross-entropy loss in the optimization process.

Experiments

In order to pre-train and evaluate LayoutLMv2 models, we select datasets in a wide range from the visually-rich document understanding area. Following LayoutLM, we use IIT-CDIP Test Collection (Lewis et al., 2006) as the pre-training dataset. Six datasets are used as down-stream tasks. The FUNSD (Jaume et al., 2019), CORD (Park et al., 2019), SROIE (Huang et al., 2019) and Kleister-NDA (Graliński et al., 2020) datasets define entity extraction tasks that aim to extract the value of a set of pre-defined keys, which we formalize as a sequential labeling task. RVL-CDIP (Harley et al., 2015) is for document image classification. DocVQA (Mathew et al., 2021), as the name suggests, is a dataset for visual question answering on document images. Statistics of datasets are shown in Table 1. Refer to the Appendix for details.

2 Settings

Following the typical pre-training and fine-tuning strategy, we update all parameters including the visual encoder layers, and train whole models end-to-end for all the settings. Training details can be found in the Appendix.

We train LayoutLMv2 models with two different parameter sizes. We use a 12-layer 12-head Transformer encoder and set hidden size $d=768$ in LayoutLMv2BASE. While in the LayoutLMv2LARGE, the encoder has 24 Transformer layers with 16 heads and $d=1024$. Visual backbones in the two models are based on the same ResNeXt101-FPN architecture. The numbers of parameters are 200M and 426M approximately for LayoutLMv2BASE and LayoutLMv2LARGE, respectively.

For the encoder along with the text embedding layer, LayoutLMv2 uses the same architecture as UniLMv2 (Bao et al., 2020), thus it is initialized from UniLMv2. For the ResNeXt-FPN part in the visual embedding layer, the backbone of a Mask-RCNN (He et al., 2017) model trained on PubLayNet (Zhong et al., 2019) is leveraged.“MaskRCNN ResNeXt101_32x8d FPN 3X” setting in https://github.com/hpanwar08/detectron2 The rest of the parameters in the model are randomly initialized.

During pre-training, we sample pages from the IIT-CDIP dataset and select a random sliding window of the text sequence if the sample is too long. We set the maximum sequence length $L=512$ and assign all text tokens to the segment [A]. The output shape of the average pooling layer is set to $W=H=7$, so that it transforms the feature map into $49$ visual tokens. In MVLM, 15% text tokens are masked among which 80% are replaced by a special token [MASK], 10% are replaced by a random token sampled from the whole vocabulary, and 10% remains the same. In TIA, 15% of the lines are covered. In TIM, 15% images are replaced, and 5% are dropped.

Fine-tuning LayoutLMv2

We use the [CLS] output along with pooled visual token representations as global features in the document-level classification task RVL-CDIP. For the extractive question answering task DocVQA and the other four entity extraction tasks, we follow common practice like Devlin et al. (2019) and build task specified head layers over the text part of LayoutLMv2 outputs.

In the DocVQA paper, experiment results show that the BERT model fine-tuned on the SQuAD dataset (Rajpurkar et al., 2016) outperforms the original BERT model. Inspired by this fact, we add an extra setting, which is that we first fine-tune LayoutLMv2 on a question generation (QG) dataset followed by the DocVQA dataset. The QG dataset contains almost one million question-answer pairs generated by a generation model trained on the SQuAD dataset.

Baselines

3 Results

RVL-CDIP

Table 3 shows the classification accuracy on the RVL-CDIP dataset, including text-only pre-trained models, the LayoutLM family as well as several image-based baseline models. As shown in the table, both the text and visual information are important to the document image classification task because document images are text-intensive and represented by a variety of layouts and formats. Therefore, we observed that the LayoutLM family outperforms those text-only or image-only models as it leverages the multi-modal information within the documents. Specifically, the $\textrm{LayoutLMv2}_{\rm LARGE}$ model significantly improves the classification accuracy by more than 1.2% point over the previous SOTA results, which achieves an accuracy of 95.64%. This also verifies that the pre-trained LayoutLMv2 model benefits not only the information extraction tasks in document understanding but also the document image classification task through effective multi-model training.

DocVQA

Table 4 lists the Average Normalized Levenshtein Similarity (ANLS) scores on the DocVQA dataset of text-only baselines, LayoutLM family models, and the previous top-1 on the leaderboard. With multi-modal pre-training, LayoutLMv2 models outperform LayoutLM models and text-only baselines by a large margin when fine-tuned on the train set. By using all data (train + dev) as the fine-tuning dataset, the $\textrm{LayoutLMv2}_{\rm LARGE}$ single model outperforms the previous top-1 on the leaderboard which ensembles 30 models.Unpublished results, the leaderboard is available at https://rrc.cvc.uab.es/?ch=17&com=evaluation&task=1 Under the setting of fine-tuning $\textrm{LayoutLMv2}_{\rm LARGE}$ on a question generation dataset (QG) and the DocVQA dataset successively, the single model performance increases by more than 1.6% ANLS and achieves the new SOTA.

4 Ablation Studies

To fully understand the underlying impact of different components, we conduct an ablation study to explore the effect of visual information, the pre-training tasks, spatial-aware self-attention mechanism, as well as different text-side initialization models. Table 5 shows model performance on the DocVQA validation set. Under all the settings, we pre-train the models using all IIT-CDIP data for one epoch. The hyper-parameters are the same as those used to pre-train $\textrm{LayoutLMv2}_{\rm BASE}$ in Section 3.2. “LayoutLM” denotes the vanilla LayoutLM architecture in (Xu et al., 2020), which can be regarded as a LayoutLMv2 architecture without visual module and spatial-aware self-attention mechanism. “X101-FPN” denotes the ResNeXt101-FPN visual backbone described in Section 3.2.

We first evaluate the effect of introducing visual information. From #1 to #2a, we add the visual module without changing the pre-training strategy, where results show that LayoutLMv2 pre-trained with only MVLM can leverage visual information effectively. Then, we compare the two cross-modality alignment pre-training tasks TIA and TIM. According to the four results in #2, both tasks improve the model performance substantially, and the proposed TIA benefits the model more than the commonly used TIM. Using both tasks together is more effective than using either one alone. According to this observation, we keep all the three pre-training tasks and introduce the spatial-aware self-attention mechanism (SASAM) to the model architecture. Compare the results #2d and #3, the proposed SASAM can further improve the model accuracy. Finally, in settings #3 and #4, we change the text-side initialization checkpoint from BERT to UniLMv2, and confirm that LayoutLMv2 benefits from the better initialization.

Related Work

In recent years, pre-training techniques have become popular in both NLP and CV areas, and have also been leveraged in the VrDU tasks.

Devlin et al. (2019) introduced a new language representation model called BERT, which is designed to pre-train deep bidirectional representations from the unlabeled text by jointly conditioning on both left and right context in all layers. Bao et al. (2020) propose to pre-train a unified language model for both autoencoding and partially autoregressive language modeling tasks using a novel training procedure, referred to as a pseudo-masked language model. Our multi-modal Transformer architecture and the MVLM pre-training strategy extend Transformer and MLM used in these work to leverage visual information.

Lu et al. (2019) proposed ViLBERT for learning task-agnostic joint representations of image content and natural language by extending the popular BERT architecture to a multi-modal two-stream model. Su et al. (2020) proposed VL-BERT that adopts the Transformer model as the backbone, and extends it to take both visual and linguistic embedded features as input. Different from these vision-language pre-training approaches, the visual part of LayoutLMv2 directly uses the feature map instead of pooled ROI features, and benefits from the new TIA pre-training task.

Xu et al. (2020) proposed LayoutLM to jointly model interactions between text and layout information across scanned document images, benefiting a great number of real-world document image understanding tasks such as information extraction from scanned documents. This work is a natural extension of the vanilla LayoutLM, which takes advantage of textual, layout, and visual information in a single multi-modal pre-training framework.

Conclusion

In this paper, we present a multi-modal pre-training approach for visually-rich document understanding tasks, aka LayoutLMv2. Distinct from existing methods for VrDU, the LayoutLMv2 model not only considers the text and layout information but also integrates the image information in the pre-training stage with a single multi-modal framework. Meanwhile, the spatial-aware self-attention mechanism is integrated into the Transformer architecture to capture the relative relationship among different bounding boxes. Furthermore, new pre-training objectives are also leveraged to enforce the learning of cross-modal interaction among different modalities. Experiment results on 6 different VrDU tasks have illustrated that the pre-trained LayoutLMv2 model has substantially outperformed the SOTA baselines in the document intelligence area, which greatly benefits a number of real-world document understanding tasks.

For future research, we will further explore the network architecture as well as the pre-training strategies for the LayoutLM family. Meanwhile, we will also investigate the language expansion to make the multi-lingual LayoutLMv2 model available for different languages, especially the non-English areas around the world.

Acknowledgments

This work was supported by the National Key R&D Program of China via grant 2020AAA0106501 and the National Natural Science Foundation of China (NSFC) via grant 61976072 and 61772153.

References

Appendix

Appendix A Details of Datasets

Introduction to the dataset and task definitions along with the description of required data processing are presented as follows.

Following LayoutLM, we pre-train LayoutLMv2 on the IIT-CDIP Test Collection (Lewis et al., 2006), which contains over 11 million scanned document pages. We extract text and corresponding word-level bounding boxes from document page images with the Microsoft Read API.https://docs.microsoft.com/en-us/azure/cognitive-services/computer-vision/concept-recognizing-text

FUNSD

FUNSD (Jaume et al., 2019) is a dataset for form understanding in noisy scanned documents. It contains 199 real, fully annotated, scanned forms where 9,707 semantic entities are annotated above 31,485 words. The 199 samples are split into 149 for training and 50 for testing. The official OCR annotation is directly used with the layout information. The FUNSD dataset is suitable for a variety of tasks, where we focus on semantic entity labeling in this paper. Specifically, the task is assigning to each word a semantic entity label from a set of four predefined categories: question, answer, header, or other. The entity-level F1 score is used as the evaluation metric.

CORD

We also evaluate our model on the receipt key information extraction dataset, i.e. the public available subset of CORD (Park et al., 2019). The dataset includes 800 receipts for the training set, 100 for the validation set, and 100 for the test set. A photo and a list of OCR annotations are equipped for each receipt. An ROI that encompasses the area of receipt region is provided along with each photo because there can be irrelevant things in the background. We only use the ROI as input instead of the raw photo. The dataset defines 30 fields under 4 categories and the task aims to label each word to the right field. The evaluation metric is entity-level F1. We use the official OCR annotations.

SROIE

The SROIE dataset (Task 3) (Huang et al., 2019) aims to extract information from scanned receipts. There are 626 samples for training and 347 samples for testing in the dataset. The task is to extract values from each receipt of up to four predefined keys: company, date, address, or total. The evaluation metric is entity-level F1. We use the official OCR annotations and results on the test set are provided by the official evaluation site.

Kleister-NDA

Kleister-NDA (Graliński et al., 2020) contains non-disclosure agreements collected from the EDGAR database, including 254 documents for training, 83 documents for validation, and 203 documents for testing. This task is defined to extract the values of four fixed keys. We get the entity-level F1 score from the official evaluation tools.https://gitlab.com/filipg/geval Words and bounding boxes are extracted from the raw PDF file. We use heuristics to locate entity spans because the normalized standard answers may not appear in the utterance. As the labeled answers are normalized into a canonical form, we apply post-processing heuristics to convert the extracted date information into the “YYYY-MM-DD” format, and company names into the abbreviations such as “LLC” and “Inc.”.

RVL-CDIP

RVL-CDIP (Harley et al., 2015) consists of 400,000 grayscale images, with 8:1:1 for the training set, validation set, and test set. A multi-class single-label classification task is defined on RVL-CDIP. The images are categorized into 16 classes, with 25,000 images per class. The evaluation metric is the overall classification accuracy. Text and layout information is extracted by Microsoft OCR.

DocVQA

As a VQA dataset on the document understanding field, DocVQA (Mathew et al., 2021) consists of 50,000 questions defined on over 12,000 pages from a variety of documents. Pages are split into the training set, validation set, and test set with a ratio of about 8:1:1. The dataset is organized as a set of triples $\langle$page image, questions, answers$\rangle$. Thus, we use Microsoft Read API to extract text and bounding boxes from images. Heuristics are used to find given answers in the extracted text. The task is evaluated using an edit distance based metric ANLS (aka average normalized Levenshtein similarity). Given that human performance is about 98% ANLS on the test set, it is reasonable to assume that the found ground truth which reaches over 97% ANLS on training and validation sets is good enough to train a model. Results on the test set are provided by the official evaluation site.

Appendix B Model Training Details

We pre-train LayoutLMv2 models using Adam optimizer (Kingma and Ba, 2015; Loshchilov and Hutter, 2019), with the learning rate of $2\times 10^{-5}$, weight decay of $1\times 10^{-2}$,and $(\beta_{1},\beta_{2})=(0.9,0.999)$. The learning rate is linearly warmed up over the first $10\%$ steps then linearly decayed. $\textrm{LayoutLMv2}_{\rm BASE}$ is trained with a batch size of $64$ for $5$ epochs, and $\textrm{LayoutLMv2}_{\rm LARGE}$ is trained with a batch size of $2048$ for $20$ epochs on the IIT-CDIP dataset.

Fine-tuning for Visual Question Answering

We treat the DocVQA as an extractive QA task and build a token-level classifier on top of the text part of LayoutLMv2 output representations. Question tokens, context tokens and visual tokens are assigned to segment [A], [B] and [C], respectively. The maximum sequence length is set to $L=384$.

Fine-tuning for Document Image Classification

This task depends on high-level visual information, thereby we leverage the image features explicitly in the fine-tuning stage. We pool the visual embeddings into a global pre-encoder feature, and pool the visual part of LayoutLMv2 output representations into a global post-encoder feature. The pre and post-encoder features along with the [CLS] output feature are concatenated and fed into the final classification layer.

Fine-tuning for Sequential Labeling

We formalize FUNSD, SROIE, CORD, and Kleister-NDA as the sequential labeling tasks. To fine-tune LayoutLMv2 models on these tasks, we build a token-level classification layer above the text part of the output representations to predict the BIO tags for each entity field.

Appendix C Detailed Experiment Results

Tables list per-task detailed results for the four entity extraction tasks, with Table 6 for FUNSD, Table 7 for CORD, Table 8 for SROIE, and Table 9 for Kleister-NDA.