Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Up to 36x Speedup: Mask-based Parallel Inference Paradigm for Key Information Extraction in MLLMs

Xinzhong Wang, Ya Guo, Jing Li, Huan Chen, Yi Tu, Yijie Hong, Gongshen Liu, Huijia Zhu

TL;DR

The paper tackles the latency bottleneck in key information extraction (KIE) from visually-rich documents by proposing PIP, a mask-based parallel inference paradigm that decouples output slots and generates all masked tokens in a single forward pass. It introduces a two-stage training regime—mask pre-training with bidirectional attention on a large image-caption corpus, followed by KV supervised fine-tuning on a curated KV dataset—to enable effective parallel decoding. Empirically, PIP achieves 5-36× inference speedups with negligible accuracy loss and sets SOTA performance on SROIE and CORD, while remaining competitive on FUNSD. The approach is demonstrated across five public benchmarks, with attention analyses showing tokens attend to distinct visual regions, supporting scalable real-time KIE for industry deployment.

Abstract

Key Information Extraction (KIE) from visually-rich documents (VrDs) is a critical task, for which recent Large Language Models (LLMs) and Multi-Modal Large Language Models (MLLMs) have demonstrated strong potential. However, their reliance on autoregressive inference, which generates outputs sequentially, creates a significant efficiency bottleneck, especially as KIE tasks often involve extracting multiple, semantically independent fields. To overcome this limitation, we introduce PIP: a Parallel Inference Paradigm for KIE. Our approach reformulates the problem by using "[mask]" tokens as placeholders for all target values, enabling their simultaneous generation in a single forward pass. To facilitate this paradigm, we develop a tailored mask pre-training strategy and construct large-scale supervised datasets. Experimental results show that our PIP-models achieve a 5-36x inference speedup with negligible performance degradation compared to traditional autoregressive base models. By substantially improving efficiency while maintaining high accuracy, PIP paves the way for scalable and practical real-world KIE solutions.

Up to 36x Speedup: Mask-based Parallel Inference Paradigm for Key Information Extraction in MLLMs

TL;DR

The paper tackles the latency bottleneck in key information extraction (KIE) from visually-rich documents by proposing PIP, a mask-based parallel inference paradigm that decouples output slots and generates all masked tokens in a single forward pass. It introduces a two-stage training regime—mask pre-training with bidirectional attention on a large image-caption corpus, followed by KV supervised fine-tuning on a curated KV dataset—to enable effective parallel decoding. Empirically, PIP achieves 5-36× inference speedups with negligible accuracy loss and sets SOTA performance on SROIE and CORD, while remaining competitive on FUNSD. The approach is demonstrated across five public benchmarks, with attention analyses showing tokens attend to distinct visual regions, supporting scalable real-time KIE for industry deployment.

Abstract

Key Information Extraction (KIE) from visually-rich documents (VrDs) is a critical task, for which recent Large Language Models (LLMs) and Multi-Modal Large Language Models (MLLMs) have demonstrated strong potential. However, their reliance on autoregressive inference, which generates outputs sequentially, creates a significant efficiency bottleneck, especially as KIE tasks often involve extracting multiple, semantically independent fields. To overcome this limitation, we introduce PIP: a Parallel Inference Paradigm for KIE. Our approach reformulates the problem by using "[mask]" tokens as placeholders for all target values, enabling their simultaneous generation in a single forward pass. To facilitate this paradigm, we develop a tailored mask pre-training strategy and construct large-scale supervised datasets. Experimental results show that our PIP-models achieve a 5-36x inference speedup with negligible performance degradation compared to traditional autoregressive base models. By substantially improving efficiency while maintaining high accuracy, PIP paves the way for scalable and practical real-world KIE solutions.
Paper Structure (15 sections, 4 equations, 3 figures, 3 tables)

This paper contains 15 sections, 4 equations, 3 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Methodology
  3. Task Formulation
  4. Model Training Process
  5. Mask Pre-training
  6. KV Supervised Fine-Tuning
  7. Experiment
  8. Experiments Setup
  9. Results
  10. Compared with SOTA Models
  11. Compared with Base Models
  12. GPU Memory Consumption
  13. Attention Visualization
  14. Conclusion
  15. Acknowledgements

Figures (3)

  • Figure 1: Comparison of two inference paradigms: (a) Traditional autoregressive inference, which generates tokens sequentially one by one; (b) Our PIP-Models, where different "[mask]" tokens independently attend to distinct image regions and generate all tokens in parallel.
  • Figure 2: Our overall training processes: (a) represents the mask pre-training phase, and (b) denotes the KV supervised fine-tuning stage.
  • Figure 3: The visualization of attention for each token in our PIP-Models when outputting "193.00", demonstrating the model's focus on different regions of the image.