近期以 DeepSeek OCR 为代表的 LLM 驱动端到端 OCR 方案凭借语言先验大幅提升识别精度,但存在致命短板:文本输出越长,KV 缓存持续膨胀,显存占用与推理速度同步恶化,处理多页 PDF、长篇报告时只能分页分段识别,上下文连贯性大打折扣。
百度全新开源 Unlimited-OCR 给出创新解法,自研 R-SWA 参考滑动窗口注意力机制,模拟人类阅读的工作记忆逻辑,全程维持恒定 KV 缓存,彻底解决长序列算力衰减问题。模型沿用高压缩视觉编码器,在 32K 上下文内单次推理即可完成数十页文档统一解析,不仅适配各类图文、表格、扫描件,R-SWA 架构还可迁移至 ASR、机器翻译等长序列任务,配套完整代码、权重开源,兼容 Transformers、vLLM、SGLang 主流推理框架,大幅降低长文档结构化解析落地门槛。
可视化展示
相关链接
-
论文:https://arxiv.org/pdf/2606.23050 -
仓库:https://github.com/baidu/Unlimited-OCR -
模型:https://huggingface.co/baidu/Unlimited-OCR
论文介绍
近期,以 DeepSeek OCR 为代表的端到端 OCR 模型再次引发了业界的广泛关注。一种普遍的观点认为,利用大语言模型(LLM)作为解码器,能够让模型借助语言的先验分布,从而提升 OCR 性能。然而,这种方法也存在明显的弊端:随着输出序列变长,累积的 KV 缓存会导致内存消耗激增,并使生成速度逐渐变慢。这与人类形成了鲜明对比,因为人类在执行长序列抄写任务时,并不会出现这种效率下降的现象。
论文提出了 Unlimited OCR 模型,旨在模拟人类解析过程中的工作记忆机制。该模型以 DeepSeek OCR 为基准,将解码器中的所有注意力层替换为我们提出的“参考滑动窗口注意力”(Reference Sliding Window Attention,简称 R-SWA);这种机制在降低注意力计算成本的同时,确保了整个解码过程中 KV 缓存大小保持恒定。结合 DeepSeek OCR 编码器的高压缩率与我们设计的恒定 KV 缓存机制,Unlimited OCR 能够在标准的 32K 最大序列长度限制下,通过单次前向传播完成数十页文档的转写任务。更重要的是,R-SWA 是一种通用的解析注意力机制——除了 OCR 之外,它同样适用于自动语音识别(ASR)、机器翻译等多种任务。相关代码与模型权重已在 http://github.com/baidu/Unlimited-OCR 公开。
方法概述
受人类抄书过程的启发,论文提出了“无限OCR”(Unlimited OCR)模型。该模型采用统一的端到端架构,由编码器和MoE-LLM解码器组成,其中所有注意力机制均采用R-SWA。其KV缓存被实现为一个容量为 的队列:每生成一个新的token,队列中对应于第 个 token 的 KV 项便会被移除,从而确保在生成过程中计算成本和内存占用均不会随时间推移而增加。
参考滑动窗口注意力(R-SWA)示意图。每个生成的 Token 都会关注所有参考 Token(OCR 任务中的视觉 Token)以及前
个输出 Token(默认值为 128)。与标准的全局注意力(full attention)相比,R-SWA 在整个解码过程中保持恒定的 KV 缓存大小。与常规滑动窗口注意力(vanilla SWA)相比,它将视觉 Token 排除在状态转换过程之外,从而保留了视觉 Token 的保真度,避免了渐进式模糊现象。
实验结果
OmniDocBench (v1.5/v1.6) 上的对比。表中的所有模型均为端到端 VLM 架构。v1.5 主要用于与经典的端到端算法及基线模型 DeepSeek OCR 进行对比;v1.6 则主要与当前的端到端 SOTA 模型进行对比。除本文提出的 Unlimited OCR 外,其余模型均选自 OmniDocBench 仓库。
针对九种文档类型,对 Unlimited OCR 与 DeepSeek-OCR 系列进行了详细的子类别对比。R-order 表示阅读顺序(Reading Order)。所有指标均为编辑距离,数值越低越好。红色单元格表示 DeepSeek-OCR 或 DeepSeek-OCR 2 在相应指标上的表现优于 Unlimited OCR。
长篇幅 OCR 的性能表现。在不同页数条件下测试了 Distinct-n 和编辑距离指标。对于 Distinct-n 指标,数值越高越好。
理论推理性能上限对比。比较了 DeepSeek OCR 和 Unlimited OCR 在不同输出长度下的 TPS 上限。
使用教程
在NVIDIA GPU上使用Huggingface Transformer进行推理。测试环境:Python 3.12.3 + CUDA 12.9。
torch==2.10.0
torchvision==0.25.0
transformers==4.57.1
Pillow==12.1.1
matplotlib==3.10.8
einops==0.8.2
addict==2.4.0
easydict==1.13
pymupdf==1.27.2.2
psutil==7.2.2
import os
import torch
from transformers import AutoModel, AutoTokenizer
model_name = 'baidu/Unlimited-OCR'
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(
model_name,
trust_remote_code=True,
use_safetensors=True,
torch_dtype=torch.bfloat16,
)
model = model.eval().cuda()
# ── Single image supports two configs: gundam or base ──
# gundam: base_size=1024, image_size=640, crop_mode=True
# base: base_size=1024, image_size=1024, crop_mode=False
model.infer(
tokenizer,
prompt='<image>document parsing.',
image_file='your_image.jpg',
output_path='your/output/dir',
base_size=1024, image_size=640, crop_mode=True,
max_length=32768,
no_repeat_ngram_size=35, ngram_window=128,
save_results=True,
)
# ── Multi page / PDF only uses base (image_size=1024) ──
model.infer_multi(
tokenizer,
prompt='<image>Multi page parsing.',
image_files=['page1.png', 'page2.png', 'page3.png'],
output_path='your/output/dir',
image_size=1024,
max_length=32768,
no_repeat_ngram_size=35, ngram_window=1024,
save_results=True,
)
# ── PDF (convert pages to images, then multi-page parsing) ──
import tempfile, fitz # PyMuPDF
def pdf_to_images(pdf_path, dpi=300):
doc = fitz.open(pdf_path)
tmp_dir = tempfile.mkdtemp(prefix='pdf_ocr_')
mat = fitz.Matrix(dpi / 72, dpi / 72)
paths = []
for i, page in enumerate(doc):
out = os.path.join(tmp_dir, f'page_{i+1:04d}.png')
page.get_pixmap(matrix=mat).save(out)
paths.append(out)
doc.close()
return paths
model.infer_multi(
tokenizer,
prompt='<image>Multi page parsing.',
image_files=pdf_to_images('your_doc.pdf', dpi=300),
output_path='your/output/dir',
image_size=1024,
max_length=32768,
no_repeat_ngram_size=35, ngram_window=1024,
save_results=True,
)
结论
在本技术报告中,我们提出了 Unlimited OCR 模型,并介绍了 R-SWA 算法,以支持其进行长序列解析的能力。我们验证了,当端到端模型解码器中的所有标准注意力机制被替换为基于因果参考的 SWA(R-SWA)时,模型在解析任务上的性能不会受损。这表明模型学会了将历史输出中的有用信息持续传递到当前窗口中;这种“软性遗忘”机制与人类在抄写书籍时的行为模式相吻合。我们相信,R-SWA 未来将应用于更多任务,从而使注意力计算和内存占用不再成为长序列解析领域的瓶颈。
感谢你看到这里,添加小助手 AIGC_Tech 加入官方 AIGC读者交流群,下方扫码加入 AIGC Studio 星球,获取前沿AI应用、AIGC实践教程、大厂面试经验、AI学习路线以及IT类入门到精通学习资料等,欢迎一起交流学习💗~

