apple-aiml-research/ml-ane-transformers

Apple Neural Engine (ANE) Transformers

何がこれか – ane_transformers は、AppleのNeural Engine (ANE) で効率的に実行できるように最適化されたTransformerモデルの参照実装を提供する小さなオープンソースライブラリです。以下の2つのパッケージを提供しています：

• ane_transformers.reference – ANE上で効率的に実行するために必要な低レベルの変更を示す、独立して読みやすい実装。
• ane_transformers.huggingface – DistilBert などの特定のHugging Faceモデルクラスのドロップイン置き換えであり、同じPython APIを維持しながらこれらの最適化を組み込みます。

なぜ重要か – AppleのA14、A15、M1、M2などには専用のニューラルプロセッシングユニットが搭載されています。計算グラフを再構成（オペレーションの結合、埋め込み参照をCPUに移動、ANEの命令セットにモデルを適合）することで、このライブラリは、Core ML経由でデバイス上で実行した場合、通常のPyTorch/torch-scriptモデルと比較して最大10倍の低遅延と最大14倍のピークメモリ削減を実現すると主張しています。

使い方 – READMEは、Hugging Faceの distilbert-base-uncased-finetuned-sst-2-english モデルを完全にチュートリアル形式で説明しています：

1. transformers.AutoModelForSequenceClassification で標準モデルを読み込む。
2. ane_transformers.huggingface.distilbert.DistilBertForSequenceClassification に置き換え、元のstate dictをコピーする。
3. 入力のトークン化、torch.jit.trace によるモデルのトレース、coremltools.convert を使用したトレース済みモデルからCore ML mlprogram パッケージへの変換。
4. 生成された .mlpackage をXcodeプロジェクトにドロップし、XcodeのパフォーマンスタブでiPhoneまたはMac上でスピードアップを検証する。

インストール – パッケージはPyPIに公開されています：

pip install ane_transformers          # 事前ビルド済みのwheel（最も高速）
# または、開発用に：
pip install -e .

tokenizers 依存関係でビルドに失敗した場合、READMEにはRustコンパイラを用いた既知の修正方法が記載されています。

対応ハードウェア – 最適化はANEを搭載したApple Silicon（iPhone 12/A14以降、A14+搭載iPad、M1以降のMac）を対象としています。古いハードウェアで実行するとユニットテストが警告を出しますが、生成されたCore MLモデルは依然として実行可能（ただし、宣伝されたスピードアップは得られない）です。

制限 / 注意点

• 現在、Hugging Faceモデルの数はわずか（例：DistilBERT）。他のアーキテクチャへの拡張には追加作業が必要です。
• 最適化されたモデルには、一度限りのコンパイルコストが追加されます。初回のロードは遅くなります。
• 606のオペレーションのうち4つはCPU上に残っています（埋め込み参照）が、これはこのモデルサイズでは意図的な設計です。
• パフォーマンス向上はシーケンス長やバッチサイズに依存します。READMEでは、長いシーケンス（例：512トークン、バッチ8）で最大10倍の遅延削減が報告されています。

誰がこのプロジェクトに関心を持つべきか – iOS/macOS開発者で、ローカルにNLPモデルを配信し、可能な限り低いデバイス内遅延とメモリ使用量を必要とする人、およびAppleのANE向けTransformerグラフの最適化方法に関心を持つ研究者。