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 は、2 つの主要コンポーネントからなるハイブリッドアーキテクチャを採用しています:

  • Vision Encoder: NaViT アーキテクチャに基づく 1.2B パラメータのエンコーダで、スクラッチから学習されています。このコンポーネントは Core ML で実行されます。
  • LM Backbone: Qwen2.5-1.5B バックボーンで、MLX フレームワークを通じて実行されます。

変換プロセス: PyTorch から Core ML へ

PyTorch から Core ML へのモデル変換は、主に 2 つのステップで構成されます。PyTorch 実行グラフを取得する（torch.jit.trace または torch.export を使用）ことと、そのグラフを coremltools を用いて .mlpackage にコンパイルすることです。

デバイス上での使用のためのモデル簡素化

変換を成功させるために、単一画像処理に不要な機能を除去してモデルを簡素化しました:

  • Single Image Processing: モデルは動画やバッチではなく、1 回に 1 枚の画像を処理するように変更されました。これはデバイス上のアプリケーションで一般的な最適化です。
  • Attention Implementation: すべての注意機構のバリエーションを除去し、iOS 18 の Core ML がサポートする標準の scaled_dot_product_attention（sdpa）演算子に統一しました。
  • Removal of Sliding Window Attention: モデルは Sliding Window Attention を必要としないため、sdpa との実装衝突を避けるために無効化しました。

変換中の技術的バグ修正

ビジョンエンコーダの変換中にいくつかの技術的課題が発生し、解決されました:

  • Dtype Mismatches: coremltools は torch.arange の dtype 引数を無視し、デフォルトで int32 になります。これを解決するために、行列乗算時に fp32 テンソルと互換性を保つよう明示的なキャストを追加しました。
  • Repeat Interleave Issues: flash_attention_2 で可変長シーケンスをマスクするために使用されていた repeat_interleave 呼び出しは、単一画像処理では不要なため削除しました。
  • Masking Logic: Neural Engine をサポートするため、bool テンソルがサポートされていないので、ブールマスクをすべてゼロのフロートマスクに置き換えました。
  • Dynamic Control Flow: grid_thw を反復するループを削除し、動的制御フローを排除しました。これは一般的に ML コンパイラでサポートされていません。

初期ベンチマークとパフォーマンス

初期の変換結果では、モデルは元の PyTorch の精度と一致し、最大差は 0.006000518798828125、平均差は 1.100682402466191e-05 でした。しかし、初期の FLOAT32 バージョンはサイズが 5GB を超え、GPU 上でビジョンエンコーダの単一フォワードパスに 1 秒以上かかるため、実際のデバイス展開向けにさらなる量子化と最適化が必要であることが示されました。

Sources