用状态空间模型替代Transformer，高效处理长文档
DocMamba: Efficient Document Pre-training with State Space Model
- 采用状态空间模型实现线性计算复杂度，替代传统自注意力
- 在FUNSD、CORD等数据集上达新最优，速度更快、内存更少
- 适合需要处理长文档或追求推理效率的研究与应用
近年来，视觉丰富的文档理解受到广泛关注。基于Transformer的预训练模型虽性能优越，但其自注意力机制的二次计算复杂度限制了处理长文档的效率。本文提出DocMamba，一种基于状态空间模型的新框架，将计算复杂度降至线性，同时保持全局建模能力。为提升文档处理效果，引入分段优先双向扫描（SFBS）以捕捉连续语义信息。实验表明，DocMamba在FUNSD、CORD和SORIE等下游数据集上达到新最优表现，显著提升处理速度并降低内存占用。特别地，在HRDoc上的实验验证了其在长度外推方面的潜力。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。