AppleのNeural Engineで動作するSOTA OCRモデル「dots.ocr」をCore MLとMLXで実装——3Bパラメータの高精度オンデバイスOCR実現への道
2025年、Appleのハードウェア性能とモデルの精度が進化し、オンデバイス上で競争力のあるAIモデルを実行することが現実味を帯びてきた。RedNoteが開発した30億パラメータのOCRモデル「dots.ocr」は、Gemini 2.5 Proを上回る性能を実現し、オンデバイスでのOCR利用が本格化する可能性を示している。オンデバイス実行の利点は、APIキーの漏洩リスクゼロ、コストゼロ、ネットワーク不要といった点にあるが、限られた計算資源と電力予算を考慮する必要がある。 この課題を解決する鍵は、Appleの独自AIアクセラレータ「Neural Engine」である。2017年以降、すべてのAppleデバイスに搭載されており、CPUやGPUに比べて最大12倍の電力効率を発揮する。しかし、Neural EngineはCore MLという閉鎖的なフレームワークを通じてのみアクセス可能で、PyTorchモデルをCore MLに変換する際には技術的ハードルが存在する。また、MLXというGPU向けの柔軟なフレームワークも併用することで、開発の柔軟性が向上する。 本記事では、dots.ocrのオンデバイス化プロセスを3段階に分けて解説する。まず、モデル構造の理解と簡略化。dots.ocrは12億パラメータの視覚エンコーダ(NaViTベース)と15億パラメータの言語モデル(Qwen2.5)から構成され、視覚エンコーダはCore ML、言語モデルはMLXで処理する。変換の第一歩として、複数画像や動画処理を単一画像処理に簡略化。さらに、複数のアテンション実装をシンプルなscaled_dot_product_attention(sdpa)に統一し、Core MLとの互換性を確保。 次に、変換プロセスで発生するバグを一つずつ修正。最初はarangeのデータ型がint32に強制される問題があり、castを追加して解決。次にrepeat_interleaveで定数の不一致が発生し、変数長シーケンス処理のコードを削除。さらに、inplace_fillの動的インデックス制限や、reshapeの形状不一致も、ループの削除とマスクの型変換で対処。最終的に、PyTorchとCore MLでの出力差は最大0.006以内と非常に小さい結果を得た。 しかし、モデルサイズは5GBを超え、実用には不適。1回の推論で1秒以上かかることも判明。この問題は、第2段階でCore MLとMLXの統合、第3段階で量子化や動的形状対応といった最適化により解決を目指す。このプロセスは、他のモデルのオンデバイス化にも応用可能であり、開発者への実用的な指針となる。
