SOTA OCR with Core ML and dots.ocr

Hugging Face has detailed the process of converting dots.ocr, a 3B parameter OCR model from RedNote, for on-device execution using Apple's Core ML and MLX frameworks. The model is significant because it surpasses Gemini 2.5 Pro on the OmniDocBench benchmark, enabling high-performance OCR without the need for API keys, network connectivity, or recurring costs.

Hardware Acceleration via the Neural Engine

Apple's Neural Engine (NE) provides a highly power-efficient alternative to the CPU and GPU for AI workloads. Testing indicates that the Neural Engine is 12x more power efficient than the CPU and 4x more power efficient than the GPU, making it the ideal target for on-device models with limited power budgets. However, the Neural Engine is only accessible through Core ML, a closed-source framework, which often requires complex conversion processes from PyTorch.

Model Architecture: dots.ocr

dots.ocr utilizes a hybrid architecture consisting of two primary components:

  • Vision Encoder: A 1.2B parameter encoder based on the NaViT architecture, trained from scratch. This component is run via Core ML.
  • LM Backbone: A Qwen2.5-1.5B backbone, which is run via the MLX framework.

The Conversion Process: PyTorch to Core ML

Converting a model from PyTorch to Core ML involves two primary steps: capturing the PyTorch execution graph (using torch.jit.trace or torch.export) and compiling that graph into an .mlpackage using coremltools.

Simplifying the Model for On-Device Use

To ensure a successful conversion, the model was simplified by removing features not essential for single-image processing:

  • Single Image Processing: The model was modified to process one image at a time rather than videos or batches, a common optimization for on-device applications.
  • Attention Implementation: All attention variants were removed in favor of the standard scaled_dot_product_attention (sdpa) operator, which is supported by Core ML in iOS 18.
  • Removal of Sliding Window Attention: Since the model does not require Sliding Window Attention to function, it was disabled to avoid implementation conflicts with sdpa.

Technical Bug Fixes during Conversion

Several technical hurdles were encountered and resolved during the conversion of the vision encoder:

  • Dtype Mismatches: coremltools ignores the dtype argument in torch.arange, defaulting to int32. This was fixed by adding an explicit cast to ensure compatibility with fp32 tensors during matrix multiplication.
  • Repeat Interleave Issues: The repeat_interleave call used for masking variable length sequences in flash_attention_2 was removed, as it is unnecessary when processing a single image.
  • Masking Logic: To support the Neural Engine, boolean masks were replaced with float masks of all zeros, as the Neural Engine does not support bool tensors.
  • Dynamic Control Flow: A loop iterating over grid_thw was removed to eliminate dynamic control flow, which is generally unsupported by ML compilers.

Initial Benchmarking and Performance

Initial conversion results showed that the model matched original PyTorch precision with a max difference of 0.006000518798828125 and a mean difference of 1.100682402466191e-05. However, the initial FLOAT32 version of the model exceeded 5GB in size and required over one second for a single forward pass of the vision encoder on the GPU, indicating a need for further quantization and optimization for production on-device deployment.

Sources