LMDX

Analysis

Why This Paper Matters

Document information extraction (IE) is a critical task for automating workflows in finance, legal, healthcare, and logistics. Traditional approaches either rely on layout-specific models (e.g., LayoutLM, FormNet) that require extensive pre-training on document images, or on generic LLMs that struggle with spatial understanding and often hallucinate. LMDX bridges this gap by providing a lightweight methodology to adapt any off-the-shelf LLM for document IE, incorporating layout information through coordinate tokens and enforcing grounding to eliminate hallucinations. This is particularly impactful for practitioners who need to extract structured data from invoices, forms, or receipts without building custom models from scratch.

How It Works

LMDX operates in several stages. First, an OCR engine extracts words and their bounding boxes from the document image. The document is then split into chunks small enough for the LLM's context window.

Prompt Generation is the core innovation. Each chunk is converted into a prompt containing:

• Document Representation: Segment texts concatenated with coordinate tokens in the format <segment text> XX|YYsegment. Coordinates are normalized and quantized into B buckets (B=100 in practice), using line-level segments with [x_center, y_center].
• Task Description: A short instruction like "From the document, extract the text values and tags of the following entities:"
• Schema Representation: A structured JSON defining entity types, their occurrence (single/multiple), and hierarchy. For example, {"foo": "", "bar": [{"baz": []}]}.

The LLM is prompted to output a completion in the same JSON structure, where leaf entity values include both the extracted text and the segment identifier (e.g., <text> XX|YYsegment). Missing entities are filled with null (singular) or [] (repeated).

Inference and Decoding: For each chunk, K completions are sampled via top-K sampling. Raw outputs are parsed as JSON; invalid ones are discarded. Entity values are grounded by verifying that the extracted text exactly matches the referenced segment. Predictions from multiple completions are merged via majority voting per chunk, then concatenated across chunks. For hierarchical entities, voting is flattened across all levels.

Training: The authors fine-tune PaLM 2-S in two phases. First, a base entity extractor is trained on a mixture of synthetic and public documents to learn extraction syntax and schema semantics. Second, the model is specialized on the target benchmark with no overlap in documents or schemas.

Results

LMDX with PaLM 2-S achieves state-of-the-art on VRDU across all regimes. Notably, it matches the performance of baselines trained on 10-50 documents even in zero-shot settings. For example, on Ad-Buy Form Mixed Template, it scores 39.74% Micro-F1 vs. 40.68% for FormNet with 50 training documents. On Registration Single Template, it reaches 73.81% Micro-F1 vs. 74.22% for FormNet with 10 documents.

Data efficiency is a standout: LMDX reaches 5.06% of its peak performance at just 10 training documents for Registration Form Mixed Template, while LayoutLMv2 only achieves 19.75% of its peak. Generalization to unseen templates is also superior, with less than 5% Micro-F1 drop compared to 19-27% for LayoutLMv2.

On CORD, LMDX is only 4.03% from its peak performance at 10 documents, versus 22.34% for LayoutLMv3_LARGE. For hierarchical entity grouping, LMDX achieves 72.09% F1 with full data, a 46.63% improvement over LayoutLMv2.

Significance

LMDX democratizes document IE by enabling any LLM to handle layout and grounding without architectural changes. Its data efficiency makes it practical for real-world scenarios where labeled documents are scarce. The grounding mechanism is a critical safety feature for production systems, as it prevents hallucinated extractions. Future work could explore extending LMDX to multi-page documents or integrating with vision-language models for richer layout understanding.