PM 3- KIE : A Probabilistic Multi-Task Meta-Model for Document Key Information Extraction (2025)

Abstract: Key Information Extraction (KIE) from visually rich documents is commonly approached as either fine-grained token classification or coarse-grained entity extraction. While token-level models capture spatial and visual cues, entity-level models better represent logical dependencies and align with real-world use cases.We introduce PM3-KIE, a probabilistic multi-task meta-model that incorporates both fine-grained and coarse-grained models. It serves as a lightweight reasoning layer that jointly predicts entities and all appearances in a document. PM3-KIE incorporates domain-specific schema constraints to enforce logical consistency and integrates large language models for semantic validation, thereby reducing extraction errors.Experiments on two public datasets, DeepForm and FARA, show that PM3-KIE outperforms three state-of-the-art models and a stacked ensemble, achieving a statistically significant 2% improvement in F1 score.

## introduction key information extraction (kie) focuses on identifying structured key-value pairs from visually rich documents (vrds) @xcite based on a predefined schema that specifies the target information types. this capability is essential for automating business document processing across industries such as finance and law. automating information extraction significantly reduces operational costs: processing a single invoice, for example, can cost $13.11 and takes up to eight days @xcite . despite recent advances, extracting structured data remains challenging for state-of-the-art models, particularly for documents with complex schemas or semi-structured layouts @xcite . kie can generally be approached through two distinct paradigms, as illustrated in figure 1 : fine-grained token classification the finegrained task operates at the token level (illustrated at the bottom of figure 1 ), where a label is assigned to each token. fine-grained models typically leverage multi-modal transformer encoders that integrate textual, spatial, and visual features, such as layoutlmv3 @xcite and lilt @xcite . while effective in capturing tokenlevel spatial and visual relationships, these models require large annotated datasets for training. coarse-grained entity extraction in contrast, coarse-grained entity extraction targets high-level entities within a document (illustrated at the top of figure 1 ). coarse-grained models may structure outputs as key-value pairs in formats such as @xcite , retrieve entity values @xcite , or classify entities among predefined candidates @xcite . these models excel at capturing logical dependencies and generating coherent outputs and align more closely with business document processing needs like the extraction of invoice elements. both approaches offer distinct advantages, highlighting the need for models that effectively integrate them. @xcite address this by proposing a knowledge distillation framework that combines knowledge from both model types. however, their method relies on model classification logits, restricting it to specific architectures and making it incompatible with coarse-grained models that generate structured output, such as recent large language model (llm)-based approaches moreover, existing methods often fail to enforce structural constraints defined by the extraction schema. certain elements, such as document identification numbers, are mandatory, while others are optional. current models frequently omit required information, leading to schema violations. while heuristic post-processing can mitigate some of these errors, it does not ensure full consistency. additionally, state-of-the-art approaches struggle with syntactic errors, such as incomplete address information @xcite and semantic errors, such as implausible entity values (see appendix f). to address these limitations, we introduce pm3-kie, a probabilistic multi-task meta-model. pm3-kie integrates fine-grained and coarse-grained approaches using a lightweight probabilistic reasoning layer based on probabilistic soft logic (psl) @xcite , ensuring robust inference while maintaining computational efficiency. pm3-kie incorporates four key innovations: we evaluate pm3-kie on two public datasets, deepform and fara @xcite , demonstrating significant improvements over state-of-theart models and a stacked ensemble. our experiments cover out-of-distribution and low-resource settings, highlighting pm3-kie's robustness across diverse real-world challenges such as small training sizes and unknown document formats during application time. our contributions can be summarized as following: the remainder of this paper is structured as follows: section 2 reviews related work, followed by an outline of the kie problem in section 3. next, section 4 describes the pm3-kie architecture, while section 5 presents experimental results. finally, section 6 concludes with a discussion. ## kie problem statement in this section, we first review traditional kie tasks: fine-grained token classification and coarse-grained entity extraction. we then present a unified problem formulation that combines these tasks into a joint extraction framework, enabling the simultaneous extraction of entities and their mentions in a document. finally, we extend this to a metamodeling task for joint prediction conditioned on fine-grained and coarse-grained models. ## conclusion and discussion we propose pm3-kie, a probabilistic multi-task meta-model for kie that integrates fine-grained token classification and coarse-grained entity extraction models. pm3-kie flexibly combines arbitrary models through a probabilistic reasoning layer that enforces schema consistency, reducing errors such as missing required elements. additionally, it incorporates llm-based validation to ensure logical adherence and semantic plausibility. experimental results on two public datasets, deepform and fara, demonstrate that pm3-kie outperforms state-of-the-art methods across in-distribution, outof-distribution, and low-resource scenarios, highlighting its effectiveness and robustness when facing challenges commonly encountered in practice.