LMDX: Language Model-based Document Information Extraction and Localization (2024)

Abstract: Large Language Models (LLM) have revolutionized Natural Language Processing (NLP), improving state-of-the-art and exhibiting emergent capabilities across various tasks. However, their application in extracting information from visually rich documents, which is at the core of many document processing workflows and involving the extraction of key entities from semi-structured documents, has not yet been successful. The main obstacles to adopting LLMs for this task include the absence of layout encoding within LLMs, which is critical for high quality extraction, and the lack of a grounding mechanism to localize the predicted entities within the document. In this paper, we introduce Language Model-based Document Information EXtraction and Localization (LMDX), a methodology to reframe the document information extraction task for a LLM. LMDX enables extraction of singular, repeated, and hierarchical entities, both with and without training data, while providing grounding guarantees and localizing the entities within the document. Finally, we apply LMDX to the PaLM 2-S and Gemini Pro LLMs and evaluate it on VRDU and CORD benchmarks, setting a new state-of-the-art and showing how LMDX enables the creation of high quality, data-efficient parsers.

text: Findings of the Association for Computational Linguistics ACL 2024, pages 15140–15168 August 11-16, 2024 ©2024 Association for Computational Linguistics
Large Language Models (LLM) have revolutionized Natural Language Processing (NLP), improving state-of-the-art and exhibiting emergent capabilities across various tasks. However, their application in extracting information from visually rich documents, which is at the core of many document processing workflows and involving the extraction of key entities from semi-structured documents, has not yet been successful. The main obstacles to adopting LLMs for this task include the absence of layout encoding within LLMs, which is critical for high quality extraction, and the lack of a grounding mechanism to localize the predicted entities within the document. In this paper, we introduce Language Model-based Document Information EXtraction and Localization (LMDX), a methodology to reframe the document information extraction task for a LLM. LMDX enables extraction of singular, repeated, and hierarchical entities, both with and without training data, while providing grounding guarantees and localizing the entities within the document. Finally, we apply LMDX to the PaLM 2-S and Gemini Pro LLMs and evaluate it on VRDU and CORD benchmarks, setting a new state-of-the-art and showing how LMDX enables the creation of high quality, data-efficient parsers.

heading: 1 Introduction
text: The recent advent of transformers (Vaswani et al., 2017) and self-supervised pretraining procedures has led to significant progress in Visually Rich Document (VRD) Understanding. Within that field, the task of document information extraction (IE), which consists of extracting key entities within a semi-structured document (e.g. invoice, tax form, paystub, receipt, etc) given a predefined schema, has received a lot of attention from industry and
* Correspondence to <vperot@google.com>.
academia due to its importance and wide applicability to intelligent document processing workflows. However, document information extraction still remains challenging for mainstream systems.
In particular, information in semi-structured forms is organized in complex layout across many possible templates, which requires deep understanding of the document context, spatial alignment among the different segments of text, and tabular arrangement of hierarchical entities (we define hierarchical entities as entities that are composed of logically grouped leaf entities, e.g. line items on an invoice composed of item description, quantity and price, or deduction items on a paystub composed of deduction type and amount, etc.). Moreover, since some business document automation workflows require certain level of accuracy, they are often integrated with human-in-the-loop interactions for auditing and correction of predictions, requiring knowing the precise location of extracted entities to make it a tractable task for a human rater. Finally, since a quasi-infinite number of document types exist, and that organizations have limited annotation resources, most parsers are built with very small amount of training data. From those complexities emerge the following desiderata of document information extraction systems: they should (1) support high-quality extraction of singular, repeated, and hierarchical entities, while (2) localizing those entities in the document, and (3) do so with very low or no data annotation cost.
So far, no publicly disclosed system has been able to address all of those desiderata. Current mainstream document IE systems are based on sequence-tagging and sequence-generation. Sequence-tagging approaches (Palm et al., 2017; Lee et al., 2021, 2022, 2023a; Wang et al., 2023c) classifies each token into Inside-Outside-Begin (IOB) tags (Ramshaw and Marcus, 1995), which support extraction and localization of leaf entities. However, it is non-trivial to extend these meth-
15140
ods for hierarchical entities. Sequence-generation based methods (Powalski et al., 2021; Kim et al., 2022) treats extraction as text generation with autoregressive decoders (Sutskever et al., 2014). Although this line of work allows to predict hierarchical entities, it does not allow localizing entities in the document. Moreover, both categories require significant human annotation cost to ensure a high quality extraction. Thus, a unified IE framework addressing all three desiderata is highly valuable.
In parallel, Large Language Models (LLMs) (OpenAI, 2023a; Google et al., 2023; Hoffmann et al., 2022) have revolutionized Natural Language Processing, showing the capabilities to solve diverse tasks with an instruction (Wei et al., 2022) or a few demonstrations attached to the prompt (Brown et al., 2020). This paradigm shift opens the possibility of extracting entities while addressing all the aforementioned desiderata, but using LLMs for VRD IE has been underexplored. Wang et al. (2023a) proposes a document-centric LLM and frames IE as a question-answer task, enabling zero-shot extraction but lacking support of localization and hierarchical entity extraction. Furthermore, this can suffer from hallucinations, a common issue with LLMs (Huang et al., 2023).
This motivates us to introduce Language Modelbased Document Information EXtraction and Localization (LMDX), a principled methodology for leveraging existing LLMs for information extraction and localization on visually-rich documents, meeting all three identified desiderata of extraction systems and detailed in Figure 1. A comparison of LMDX characteristics and other document information extraction systems can be found at Table 1. Our contributions can be summarized as follows:
• We present a principled recipe that enables LLMs to perform the document IE task on leaf and hierarchical entities with precise entity localization, including without any training data, and using only the simple text-in, textout interface that is applicable to LLMs.
• We propose a layout encoding scheme that communicates spatial information to the LLM without any changes to its architecture.
• We introduce a decoding algorithm that transforms the responses from the LLM into extracted entities and their corresponding bounding boxes on the document, while discarding any LLM hallucination.
• We systematically evaluate the data efficiency of LMDX across multiple public benchmarks, establishing a new state-of-the-art, and provide extensive study of the different core designs to demonstrate their effectiveness.

heading: 2 Related Work
text: Framing IE from VRD. Information Extraction from VRD is a complex task that can be framed in a variety of ways. Many approaches divide the problem in two sub-tasks: a text recognition/serialization step, typically achieved by an Optical Character Recognition (OCR) service, followed by a parsing step, which finds the relevant entity values from the recognized text. Xu et al. (2021); Appalaraju et al. (2021) frame this parsing step as Named Entity Recognition (NER), encoding each token with a transformer encoder and classifying each document token into IOB tags, allowing extraction and localization of leaf entities only. Other approaches treat extraction as a sequence generation problem. Powalski et al. (2021) adds an auto-regressive decoder on top of a text-layoutimage encoder, all initialized from T5 (Raffel et al., 2020). This enables to predict hierarchical entities, but not localize entities in the document. While LMDX still frames VRD IE as a sequence generation task, our work contrasts with prior work by combining the advantages of the different framings as shown in Table 1, supporting hierarchical entities, zero-shot extraction and localizing the entities through the introduction of coordinate tokens. VRD Representation Learning. As VRDs contain both textual and visual elements whose spatial position is crucial for their understanding, many works explore custom architectures and pretraining strategies to learn the relation between textual, layout and image modalities (Lee et al., 2023a; Appalaraju et al., 2023; Zhang et al., 2022). Xu
et al. (2020) uses a separate image encoder before adding the output as feature to the token encodings, while Huang et al. (2022) jointly models the page image patches alongside the tokens, using a wordpatch alignment self-supervised pretraining task to learn the connection between the modalities. Hong et al. (2021) proposes to encode the relative 2D distances of text blocks in the attention of the transformer, and learning from unlabeled documents with an area-masking strategy. Kim et al. (2022); Lee et al. (2023b) foregoes the text modality completely, using a Vision Transformer encoder with an auto-regressive decoder pretrained on a pseudoOCR and region masking task on large document image corpora. Unlike prior work, LMDX encodes the layout modality solely through text coordinate tokens, hence allows reusing LLMs with no architecture change and foregoing expensive vision encoders, while achieving state-of-the-art results. LLMs for Extraction has mostly been studied in the text domain (Keraghel et al., 2024), either generally (Laskar et al., 2023) or domain-specific (De Toni et al., 2022; Hu et al., 2024). Wang et al. (2023b) uses a LLM to insert special tokens to mark the boundaries of target entities. Ashok and Lipton (2023) proposes a NER framework with in-context learning demonstrations, prompting the LLM to output an entities list with explanations justifying its matches with the provided entity definitions. Yet, LLMs remain underexplored for IE on VRDs. Wang et al. (2023a) uses unlabeled document corpora and turns existing labeled VRD understanding datasets in instruction tuning format, building a layout-aware LLM with various document understanding capabilities. In contrast, LMDX focuses on IE specifically, with an emphasis on hierarchical entity and entity localization support.

heading: 3 LMDX Methodology
text: Overall, our pipeline is divided into four stages: chunking, prompt generation, LLM inference and decoding, detailed in the following sections. An overview with a simple example can be found in Figure 1, with the input and output of each stage showcased. In this example, the target extraction schema contains two leaf entity types retailer and subtotal, and one hierarchical entity type line_item, composed of a product_id and a product_price. Input Document. The input to our pipeline is the document’s text segments (lines and words) along with their corresponding spatial position (bounding box) on the pages, typically obtained with an OCR service or a PDF rendering engine.

heading: 3.1 First Stage: Chunking
text: While some LLMs support long context (hundreds of thousands of tokens), not all LLMs can fit the entire document within its prompt, as documents can be hundreds of pages long. Thus, the document is divided into document chunks so that each is small
enough to be processed by the LLM. To achieve this, we first divide the document into individual pages, then we iteratively remove the last line segments until the prompt containing this chunk is below the maximum input token length of the LLM. Lastly, we group those removed lines as a new document page, and repeat the same logic until all chunks are below the input token limit of the LLM. At the end of this stage, we have N chunks. The decision to first divide the document by page stems from the observation that entities rarely cross page boundaries, and as such this chunking scheme will have minimal impact on the final extraction quality. The algorithm is described in pseudo-code in Appendix A.1.

heading: 3.2 Second Stage: Prompt Generation
text: The prompt generation stage takes in the N document chunks and creates a LLM prompt for each of them. As seen in Figure 2, our prompt design contains the document representation, a description of the task, and the target schema representation containing the entities to extract. XML-like tags are used to define the start and end of each component.
Document Representation. The chunk content is represented in the prompt as the concatenation of all its segment texts (lines or words from OCR), suffixed with the coordinates of those segments (derived from the bounding boxes) in the following format: <segment text> XX|YYsegment. Coordinate tokens, XX and YY, are built by normalizing the segment’s X and Y coordinates, and quantizing them in B buckets, assigning the index of that bucket as the token for that coordinate.
Encoding the coordinates as tokens within the prompt allows us to communicate the layout modality to the LLM, without any change to its architecture. There are many variations to that scheme: using line versus words as segment, the granularity of the quantization, and the number of coordinates to use per segment (e.g. [xcenter, ycenter] versus [xmin, ymin, xmax, ymax]). Appendix A.4 shows
how those variations affect the prompt token length. Experimentally, we’ve found using line-level segments with 2 coordinates [xcenter, ycenter] and B = 100 quantization buckets worked best, as detailed in Appendix A.12. Hence, we’ve adopted that coordinate tokenization scheme in our experiments. Task Description. The task description is simply a short explanation of the task to accomplish. In our experiments, we set it to the following: From the document, extract the text values and tags of the following entities:. Schema Representation. The schema is represented as a structured JSON object, where the keys are the entity types to be extracted, and the values correspond to their occurrence (single or multiple) and sub-entities (for hierarchical entities). For instance, {"foo": "", "bar": [{"baz": []}]} means that the LLM should extract only a single entity of type foo and multiple hierarchical entities of type bar, that could each hold multiple entities of type baz.
After this step, we have N prompts, one for each document chunk. An example of a prompt on a document can be found in Appendix A.8, Figure 10.

heading: 3.3 Completion Targets
text: In this section, we describe the expected LLM completion format, which can be observed in LLM Completions section of Figure 1. Like the schema, the completion is a JSON structured object with the keys being the entity types, and values being the extracted information from the document chunk. JSON was chosen as a format for the completion and schema since it supports hierarchical objects (hence hierarchical entities), is very token-efficient, and emitting JSON is within mainstream LLMs’ capabilities (Sengottuvelu, 2023; OpenAI, 2023b). Note that the keys in the completion have the same ordering, occurrence and class (hierarchical or leaf) as the entity types in the schema. The values of leaf entities must follow a specific format:
<text on segment1> XX|YYsegment1\n <text on segment2> XX|YYsegment2\n ...
An entity can span multiple (potentially disjoint) text segments (lines or words). For each segment of the entity, the value contains the entity text on that segment, along with the coordinate tokens of that segment, which act as a segment identifier, uniquely identifying the segment, and allowing us to localize the entities and ground the model prediction (e.g. making sure the extracted value is not a hallucination), as will be detailed in Section 3.5. Fi-
nally, missing entity types are explicitly completed by the model with null for singular types, and [] for repeated types. Samples of completions can be found in Appendix A.8, Figure 10.

heading: 3.4 Third Stage: LLM Inference
text: In this stage of the pipeline, we run inference on the LLM with the N prompts. For each prompt, we sample K completions from the LLM (for a total of NK completions for the entire document) using TopK sampling. This randomness in the sampling allows to do error correction (e.g. if a response is not valid JSON, have hallucinated segment coordinate identifier, etc), and increase the extraction quality as will be shown in Section 4.3. We use a fixed random seed to get a deterministic inference.

heading: 3.5 Fourth Stage: Decoding
text: In this stage (Decoding in Figure 1), we parse the LLM completions into entities and their locations. Conversion to structured entities. We begin by parsing each LLM completion as a JSON object. Completions that fail to parse are discarded. For each key-value pair in the JSON object, we interpret the key as the entity type and parse the value to get the entity text and bounding box (as detailed in the next paragraph). Predicted entity types that are not in the target extraction schema are discarded. If the model unexpectedly predicts multiple values for single-occurrence entity types, we use the most frequent value as the final predicted value. Hierarchical JSON objects are recursively parsed as hierarchical entities in a similar manner. This algorithm is described in pseudo-code in Appendix A.3. Entity Value Parsing. We expect the JSON value to include both text extractions and segment identifiers for each predicted entity, as described in Section 3.3. We first parse the value into its (segment text, segment identifier) pairs. For each pair, we look up the corresponding segment in the original document using the segment identifier and verify that the extracted text is exactly included on that segment. The entity is discarded if that verification fails, ensuring LMDX discards all LLM hallucinations. Finally, once we have the entity location on all its segments, we get the entity bounding box by computing the smallest bounding box encompassing all the words included in the entity. Entity values with any segments that fail to ground (invalid entity value format, non-existent segment identifier, or non-matching segment text) in the original document are discarded. The entity value
parsing algorithm is described in pseudo-code in Appendix A.2, and parsing errors rates are detailed in Appendix A.10. Prediction Merging. We first merge the predicted entities for the same document chunk from the K LLM completions through majority voting (Wang et al., 2022). For each entity type, we gather the predicted entities, including empty predictions, across the K completions. The most common prediction(s) are selected as the predicted value for that entity type. We then merge the predictions among the N document chunks by concatenating them to obtain the document level predictions. Prediction Merging for hierarchical entities. For hierarchical entities, we use the entire predicted tree value from a single LLM completion, as this method best preserves the parent-child relationship predicted by the model. For each top-level hierarchical entity type, we perform majority voting on all affiliated leaf, intermediate and top-level entity types among K completions as if they are flattened. We then equally tally the votes to determine which completion to use for the prediction, and select the most common one for that hierarchical entity.

heading: 4 Evaluation
text: We seek to evaluate the effectiveness of LMDX on public IE benchmarks, and apply it to two distinct LLMs to validate the generality of the methodology: PaLM 2-S (Google et al., 2023) and Gemini Pro (Anil et al., 2023), which we call LMDXPaLM 2-S and LMDXGemini Pro respectively.
Firstly, starting from their original checkpoint, we finetune those LLMs on the prompts and completions detailed in Section 3.2 and 3.3 on a data mixture containing a variety of (document, schema, extraction) tuples. In particular, this data mixture contains the Payment dataset (Majumder et al., 2020), along with a diverse set of publicly available PDF form templates obtained from government websites that we filled with synthetic data using an internal tool, and annotated for schema and entities to extract. The goal of this tuning is to obtain a Base Entity Extractor checkpoint by training the model to learn the IE task along with our desired extraction syntax. No document or schema contained in the base extraction training phase overlap with the documents and schemas used in our target benchmarks, hence we use those LLMs for zeroshot information extraction evaluation on the target benchmarks.
Finetuned Performance. We are also interested in evaluating how data-efficient LMDX is (e.g. how quickly it learns information extraction on a new document type). To answer this, starting from the Base Entity Extractor checkpoint, we finetune the LLM directly on the target benchmark. Parameters. For training, we finetune using a batch size of 8, a dropout probability of 0.1 and a learning rate of 10−6 with a standard cross-entropy loss for 4000 steps on TPUv4 (Jouppi et al., 2023). Once training is done, for finetuned experiments, we select the checkpoint with the lowest loss on the dev set, and report performance on the test set. For LLM inference, we use a temperature of 0.5 and a TopK of 40, sampling 16 responses for each chunk processed by the LLM, as described in Section 3.4. Finally, for both training and inference, we use an input token length of 6144 and output token length of 2048. We use line-level segments and only two coordinates [xcenter, ycenter] with 100 quantization buckets, as supported by Appendix A.12.

heading: 4.1 Benchmarks
text: Visually Rich Document Understanding (VRDU). Wang et al. (2023d) introduces two public visually-rich documents IE benchmarks: Registration Form, containing 6 semantically rich entity types, and Ad-buy Form, containing 14 entity types with one hierarchical line_item entity type. For each benchmark, VRDU proposes samples of 10, 50, 100 and 200 train documents with high-quality OCR1 which we use to evaluate the data efficiency of LMDX. It also offers different tasks to evaluate the generalization powers of extraction systems: Single Template Learning (STL) where train/test splits share the same single template, Unseen Template Learning (UTL) where train/test contain disjoint sets of templates, and Mixed Template Learning (MTL) where train/test contain overlapping sets of templates. We report Micro-F1 through the provided evaluation tool. For VRDU, we only run the finetuning experiments using LMDXPaLM 2-S, given the significant cost of finetuning on all its tasks and train split sizes. Consolidated Receipt Dataset (CORD).2 Park et al. (2019) introduces a benchmark of Indonesian receipts from shops and restaurants, with a target schema of 30 fine-grained entities, grouped into menu, total and subtotal hierarchical entities. We adopt the evaluation tool from prior work (Kim
1https://cloud.google.com/vision/docs/ocr 2https://huggingface.co/datasets/naver-clova-ix/cord-v1
et al., 2022) and report Micro-F1 on that benchmark. For our experiments, we use the official train (|D| = 800), dev and test splits. To evaluate LMDX’s data efficiency, we further sample the first |D| = 10/50/100/200 documents from the train split. For each of those data setup, we finetune LMDX for 12000 steps. For comparison, we also train and evaluate state-of-the-art baselines LayoutLMv3LARGE and Donut. Those baselines are detailed in Appendix A.7. Baselines In the zero-shot setting, we compare LMDX to other LLM baselines: GPT-3.5, Gemini Pro (Anil et al., 2023), and PaLM 2-S (Google et al., 2023) that we prompt with the raw OCR text and IE instruction (called GPT-3.5+OCR, Gemini Pro+OCR and PaLM 2-S+OCR respectively). We also compare LMDX with strong Vision-Language models: LLaVA-v1.5-13B (Liu et al., 2023), Gemini Pro, and GPT-4V (OpenAI, 2023a) that we prompt with the document page image and IE instructions (called LLaVA-v1.5-13B+Image, GPT4V+Image, and Gemini Pro+Image). Those baselines are fully detailed in Appendix A.6. Unlike LMDX, those large model baselines do not localize their predictions. In the finetuned setting, we compare LMDXPaLM 2-S to popular VRD IE baselines. For VRDU, we compare to its published baselines (Wang et al., 2023d), LayoutLM/v2/v3 and FormNet. For CORD, we train and evaluate state-of-the-art baselines LayoutLMv3 and Donut. Those baselines are detailed in Appendix A.7. Following prior work (Lee et al., 2022, 2023a; Xu et al., 2020, 2021), for all models leveraging the text modality, we use the benchmarks’ provided OCR, ensuring a fair comparison.

heading: 4.2 Results
text: VRDU results are presented in Table 2. In zero-shot (|D| = 0) setting, LMDXPaLM 2-S and LMDXGemini Pro have higher extraction quality than all other large models baselines, including the ones using the same LLM and OCR or image, showing improvements brought by the LMDX methodology itself. In finetuned setting on VRDU, LMDXPaLM 2-S is much more data efficient than the baselines: it is at 5.06% Micro-F1 of its peak performance at 10 training documents for Registration Form Mixed Template (87.72% vs 92.78% MicroF1) while LayoutLMv2, the strongest finetuned baseline, is within 19.75% of its peak performance (69.44% vs 89.19% Micro-F1), showcasing that it learns extraction on a new document type much
faster. Moreover, LMDXPaLM 2-S generalizes better to unseen templates than finetuned baselines: on Registration Form, LMDXPaLM 2-S has a drop less than 5% Micro-F1 on Unseen Template compared to Single Template across data regimes, while LayoutLMv2 see a drop between 19% and 27%.
On CORD (in Table 3), we observe similar trends, reaching state-of-the-art on all but one data regime, highlighting the generality of the results.
Performance on Hierarchical Entities. To showcase extraction quality on hierarchical entities, we display in Table 2 the F1 score on Ad-buy Form Mixed’s line_item entity type. Overall, LMDX has much higher Line Item F1 than the baselines for all data regimes. In particular, LMDXPaLM 2-S has similar Line Item F1 at zero-shot than the best finetuned baseline at 200 train documents (21.21%
versus 25.46% respectively). With all the training data, LMDXPaLM 2-S scores a 72.09% Line Item F1, an absolute improvement of 46.63% over the best baseline LayoutLMv2. Finally, as LMDX encodes the layout modality, it possesses much higher zero-shot Line Item F1 than large models baselines. Entity Localization Accuracy. In order to evaluate the localization quality independently of the extraction quality, we compute the Localization Accuracy of LMDX and all baselines that can localize entities using the formula: AccuracyLocalization = NE+L NE
where NE+L is the number of entities correctly extracted and localized, and NE is the number of entities correctly extracted. Since LMDX localizes at the line level, localization verification is done at the line-level as well, i.e. localization is considered correct if the prediction bounding box is covered by the groundtruth line-level bounding box by more than 80%. We present the results in Table 4. Overall, LMDXPaLM 2-S and LMDXGemini Pro can localize their predictions reliably at the linelevel with the segment identifiers, with 88%-94% accuracy at zero-shot, and 98%-99% in finetuned cases, which is slightly higher than baselines.

heading: 4.3 Ablation Study
text: In this section, we ablate different facets of the LMDX methodology to highlight their importance. The results can be found in Table 5. For all ablations, we evaluate LMDXPaLM 2-S on the VRDU Ad-Buy Form Mixed Template task at |D| = 10 data size, only changing the ablated facet.

heading: Effects of Base Entity Extraction Training. In
text: this ablation, we remove the initial training on the varied data mixture and directly finetune on the VRDU target task. As seen in Table 5, skipping that training leads to -11.44% micro-F1 as the model has to learn from scratch the task, the desired com-
pletion format and coordinate tokens’ semantics.
Effects of Coordinate Tokens. In this ablation, we replace the coordinate tokens, which communicate the position of each line within the document, by the index of that line. This index still acts as a unique identifier for the line segment (required for entity localization) but does not communicate any position information. An example of a prompt with line index can be found in Appendix A.8 Figure 11, and per-entity F1 can be found in Appendix A.14 Table 9. As seen in Table 5, the coordinate tokens are crucial to quality, leading +14.98% micro-F1.
Effects of Sampling Strategy. In this ablation, we discard our strategy of sampling K = 16 completions per chunk, and instead sample a single response. As seen in Table 5, this leads to a 1.5% drop in micro-F1. While overall minor for quality, the sampling strategy corrects extraction format mistakes (see parsing error rates in Appendix A.10), leading to a successful extraction on all documents.
Effects of Missing Entity Types. In this ablation, we study the effect of having the model’s completions skip missing entity types in the completions instead of explicitly outputting ”type” : null for those (See example in Appendix A.8, Figure 12). As seen in Table 5, this leads to a 6.77% drop in micro-F1 over explicitly outputting missing types. We hypothesize that this is due to the fact that having completions skip missing types means that, during response generation, the model has to choose with a single token computation budget within the N remaining entity types which one is the next present (essentially a N -way classification). Explicitly emitting missing entity types means the model only has to copy the types directly from the schema in the prompt, and has to do 2-way classification within a single token computation budget to declare if an entity is present or not (e.g. emit token null if entity is missing or ” if present), which is an easier task.

heading: 4.4 In-context Learning Performance
text: In this section, we study how in-context learning (ICL) compares to finetuning. To do so, we test two methodologies: Random, which randomly selects |D| documents and extractions from the train set, and Nearest Neighbors, which uses similarity based on SentenceT5 embeddings (Ni et al., 2021) to retrieve |D| documents to add in the LLM context. The results on CORD for LMDXPaLM 2-S and LMDXGemini Pro are shown in Figure 3. Overall, while both methods increase the performance significantly, nearest neighbors shows a clear advantage, matching the best random ICL performance with only a single in-context example (86.43% versus 86.57% micro-F1 for LMDXPaLM 2-S), and matching the finetuned performance at |D| = 10 examples (90.33% versus 90.02% micro-F1), as examples from the same template are retrieved (see Appendix A.9). Beyond |D| = 10, the quality plateaus as no more example fit in the prompt.

heading: 5 Conclusion
text: In this paper, we have introduced LMDX, a methodology that enables using LLMs for information extraction on visually rich documents. With its coordinate tokens and decoding strategy, LMDX allows the high-quality extraction of singular, repeated and hierarchical entities, while localizing the entities in the document. LMDX is data efficient, and even allows extraction at zero-shot on entirely new document types and schemas. LMDX can benefit from orthogonal research fields, and we continue the discussion in the Limitations section.

heading: Acknowledgements
text: The authors would like to thank Tania BedraxWeiss, Riham Mansour, Slav Petrov, Yunhsuan Sung, Mike Kwong and Chun-Liang Li for their valuable feedback on the experiments and paper. The authors also thank Nikolai Glushnev for discussions, along with the help naming LMDX.

heading: Limitations
text: We acknowledge the limitations of LMDX from the following aspects to inspire future research in the field of information extraction and localization.
Firstly, LMDX’s input is text lines and their bounding boxes, usually coming from OCR. This means that LMDX can not extract non-textual entities thus would not be able to extract an entity that would be an image embedded in a document (e.g. a product_image entity in a product webpage). This also limits performance in high-data scenarios, as all page image information is discarded. Furthermore, such input means that LMDX is sensitive to errors from the OCR process (wrong reading order, incorrect line grouping, undetected text and erroneously recognized characters). Qualitatively, we have found that a common error type for LMDX is caused by OCR grouping multiple semantically different segments together (we give a deeper analysis with concrete examples in Appendix A.11). Thus, techniques aiming to improve this error type would be a worthwhile future research direction.
Moreover, LMDX’s localization mechanism is applied at the line level, where we verify that the predicted text is indeed present on the line. If the entity text appears multiple times on the line, we don’t have a definitive way to choose the correct text. Thus, LMDX’s localization and bounding boxes are not reliable beyond line-level granularity. While sufficient for greatly speeding up human-in-the-loop interactions like prediction auditing/review, getting entity bounding boxes precise at character-level would be beneficial, more natural-looking, and a worthy research direction.
Lastly, LMDX relies on LLMs supporting thousands of tokens in input and output (as detailed in Appendix A.4), which is both computationally expensive and slow, requiring the use of hardware acceleration for acceptable latency and throughput. We showcase a latency comparison between popular solutions in Appendix A.13. General research in accelerating LLM inference (Shazeer, 2019; Ainslie et al., 2023; Leviathan et al., 2023; Hong et al., 2023) would make LMDX more costeffective in production setting. Specifically for LMDX, the coordinate tokens represent a large part of the total number of tokens, so research on minimizing their number (e.g. by introducing dedicated coordinate tokens within the LLM vocabulary) would yield significant savings and be a worthwhile direction for future work.

heading: A Appendix

heading: A.1 Chunking algorithm
text: Algorithm 1 Document Chunking 1: function CHUNK(D, L, F ) . D is a document containing multiple pages. L is token limit. 2: . F is a function that outputs prompt token length given some segments (e.g. lines). 3: C = φ . C is to record all produced chunks. 4: for i = 1 to |D.pages| do 5: S = D.pages[i].segments 6: while S 6= φ do 7: for j = |S| to 1 do . Start pruning from the end of the page. 8: if F (S[1 : j]) ≤ L then 9: C = C ∪ {S[1 : j]} 10: S = S[j + 1 : |S|] . Continue to work on the rest of the segments. 11: Exit for loop 12: end if 13: end for 14: end while 15: end for 16: return C 17: end function
A.2 Entity value parsing algorithm
Algorithm 2 Entity Value Parsing 1: function PARSEENTITYVALUE(D, E) . D is a document chunk. 2: . E is raw extraction results for one entity type parsed from one LLM sample. 3: G = φ . G is to record all parsed entity values. 4: R = Regex(“(\d\d\|\d\d)”) . R is a regex that captures the segment identifiers. 5: M = {“s.x|s.y” 7→ s|s ∈ D.segments} . M holds a mapping between segment id and segment. 6: for i = 1 to |E| do 7: W = φ . W is to hold all words for this entity. 8: P = R.split(E[i]) . P is expected to be interleaved text values and segment ids. 9: for j = 1 to |P |/2 do 10: if P [j ∗ 2] /∈M then 11: Go to next i . Segment ID is hallucinated. Grounding failure. 12: end if 13: S = M [P [j ∗ 2]] . Retrieve the stored segment from M with parsed segment ID. 14: T = P [j ∗ 2− 1] . T is to hold the predicted text. 15: if T not substring of S then 16: Go to next i . Grounding failure, skip the current entity. 17: end if 18: W = W ∪ (S ∩ T ) 19: end for 20: G′.value = ⋃ w∈W w.text_value . G
′ is to hold the entity to return. 21: G′.bounding_box = {min(b.x),min(b.y),max(b.x),max(b.y)}w∈W,b=w.bounding_box 22: G = G ∪ {G′} 23: end for 24: return G 25: end function

heading: A.3 Decoding algorithm
text: Algorithm 3 Responses Decoding 1: function DECODEFORTYPE(J , T , D) . J is one or more JSON objects. 2: . T is an entity type. 3: . D is a document chunk. 4: E = φ . E is to record all parsed and grounded entities. 5: for j = 1 to |J | do 6: J ′ = J [j][T.type] . J ′ is to hold entities for T’s type before grounding. 7: if T.subtypes = φ then . T is leaf entity type. 8: E = E ∪ ParseEntityV alue(D, J ′) 9: else . T is hierarchical entity type. 10: E′.subtypes = ⋃
T ′∈T.subtypesDecodeForType(J ′, T ′, D) . E′ is hierarchical entity.
11: E = E ∪ {E′} 12: end if 13: end for 14: return E 15: end function 16: 17: function MAJORITYVOTING(T , E) . T is an entity type. 18: . E is a 2D vector of entities of type T from all LLM responses. 19: V = [0, 0, ..., 0] ∈ R|E| . V is to record all votes. 20: L = {T} 21: while L 6= φ do 22: T ′ = L[0] 23: E′ = φ 24: for j = 1 to |E| do 25: E′ = E′ ∪ {e|e ∈ E[j], e.type = T ′} . E′[j] holds entities with type T ′ from E[j]. 26: end for 27: for i = 1 to |E′| - 1 do 28: for j = i+ 1 to |E′| do 29: if E′[i] = E′[j] then 30: V [i] = V [i] + 1 31: V [j] = V [j] + 1 32: end if 33: end for 34: end for 35: L = L[1 : |L|] . Remove T ′ and inject its sub-types for recursion. 36: L = L ∪ T ′.subtypes 37: end while 38: return E[argmax(V )] . Return the entity values with the highest votes. 39: end function 40: 41: function DECODEALLSAMPLES(S, T , D) . S is all LLM response samples on D. 42: . T is a list of entity types. 43: . D is a document chunk. 44: return ⋃ T ′∈T MajorityV oting( ⋃ S′∈S DecodeForType(ParseJson(S
′), T ′, D)) 45: end function

heading: A.4 Token Length Statistics
text: Table 6 details the token length (50th and 99th percentiles) of the prompt and completion targets for LMDXPaLM 2-S for the train split of datasets used in our experiments. We select the line level segment, 2 coordinate scheme, no JSON indentation so that all datasets fit within our 6144 prompt token length and 2048 output token length.

heading: A.5 Schemas
text: In this section, we present the schemas used for the experiments of this paper. The schema for VRDU Ad-Buy Form, VRDU Registration Form, and CORD can be found in Figure 4, Figure 5 and Figure 6 respectively.

heading: A.6 Zero-shot Baselines Details
text: We compare LMDX to other Large Model baselines on all benchmarks in the zero-shot context. Those baselines are detailed below. Text-based Baselines. We evaluate the zero-shot extraction ability of multiple strong Large Language Models: GPT-3.53, Gemini Pro (Anil et al., 2023) and PaLM 2-S (Google et al., 2023). To do so, we prompt them with the raw benchmark’s OCR text (no coordinate tokens or segment identifier like for LMDX), and extraction instructions alongside the schema in JSON format. We then parse the completions as JSON to get the predicted entities directly. Note that those predicted entities are not localized within the document. A sample prompt can be observed in Figure 7. In particular for GPT-3.5, we use the gpt-3.5-turbo-1106 through OpenAI’s API. Page Image-based Baselines. We evaluate the zero-shot extraction ability of multiple strong Vision-Language Models: LLaVA-v1.5-13B (Liu et al., 2023), Gemini Pro (Anil et al., 2023), and GPT-4V (OpenAI, 2023a). The prompt includes task description, instructions and target schema represented in JSON format as text input and the document page as image input. We provide examples of valid JSON values in the task instructions. Note that those predicted entities are not localized within the document. A sample prompt can be observed in Figure 8. In particular for GPT-4V, we use the gpt-4-1106-preview through OpenAI’s API. For Gemini Pro, we use gemini-pro-vision through the VertexAI API4.

heading: A.7 CORD Baselines Details
text: LayoutLMv3 Baseline. We follow the released implementation5 for the LayoutLMv3LARGE model and the training protocol described in (Huang et al., 2022) as closely as possible. In particular, we train the model for 80 epochs for each experiment on CORD (namely, 10, 50, 100 and 200-document training sets), on the IOB tags of the leaf entities. One difference in our training is that, due to computational resource constraints, we use batch_size = 8 and learning_rate = 2 · 10−5.
As the LayoutLMv3 model can only extract leaf entities, we design and heavily optimize a heuristic algorithm to group the leaf entities into hierarchical
3https://platform.openai.com/docs/models/gpt-3-5-turbo 4https://cloud.google.com/vertex-ai/docs/generative-
ai/start/quickstarts/quickstart-multimodal 5https://github.com/microsoft/unilm/tree/master/layoutlmv3
entities menu, subtotal and total. The best heuristics we could find are as follows:
• For the subtotal and total hierarchical entity types, since they appear only once per document, we group all their extracted sub-entities under a single subtotal and total entity, respectively.
• For menu hierarchical entity type, we observe that those entities usually occur multiple times on a document, and each menu has at most one nm, num, unitprice, cnt, discountprice, price, itemsubtotal, etc sub-entities and potentially multiple sub_nm, sub_price and sub_cnt subentities. We also notice that the sub-entities aligned horizontally overwhelmingly belong to the same menu entity, and a menu entity can sometimes span over two or more consecutive horizontal lines. To leverage those observations, we perform a two-step grouping process for menu entities. First, we merge the extracted leaf sub-entities into horizontal groups, where a threshold of 0.5 on the intersection-over-union of the Y-axis was used for the determination of horizontal alignment. Second, we further merge the consecutive horizontal groups into menu entities, if and only if the horizontal groups do not have type duplication in any of the nm, num, unitprice, cnt, discountprice, price, itemsubtotal, and etc subentities (namely, those sub-entities only show up in at most one of the consecutive horizontal groups to be merged). We allow duplication of sub_nm, sub_price and sub_cnt sub-entity types. After those two steps, we obtain the final menu entities.
Donut Baseline. We follow Donut released implementation6 for the Donut benchmarking results on CORD. We use the default training configuration for all experiments on CORD (namely, 10, 50, 100 and 200-document training sets), with the following difference: we reduce batch size from 8 to 4 due to computational resource constraints, and increase the number of train epochs from 30 to 60. For each experiment, checkpoint with the lowest loss on the dev set is selected and we report performance on test set. Micro-F1 scores produced by Donut evaluation code are reported (similar to all our other models).
6https://github.com/clovaai/donut

heading: A.8 Sample Prompts and Completions
text: In this section, we present example of LMDX prompts and completions from the LLM on the VRDU Ad-Buy dataset to better showcase the format used. Figure 9 shows the original document with the line bounding boxes from OCR, Figure 10 shows the corresponding prompt and completion on that document with coordinate segment identifiers, Figure 11 shows the same prompt and completion, but with line index segment identifiers (used in ablation studies to showcase how the LLM can interpret the layout) and Figure 12 shows how the completion changes when skipping missing entity types entirely in the completion. Finally, Figure 13 and Figure 14 show how the prompts/completions changes if 4 line-level [xmin, ymin, xmax, ymax] coordinates or 2-word level [xcenter, ycenter] coordinates are used.

heading: A.9 Nearest Neighbors In-Context Learning
text: In our study, nearest neighbors leads to a significant quality gain over randomly selecting examplars. In this section, we explore why that is the case in the context of VRD information extraction. Figures 15, 16 and 17 show typical retrievals using sentenceT5 (Ni et al., 2021) embeddings7 on the OCR text for similarity. Unsurprisingly, nearest neighbors works well as it retrieves examplars from the same template as the target document, i.e. from the same merchant in the case of CORD documents (store/restaurant receipts). As those examples share the same layout, boilerplate text, and entities, it makes it easier for the model to understand the correct extraction pattern, leading to significant gains in quality. 7https://www.kaggle.com/models/google/sentencet5/frameworks/tensorFlow2/variations/st5-base

heading: A.10 Completion Parsing Error Rates
text: In this section, we report the various completion parsing error types and their occurrence rates for LMDXPaLM 2-S. Invalid JSON Formatting. This refers to cases for which Python’s json.loads(completion) fails on a LLM’s completion. As observed in Table 7, the JSON parsing error rate is below 0.3% in all training settings. Invalid Entity Value Format. This refers to cases where the leaf entity value does not follow the expected "<text-segment-1> XX|YY <text-segment-2> XX|YY" format. As observed in Table 7, the Invalid Entity Value Format Rate is below 0.05% in all training settings. Hallucination / Entity Text Not Found. This refers to cases where the segment identifier is valid, but the entity text does not appear on the predicted segment (hallucination). As observed in Table 7, the Entity Text Not Found error rate is below 0.6% in all training settings. As part of LMDX methodology, we discard any prediction whose text does not appear on the specified segment, ensuring we discard all hallucination.
Note that those numbers are computed at the completion level. Since multiple completions are sampled for each document chunk, the sampling scheme allows for correcting those errors and no document in the benchmarks fail extraction.

heading: A.11 Error Analysis
text: In this section, we perform an error analysis on the test set to identify common error patterns of LMDX. A very common error type we observe is caused by OCR lines grouping multiple semantically different segments. We show two instance of those cases observed in LMDXPaLM 2-S on the VRDU Ad-Buy Form in Figure 18. In the first example, prediction
for the entity line_item/program_desc includes text from the previous column "Channel" along with the value in the column "Description". From the OCR line bounding boxes, we can see that these two columns are grouped as the same OCR line. In the second example, the model confuses between the adjacent keys "Invoice Period" and "Flight Dates" and extracts invoice dates as flight dates. Similar to the first example, OCR line bounding boxes show that the invoice dates and the key "Flight Dates" are grouped together in the same line although they are semantically different. As LMDXPaLM 2-S uses only coarse line layout information ([xcenter, ycenter] with 100 quantization buckets), the model fails in these cases, which is a current limitation of LMDX. We believe that incorporating the image modality will make LMDX more robust to those OCR errors.

heading: A.12 Effect of Coordinate Tokenization Schemes
text: In this section, we study how different designs of coordinate tokens affect extraction quality, to determine which is the most effective. There are many ways to tokenize the coordinates of the text segments, e.g. using line versus words as segment, the number of coordinates to use per segment ([xcenter, ycenter] versus [xmin, ymin, xmax, ymax]). To establish which
scheme is the most effective, we evaluate the zeroshot performance of LMDXPaLM 2-S on VRDU Registration Form, Ad-Buy Form and CORD benchmarks with the following schemes:
• 2 Line-level [xcenter, ycenter] coordinates with B = 100 quantization buckets.
• 4 Line-level [xmin, ymin, xmax, ymax] coordinates with B = 100 quantization buckets.
• 2 Word-level [xcenter, ycenter] coordinates with B = 100 quantization buckets.
An example of prompt and completion for each scheme is given in Appendix A.8. Results quality is given in Figure 19.
On all benchmarks, 2-line level [xcenter, ycenter] coordinates obtains the best performance. This is caused by that coordinate tokenization scheme being very token efficient (see Appendix A.4). 2 Word-level [xcenter, ycenter] coordinates and 4 Linelevel [xmin, ymin, xmax, ymax] coordinates increase the tokens used by coordinates drastically, leading to multiple chunks generated for each document page, hence lowering quality since the model won’t be able to correctly predict entities spanning multiple chunks. Furthermore, the large number of coordinate tokens added by those schemes makes the text sequences used in LMDX significantly different from the natural text sequences the LLM was pretrained on, hence making the LLM less effective at interpreting them, thereby lowering extraction quality.

heading: 2 Word Coord. 4 Line Coord. 2 Line Coord.

heading: A.13 Latency Comparison
text: In Table 8, we compare the latencies of LMDX and other methodologies, LayoutLMv3 and Donut. Overall, LayoutLMv3 is the fastest model, as it is encoder-only. LMDXGemini Pro is the slowest among accelerated models, due to its reliance on an
LLM with long inputs and outputs. Nonetheless, its latency remains acceptable in a production setting.

heading: A.14 Ablation Detailed Metrics.
text: In Table 9, we present the per-entity F1 score for the coordinate ablations to highlight which entity most benefit from the introduction of layout information, and explain the large difference observed.
Overall, performance is lower across the board without coordinates, as the VRDU Ad-Buy form benchmark consists entirely of invoices, a very layout-heavy document, where the alignment of the key segments (e.g. "Contract Number: ", "Total Due:", "Product:") and value segments (e.g. "123456", "$1000", "Political Advertisement") matters a lot for the correct understanding of most entities. Thus, removing the layout modality will have a stark difference in performance. However, not all entities are affected in the same way:
• 8 out of 9 leaf entities see single digit drops, with gross_amount affected the least. gross_amount can be derived mostly without
looking at its key (e.g."Amount Due") and simply using cues commonly found in money amount entities (e.g. "$", "USD", etc).
• line_item, the hierarchical entity in AdBuy form, is affected the most. This is expected as line_item’s components (channel, program_desc, program_start_date, program_end_date, sub_amount) are always visually arranged in tables, hence understanding horizontal and vertical alignments of the different text lines (which the coordinates provide) are critical for the correct grouping of those components into line_items.
• Since line_items are by far the most common entity in VRDU Ad-Buy form (9163 occurrences), they have the most effect on the overall micro-F1.

heading: A.15 Datasets Details
text: As part of this work, we used the following datasets:
• VRDU (Wang et al., 2023d) consists of two public visually-rich documents information extraction benchmarks: Registration Form, which consists of 1915 documents recording foreign agents’ activities with the US government requiring public disclosure, and Ad-Buy Form, which consists of 641 signed receipts between a TV station and a political campaign group, requiring public disclosure to the Federal Communications Commission. All data is in English, and publicly available for research purposes.
• CORD (Park et al., 2019) consists of 1000 Indonesian receipts from various shops and restaurants, on which the authors de-identified identifiable information. The data is a mix of Indonesian and English languages, and is publicly available for research purposes.
• Payment (Majumder et al., 2020) is a dataset of 14,237 single-page invoices, annotated with Each invoice is from a different vendor, so the documents do not share any common templates. All the data is in English.
• Synthetic Forms. Finally, we generated synthetic training data using blank PDF form templates crawled from government websites, and filled with fully synthetic values using an internal tool. All generated data is in English.
All extraction benchmarks we used are publicly available for research purposes, and we limited their usage to this research work. We also checked these datasets to ensure that no personally identifiable information is involved.

heading: Future Work
text: We acknowledge the limitations of LMDX from the following aspects to inspire future research in the field of information extraction and localization. Firstly, LMDX’s input is text lines and their bounding boxes, usually coming from OCR. This means that LMDX can not extract non-textual entities thus would not be able to extract an entity that would be an image embedded in a document (e. g. a product_image entity in a product webpage). This also limits performance in high-data scenarios, as all page image information is discarded. Furthermore, such input means that LMDX is sensitive to errors from the OCR process (wrong reading order, incorrect line grouping, undetected text and erroneously recognized characters). Qualitatively, we have found that a common error type for LMDX is caused by OCR grouping multiple semantically different segments together (we give a deeper analysis with concrete examples in Appendix A. 11). Thus, techniques aiming to improve this error type would be a worthwhile future research direction. Moreover, LMDX’s localization mechanism is applied at the line level, where we verify that the predicted text is indeed present on the line. If the entity text appears multiple times on the line, we don’t have a definitive way to choose the correct text. Thus, LMDX’s localization and bounding boxes are not reliable beyond line-level granularity. While sufficient for greatly speeding up human-in-the-loop interactions like prediction auditing/review, getting entity bounding boxes precise at character-level would be beneficial, more natural-looking, and a worthy research direction. Lastly, LMDX relies on LLMs supporting thousands of tokens in input and output (as detailed in Appendix A. 4), which is both computationally expensive and slow, requiring the use of hardware acceleration for acceptable latency and throughput. We showcase a latency comparison between popular solutions in Appendix A. 13. General research in accelerating LLM inference (Shazeer, 2019; Ainslie et al. , 2023; Leviathan et al. , 2023; Hong et al. , 2023) would make LMDX more costeffective in production setting. Specifically for LMDX, the coordinate tokens represent a large part of the total number of tokens, so research on minimizing their number (e. g. by introducing dedicated coordinate tokens within the LLM vocabulary) would yield significant savings and be a worthwhile direction for future work