Baidu Unlimited-OCR로 장문서 OCR의 KV 캐시 문제 해결하기

DeepSeek-OCR처럼 LLM 디코더로 텍스트를 뽑아내는 OCR 모델은 페이지 수가 늘어날수록 골치가 아파진다. 인식 결과를 토큰 단위로 하나씩 생성하는 구조라서, 출력이 길어질수록 어텐션이 참조해야 할 KV 캐시가 계속 쌓이고 메모리 사용량과 생성 속도가 함께 나빠진다. 스캔본 몇 장이면 괜찮지만 수십 페이지짜리 PDF를 한 번에 넣으면 금세 한계에 부딪힌다.

바이두가 최근 공개한 Unlimited-OCR(baidu/Unlimited-OCR)은 이 문제를 정면으로 다룬 모델이다. DeepSeek-OCR을 기반으로 R-SWA(Reference Sliding Window Attention)라는 새 어텐션 메커니즘을 얹어서, 디코딩이 아무리 길어져도 KV 캐시 크기를 일정하게 유지한다. 그 결과 32K 토큰 컨텍스트 안에서 수십 페이지 문서를 단일 forward pass로 처리할 수 있다고 밝히고 있다.

KV 캐시가 계속 쌓이는 이유

트랜스포머 디코더는 새 토큰을 생성할 때마다 이전에 생성한 모든 토큰의 key/value를 어텐션 계산에 다시 사용한다. 이게 KV 캐시인데, 시퀀스 길이에 비례해서 커진다. OCR처럼 출력이 원문 그대로 길게 이어지는 작업에서는 이 캐시가 특히 빠르게 불어난다. 캐시가 커지면 GPU 메모리를 잡아먹는 것은 물론, 매 스텝마다 어텐션이 봐야 할 대상이 늘어나므로 생성 속도 자체도 느려진다.

R-SWA: 작업 기억을 모방한 고정 크기 캐시

R-SWA는 사람이 긴 글을 읽을 때 방금 읽은 문장 몇 개만 기억하고 그 이전 내용은 압축된 형태로만 남겨두는 것과 비슷한 방식으로 동작한다. 슬라이딩 윈도우로 최근 구간만 상세히 참조하고, 나머지는 고정 크기 상태로 관리해서 디코딩 전 구간에 걸쳐 KV 캐시 크기가 늘어나지 않게 만든다. 여기에 인코더의 높은 이미지 압축률까지 더해져, 논문에서는 32K 토큰 한도 안에서 수십 페이지짜리 문서를 하나의 forward pass로 처리할 수 있다고 설명한다. 저자들은 이 메커니즘이 OCR뿐 아니라 음성 인식이나 번역처럼 출력이 긴 다른 시퀀스 생성 작업에도 적용 가능한 범용 구조라고 밝히고 있다.

스펙과 두 가지 추론 모드

모델은 3B 파라미터, BF16 safetensors로 배포되며 라이선스는 MIT다. 이미지 처리 방식은 두 가지 프리셋으로 나뉜다.

• gundam: base_size=1024, image_size=640, crop_mode=True — 문서를 크롭해서 세부 텍스트 인식률을 높이는 모드
• base: base_size=1024, image_size=1024, crop_mode=False — 전체 페이지를 한 번에 보는 모드

추론은 Hugging Face Transformers 외에 vLLM(vllm/vllm-openai:unlimited-ocr 도커 이미지), SGLang을 통해서도 가능하다.

멀티페이지 PDF 파싱해보기

Transformers로 단일 이미지를 파싱하는 코드는 다음과 같다.

import torch
from transformers import AutoModel, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained('baidu/Unlimited-OCR', trust_remote_code=True)
model = AutoModel.from_pretrained(
    'baidu/Unlimited-OCR',
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
).eval().cuda()

result = model.infer(
    tokenizer,
    prompt='<image>document parsing.',
    image_file='page.jpg',
    base_size=1024,
    image_size=640,
    crop_mode=True,
    max_length=32768,
)

여러 페이지를 한 번에 넣고 싶다면 infer_multi()를 쓰면 된다. PDF는 PyMuPDF로 페이지별 이미지로 변환한 뒤 리스트로 넘기는 방식이라, 별도의 페이지 분할 파이프라인을 직접 짤 필요가 없다.

벤치마크와 현재 한계

llamaindex/ParseBench 기준 공개된 수치는 평균 46.17점, Text Content 86.81점, Text Formatting 0.97점이다. 순수 텍스트 인식은 준수하지만 표나 레이아웃 같은 포맷 보존 점수는 아직 낮은 편이라, 문서 구조까지 그대로 살려야 하는 워크플로우에는 후처리가 필요할 수 있다. 그래도 긴 문서를 페이지 단위로 쪼개 반복 호출하던 기존 파이프라인을 단일 호출로 대체할 수 있다는 점에서, 대량 문서를 다루는 OCR 파이프라인이라면 한 번쯤 벤치마크해볼 가치가 있다.