Trillion Labs, 자기 추측 디코딩 적용한 OCR 모델 GravityOCR 공개

Trillion Labs가 문서 이미지를 텍스트, HTML 표, LaTeX 수식으로 변환하는 OCR 모델 GravityOCR을 공개했다. 이 모델은 GLM-OCR(0.9B)의 가중치를 재사용하며, 자기 추측 디코딩(Self-Speculative Decoding) 기법을 적용해 디코딩 속도를 크게 개선했다.

핵심 기술 및 성능

GravityOCR은 별도의 초안 모델이나 보조 헤드 없이, GLM-OCR에 마스크 토큰 [M]의 학습 가능한 임베딩 하나만 추가하여 확산(Diffusion) 경로로 초안을 작성하고 자기회귀(AR) 경로로 검증하는 구조를 채택했다. 이를 통해 모델 forward 1회당 평균 9.7개의 토큰을 확정할 수 있다.

• 속도 향상: H100 GPU 기준 초당 0.730페이지를 처리하여, 동일 체크포인트의 AR 경로(0.554페이지) 대비 1.32배 빠르다.
• 정확도 유지: OmniDocBench v1.6 벤치마크에서 Overall 점수 95.16점을 기록하며, 원본 GLM-OCR(95.48)의 99.7% 수준의 정확도를 유지했다.
• 출력 동일성: 정확한 산술 연산에서는 AR 그리디 디코딩과 동일한 출력을 생성하며, bf16 서빙 환경에서도 96.6%의 완전 일치율을 보였다.

학습 및 구현 특징

GLM-OCR 체크포인트에서 미세조정 방식으로 학습되었으며, Clean Stream(AR 손실)과 Corrupted Streams(확산 손실)을 함께 학습한다. 학습 데이터는 약 1,230만 개의 영역 단위 예시(주로 영어)를 사용했다. 추론 시 라운드당 2번의 forward(초안 생성 및 검증)를 통해 최소 2개에서 최대 B+2개의 토큰을 확정한다.

라이선스 및 공개 정보

모델 가중치는 MIT 라이선스로 Hugging Face에, 서빙 및 추론 코드는 GitHub에 공개되었다. SGLang v0.5.12에 패치를 적용하여 서빙할 수 있으며, --spec_natcache 옵션을 통해 AR 경로와 바이트 단위 동일한 출력을 보장한다. 다만, 학습 코드와 데이터 파이프라인은 비공개 상태이다.