Core ML 및 dots.ocr를 활용한 최신 OCR

Hugging Face는 RedNote의 3B 파라미터 OCR 모델인 dots.ocr를 Apple의 Core ML 및 MLX 프레임워크를 사용해 온디바이스 실행으로 변환하는 과정을 상세히 설명했습니다. 이 모델은 OmniDocBench 벤치마크에서 Gemini 2.5 Pro를 능가하여 API 키, 네트워크 연결 또는 지속적인 비용 없이 고성능 OCR을 가능하게 합니다.

Neural Engine을 통한 하드웨어 가속

Apple의 Neural Engine(NE)은 AI 작업에 대해 CPU와 GPU에 비해 매우 전력 효율적인 대안을 제공합니다. 테스트 결과 Neural Engine이 CPU보다 12배 더 전력 효율적이며 GPU보다 4배 더 전력 효율적임을 보여주어 전력 예산이 제한된 온디바이스 모델에 이상적인 대상이 됩니다. 그러나 Neural Engine은 폐쇄형 프레임워크인 Core ML을 통해서만 접근 가능하며, 이는 종종 PyTorch에서 복잡한 변환 과정을 필요로 합니다.

모델 아키텍처: dots.ocr

dots.ocr는 두 가지 주요 구성 요소로 이루어진 하이브리드 아키텍처를 활용합니다:

  • Vision Encoder: NaViT 아키텍처를 기반으로 한 1.2B 파라미터 인코더로, 처음부터 학습되었습니다. 이 구성 요소는 Core ML을 통해 실행됩니다.
  • LM Backbone: Qwen2.5-1.5B 백본으로, MLX 프레임워크를 통해 실행됩니다.

변환 과정: PyTorch에서 Core ML로

PyTorch 모델을 Core ML로 변환하려면 두 가지 주요 단계가 필요합니다: PyTorch 실행 그래프를 캡처(torch.jit.trace 또는 torch.export 사용)하고, 해당 그래프를 coremltools를 사용해 .mlpackage로 컴파일합니다.

온디바이스 사용을 위한 모델 단순화

성공적인 변환을 위해 모델을 단일 이미지 처리에 필수적이지 않은 기능들을 제거하여 단순화했습니다:

  • Single Image Processing: 모델을 비디오나 배치가 아닌 한 번에 하나의 이미지만 처리하도록 수정했습니다. 이는 온디바이스 애플리케이션에서 흔히 사용하는 최적화입니다.
  • Attention Implementation: 모든 attention 변형을 제거하고, iOS 18에서 Core ML이 지원하는 표준 scaled_dot_product_attention(sdpa) 연산자를 사용했습니다.
  • Removal of Sliding Window Attention: 모델이 Sliding Window Attention을 필요로 하지 않으므로, sdpa와의 구현 충돌을 피하기 위해 이를 비활성화했습니다.

변환 중 기술적 버그 수정

Vision Encoder 변환 과정에서 여러 기술적 장애물이 발생했으며, 이를 해결했습니다:

  • Dtype Mismatches: coremltools는 torch.arange의 dtype 인자를 무시하고 기본값을 int32로 설정합니다. 이를 해결하기 위해 행렬 곱셈 시 fp32 텐서와 호환되도록 명시적인 캐스트를 추가했습니다.
  • Repeat Interleave Issues: flash_attention_2에서 가변 길이 시퀀스를 마스킹하기 위해 사용된 repeat_interleave 호출을 제거했습니다. 이는 단일 이미지를 처리할 때 불필요합니다.
  • Masking Logic: Neural Engine을 지원하기 위해 불리언 마스크를 모두 0인 float 마스크로 교체했습니다. Neural Engine은 bool 텐서를 지원하지 않기 때문입니다.
  • Dynamic Control Flow: grid_thw를 반복하는 루프를 제거하여 동적 제어 흐름을 없앴습니다. 이는 일반적으로 ML 컴파일러에서 지원되지 않습니다.

초기 벤치마킹 및 성능

초기 변환 결과는 모델이 원본 PyTorch 정밀도와 최대 차이 0.006000518798828125, 평균 차이 1.100682402466191e-05로 일치함을 보여주었습니다. 그러나 초기 FLOAT32 버전 모델은 크기가 5GB를 초과하고 GPU에서 Vision Encoder의 단일 포워드 패스에 1초 이상이 소요되어, 실제 온디바이스 배포를 위한 추가 양자화 및 최적화가 필요함을 나타냈습니다.

Sources