KV缓存压成常数:超长文档OCR的开源新解法

作者:c4t2026.11.04 06:06浏览量:0

简介:超长文档解析长期受制于KV缓存的线性膨胀与注意力的平方级开销。百度新一代开源OCR模型通过稀疏参考滑动窗口注意力(SR-SWA)与持续解码机制,在百页级单序列测试中将峰值KV缓存从47GB降至1.77GB、降幅约96%,实现数千页文档连续识别转写。本文拆解其设计思路、复杂度变化与工程落地价值。

一、从“单页识别”到“整本阅读”:一道绕不开的显存关卡

文档数字化的需求正在发生变化。过去,OCR的任务边界通常是一张图片、一页扫描件;而在真实业务里,更常见的输入是整本书籍、成套卷宗、连续的年报财报。用户期望的不再是几千份孤立的单页结果,而是一份上下文连贯、可直接检索的完整文本。

百度新一代开源OCR模型正是冲着这个场景来的。据其技术报告,它将某七部曲长篇小说全集——3600余页——作为单条序列连续识别并转写,从头读到尾不中断;换成一套三卷本、1200余页、几十万字的长篇小说,同样一口气完成。

但“连着读”的前提,是先过一道关:显存。Transformer模型在推理时要为已处理内容保留中间状态,也就是KV缓存——可以理解为模型边读边记的“小抄”。页数越多,小抄越厚:不仅占显存,每生成一个字还要翻遍整本小抄,越翻越慢。

技术报告中给出了一组百页级单序列测试的对比数据:若采用标准解码方式,把所有页面的视觉缓存同时驻留显存,峰值KV缓存占用为47GB;而该模型将其降至1.77GB,削减约96%。47GB是什么概念?已经超过不少笔记本电脑的整机内存,而这还只是缓存开销,模型权重等其他成本另算。

值得一提的是,该系列模型的上一代已凭借多页连续解析能力在开源社区积累了相当人气:累计下载量超过1000万,收获约2.86万Star,社区自制在线Demo、推理框架适配等工作持续推进。到了第二代,升级重点非常明确:在连续处理更长文档的同时,把计算与缓存的负担压下来。

二、问题的根源:KV缓存为什么随页数线性膨胀

要理解96%这个数字的含金量,得先弄清楚KV缓存的账是怎么算的。

Transformer的自注意力机制中,每个token都会产生三份向量:Query(查询)、Key(键)、Value(值)。自回归解码时,模型每生成一个新token,都需要与之前所有token的K、V做注意力计算。为了避免每一步都重算历史,工程上会把历史的K、V缓存下来——这就是KV缓存。

对纯文本模型,一个token对应一份KV;但对OCR这类视觉语言模型,问题被急剧放大:一页文档图像会先被视觉编码器切成大量图像块,映射成成百上千个视觉token。几百页文档累积下来,序列长度轻松达到数十万级别。

KV缓存占用的显存,业界通用估算公式为:

KV缓存 ≈ 2(K和V)× 层数 × 序列长度 × KV头数 × 头维度 × 每元素字节数

做一个纯示意性的量级估算(以下参数为假设值,不代表该模型实际配置):假设某模型有32层、8个KV头、头维度128、使用FP16存储,则每个token约需128KB缓存;10万个token就是约12.8GB。几百页、数十万token的序列,缓存膨胀到几十GB并不意外——这与报告中47GB的量级是吻合的。更麻烦的是,缓存增量与页数成正比:页数翻倍,缓存翻倍,直到显存见底。

计算侧同样不乐观。全量注意力的计算复杂度约为O(n²),序列越长,每一步解码要“回看”的内容越多,延迟与算力开销同步恶化。换言之,长文档OCR的瓶颈从来不是“认不出字”,而是“记不住”与“算不动”。

  1. # 示意:标准解码下KV缓存随页数线性增长
  2. kv_cache = []
  3. for page in document:
  4. visual_tokens = encode(page) # 每页成百上千个视觉token
  5. kv_cache += project_kv(visual_tokens)
  6. text = decode(kv_cache) # 每步都要回看全部历史
  7. # 显存:O(总token数),与页数成正比
  8. # 计算:注意力O(n²),随页数平方增长

三、SR-SWA:像人抄书一样“按需回看”

连续读完3600余页,靠的是什么?技术报告给出的答案是稀疏参考滑动窗口注意力(SR-SWA)。它的思路,与人类抄书高度相似。

想象一个人手工抄一整本书:他不会把全书同时摊在桌上,而是逐页精读当前页;只有当需要核对前文的人名、术语或表格时,才回头翻找特定几页。整个过程,桌面上始终只摊开一两页。

业界早已成熟的一类做法是滑动窗口注意力(SWA):把注意力限制在最近的W个token内,缓存只需保留窗口内容,显存与单步计算都可控。但它有个明显短板——窗口外的历史信息被“遗忘”,跨页依赖(比如第3页定义的缩写、第50页延续的同一张跨页表格)会丢失。

SR-SWA在滑动窗口之上增加了一条“稀疏参考”通道,可以拆成三个角色来理解:

  • 局部窗口:对当前页及邻近内容保持精细、完整的注意力,保证当下这一页的识别精度不打折;
  • 稀疏参考:将已读页面压缩为参考信息,解码时按需检索与当前内容相关的历史页面和视觉信息,补上跨页的长程依赖;
  • 减少重复计算:既然不必对全部历史做全量注意力,历史内容也无需反复重算,重复开销被裁掉。

三者叠加的效果是:缓存占用不再与总页数成正比,而是取决于窗口大小与参考信息的规模——近似常数。这正是“页数没上限”的技术来源:显存曲线被拉平了,页数自然不再是硬约束。

  1. # 示意:滑动窗口 + 稀疏参考
  2. window = deque(maxlen=W) # 只保留最近W个token的KV
  3. reference = SparseReference() # 已读页的压缩参考索引
  4. for page in document:
  5. tokens = encode(page)
  6. for t in tokens:
  7. ctx = concat(window, reference.lookup(t)) # 局部精细 + 按需回看
  8. out = decode_step(t, ctx)
  9. window.append(kv_of(t))
  10. reference.absorb(page) # 该页读毕,压缩进参考,不保留全量KV
  11. # 显存:O(W + 参考量),近似常数

四、持续解码:把单步开销与文档长度解耦

SR-SWA解决“怎么记”,持续解码机制则解决“怎么读”。

传统做法处理长文档,往往要人工切片:按页或按章切成多段,分别送入模型,最后再做结果拼接。代价是上下文断裂——跨页表格被拦腰截断、人物称谓前后不一、章节编号对不上,后续还需要额外的校对与合并流水线。

持续解码的思路是把整份文档视为一条连续序列,按页或按块逐步“读取—压缩—输出”:

  1. # 示意:持续解码主循环
  2. state = init_state()
  3. for chunk in split(document, by="page"):
  4. state = update_reference(state, chunk) # 已读内容压缩进参考
  5. yield decode(state, chunk) # 流式输出当前页文本

由此带来三个直接收益:

  1. 单步开销恒定:每一步解码只依赖窗口与参考,开销不随文档变长而水涨船高,延迟曲线平稳可预期;
  2. 上下文真正连续:跨页表格、术语一致性、编号体系在同一条序列内自然延续,不需要事后再“缝合”;
  3. 流式交付:边读边出文本,几千页的任务可以边处理边消费,不必等全部完成。

五、从47GB到1.77GB:96%削减的工程价值

回到那组测试数据:百页级单序列,峰值KV缓存从47GB降至1.77GB,降幅约96%。这组数字在工程上至少意味着四件事:

  • 显存门槛大幅下降。47GB级别的缓存意味着必须依赖多卡或高端显存卡;缓存降到1.77GB后,仅缓存部分的开销已降到常规硬件可以轻松承载的水平,长文档能力的部署弹性显著提高。
  • 并发与吞吐提升。省出来的显存可以容纳更多并发请求或更大批次,单位算力能摊薄的文档处理成本随之下降。
  • 长任务稳定性更好。峰值显存与页数解耦后,任务中途显存溢出的风险大幅降低,几百页与几千页的任务在资源画像上趋于一致,容量规划更简单。
  • 业务形态改变。“整本书一次交付”成为可能,人工切分、多任务拼接、二次校对的流水线环节被压缩,端到端延迟与人力成本同步下降。

六、落地视角:适用场景与工程建议

对计划引入超长文档解析能力的团队,以下几点值得关注:

典型适用场景:整本书籍与教材数字化、成套卷宗与合同解析、连续年报财报的结构化提取、连载内容的批量转写——凡是“输入天然连续、上下文跨页”的场景,收益最明显。

需要评估的边界:

  1. 极端长程依赖(如跨越数百页的引用)的还原质量,取决于参考压缩的保真度,建议对关键业务做抽样校验;
  2. 表格合并、页码对齐等结构化输出,应在解析后增加一致性校验环节;
  3. 窗口大小与分块粒度本质上是精度与成本的权衡旋钮,不同文档类型(印刷体书籍与密集表格)的最优配置可能不同;
  4. 长任务应配套基础可观测能力:缓存占用、吞吐、延迟等指标的监控与告警,是保障批量任务稳定运行的常规要求。

架构层面的配合:超长文档解析通常还要与通用云上能力组合——对象存储存放原始扫描件、消息队列做任务分片调度、容器平台按负载弹性扩缩、日志服务沉淀处理轨迹。这些能力与模型本身的缓存优化相互正交,组合起来才能撑起生产级的批量处理。

七、结语

从“单页识别”到“整本阅读”,OCR演进的核心矛盾始终是记忆成本。SR-SWA与持续解码给出的答案是:不追求“记得多”,而是“记得巧”——局部窗口保精度,稀疏参考补长程,流式压缩控成本,最终用近似常数的缓存换来事实上无上限的页数。

对开发者而言,这份开源实现的价值不止于OCR本身。“局部精细 + 稀疏回看 + 流式压缩”是一套可迁移的长上下文设计范式,任何需要处理超长序列的系统——代码库理解、对话历史压缩、日志分析——都能从中获得借鉴。该模型已在开源社区正式发布,感兴趣的开发者可以直接取用验证。