DocMamba: Efficient Document Pre-training with State Space Model

📄 arXiv: 2409.11887v2 📥 PDF

作者: Pengfei Hu, Zhenrong Zhang, Jiefeng Ma, Shuhang Liu, Jun Du, Jianshu Zhang

分类: cs.CL, cs.AI

发布日期: 2024-09-18 (更新: 2025-02-10)

💡 一句话要点

DocMamba：基于状态空间模型的高效文档预训练框架

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 文档理解 状态空间模型 预训练模型 长文档处理 视觉丰富文档 Mamba 线性复杂度

📋 核心要点

  1. Transformer在文档理解中表现出色，但自注意力机制的计算复杂度限制了其效率和长文档处理能力。
  2. DocMamba基于状态空间模型，将计算复杂度降至线性，并保留全局建模能力，提升长文档处理效率。
  3. DocMamba在FUNSD、CORD和SORIE等数据集上取得了SOTA结果，显著提升速度并降低内存占用，并在HRDoc上验证了长度外推能力。

📝 摘要（中文）

近年来，视觉丰富的文档理解受到了越来越多的关注。基于Transformer的预训练模型已成为主流方法，并在该领域取得了显著的性能提升。然而，自注意力机制的二次计算复杂度限制了它们的效率和处理长文档的能力。本文提出了DocMamba，一种基于状态空间模型的新型框架。它旨在将计算复杂度降低到线性，同时保留全局建模能力。为了进一步提高其在文档处理中的有效性，我们引入了Segment-First Bidirectional Scan (SFBS)来捕获连续的语义信息。实验结果表明，DocMamba在下游数据集（如FUNSD、CORD和SORIE）上取得了新的state-of-the-art结果，同时显著提高了速度并减少了内存使用。值得注意的是，在HRDoc上的实验证实了DocMamba在长度外推方面的潜力。

🔬 方法详解

问题定义：现有基于Transformer的文档理解模型，特别是处理长文档时，面临自注意力机制带来的二次方计算复杂度问题，导致效率低下和内存消耗过大。这限制了模型在实际应用中的部署和扩展能力。

核心思路：DocMamba的核心思路是利用状态空间模型（SSM）替代Transformer中的自注意力机制。SSM具有线性计算复杂度，能够更高效地处理长序列数据，同时保持全局建模能力。通过这种方式，DocMamba旨在克服Transformer在处理长文档时的效率瓶颈。

技术框架：DocMamba的整体框架包括一个基于状态空间模型的编码器。该编码器接收视觉丰富的文档作为输入，并通过一系列SSM层进行处理。为了更好地捕捉文档中的语义信息，DocMamba引入了Segment-First Bidirectional Scan (SFBS)策略。SFBS首先将文档分割成段落，然后进行双向扫描，从而捕捉连续的语义信息。最终，模型输出文档的表示，用于下游任务。

关键创新：DocMamba的关键创新在于将状态空间模型应用于文档预训练，并提出了Segment-First Bidirectional Scan (SFBS)策略。与Transformer相比，SSM具有线性计算复杂度，更适合处理长文档。SFBS策略则能够更好地捕捉文档中的连续语义信息，提升模型性能。

关键设计：DocMamba的关键设计包括：1) 使用Mamba作为核心的状态空间模型，Mamba是一种选择性状态空间模型，能够根据输入动态调整状态转移矩阵；2) 引入Segment-First Bidirectional Scan (SFBS)策略，具体实现方式未知；3) 损失函数和优化器等细节未知。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

DocMamba在FUNSD、CORD和SORIE等下游数据集上取得了新的state-of-the-art结果，表明其在文档理解任务上的优越性能。实验结果还显示，DocMamba显著提高了处理速度并减少了内存使用，验证了其高效性。此外，在HRDoc数据集上的实验证实了DocMamba在长度外推方面的潜力。

🎯 应用场景

DocMamba在视觉丰富的文档理解领域具有广泛的应用前景，例如自动化文档分析、信息抽取、智能文档处理等。该模型可以应用于金融、法律、医疗等行业，提高文档处理效率，降低人工成本。未来，DocMamba有望成为文档智能化的重要基石。

📄 摘要（原文）

In recent years, visually-rich document understanding has attracted increasing attention. Transformer-based pre-trained models have become the mainstream approach, yielding significant performance gains in this field. However, the self-attention mechanism's quadratic computational complexity hinders their efficiency and ability to process long documents. In this paper, we present DocMamba, a novel framework based on the state space model. It is designed to reduce computational complexity to linear while preserving global modeling capabilities. To further enhance its effectiveness in document processing, we introduce the Segment-First Bidirectional Scan (SFBS) to capture contiguous semantic information. Experimental results demonstrate that DocMamba achieves new state-of-the-art results on downstream datasets such as FUNSD, CORD, and SORIE, while significantly improving speed and reducing memory usage. Notably, experiments on the HRDoc confirm DocMamba's potential for length extrapolation.