開発者がApple Neural Engineの非公開APIをリバースエンジニアリングし、初めてANE上でニューラルネットワークのトレーニングを実現

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インターフェースを発見しました。これは訓練を実現するための重要なポイントであり、各重みの更新には再コンパイルが必要です。現在は、単一のトランスフォーマー層（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スターを獲得しています。