BlockBeatsの情報によると、3月3日、開発者Manjeet Singh(GitHub: maderix)がClaude Opusと協力し、Appleの非公開のプライベートAPIを逆エンジニアリングすることで、M4チップのApple Neural Engine(ANE)上で逆伝播を含むニューラルネットワークのトレーニングを初めて実現しました。ANEはAppleが推論専用に設計したアクセラレーターであり、公式にはトレーニング機能は提供されておらず、開発者はCoreMLフレームワークを通じてその推論機能を間接的に呼び出すしかできませんでした。
このプロジェクトは、CoreMLを迂回し、_ANEClient、_ANECompilerなどの40以上のプライベートクラスからIOKitカーネルドライバーまでの完全なソフトウェアスタックを直接マッピングし、メモリ上でモデルを直接コンパイルできる_ANEInMemoryModelDescriptorインターフェースを発見しました。これは重みの更新ごとに再コンパイルが必要であるため、トレーニングを実現する上で鍵となります。現在、単一のTransformer層(dim=768、seq=512)のトレーニングを実装しており、M4では1ステップあたり9.3ms、ANEの利用率は11.2%(1.78 TFLOPS、理論上のピークは15.8 TFLOPS)です。フォワードおよびバックワード伝搬の入力勾配はANEで計算され、重み勾配とAdamオプティマイザーはCPUで処理されます。
プロジェクトは、ANEの核心計算プリミティブが行列乗算ではなく畳み込みであることを発見した。1x1畳み込みを用いて行列乗算を表現すると、スループットが約3倍向上し、CoreMLを経由せずに直接呼び出すとさらに2~4倍の追加改善が得られる。Appleが公表した「38 TOPS」は誤解を招くものである。現在、プロジェクトは初期段階にあり、単一層のトレーニングのみをサポートし、合成データを使用しており、コンパイル後のリソースリークが約119回発生するため、プロセス再起動で回避している。複数層のトレーニングと実データのサポートは開発中である。プロジェクトはMITライセンスでオープンソース化され、公開から5日間で約2800スターを獲得した。
