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.

## introduction the recent advent of transformers @xcite 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 @xmath0.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 @xcite @xcite @xcite classifies each token into inside-outside-begin (iob) tags @xcite , which support extraction and localization of leaf entities. however, it is non-trivial to extend these meth-ods for hierarchical entities. sequence-generation based methods @xcite treats extraction as text generation with autoregressive decoders @xcite . 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) @xcite @xcite have revolutionized natural language processing, showing the capabilities to solve diverse tasks with an instruction @xcite or a few demonstrations attached to the prompt @xcite . this paradigm shift opens the possibility of extracting entities while addressing all the aforementioned desiderata, but using llms for vrd ie has been underexplored. @xcite 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 @xcite . 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 ## lmdx methodology 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. ## evaluation 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 @xcite , which we call lmdx palm 2-s and lmdx gemini 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 @xcite . 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 top k 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 [x center , y center ] with 100 quantization buckets, as supported by appendix a.12. ## conclusion 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.