WWDC26:Meet Core AI

Core AI を使ってオンデバイスに学習モデルの作成と改善のサイクルを回したり、オンデマンドにユーザーデバイス上でモデルをコンパイルできるらしい。そもそも、PyTorch から勉強したいと思った。

Meet Core AI


0:33 – What is Core AI

  • Core AI は Apple Intelligence をオンデバイスで動かす推論フレームワークで、WWDC26 より開発者に公開
  • モデルデプロイのライフサイクル全体をカバー
    • モデルの最適化、コンバージョン、デバッグ、アプリへの統合の高速かつ反復的サイクルに対応
    • CPU・GPU・Neural Engineすべての Apple Silicon を活用
    • モダンな Swift API、Python ツールチェーン、専用デベロッパーツールを同梱
  • 主要コンポーネント
    • CoreAI Swift フレームワーク:アプリ側の推論実行
    • coreai-torch Python パッケージ:PyTorch モデルの変換
    • Core AI Debugger:数値デバッグ用スタンドアローンアプリ
  • 対応ニーズ例:話者分離モデル、カメラ映像をもとに質問へ回答、複雑なマルチステップタスク

4:57 – Model conversion

  1. モデル変換の手順
    1. PyTorchで行動予測モデルを試作
    2. ゲームを実行しトレーニングデータを蓄積
    3. PyTorchモデルをCore AIに変換(Core AI PyTorch extensions):以下
  2. coreai-torch を使って PyTorch モデルを Core AI フォーマットへ変換する手順
    1. torch.export.export() でモデルをエクスポート、動的シェイプ(seq_len など)を指定
    2. run_decompositions() で Core AI 互換の分解テーブルを適用
    3. TorchConverter().add_exported_program() で Core AI グラフへ変換
    4. save_asset("SnakeTransformer.aimodel").aimodel ファイルとして保存
  3. 変換後の数値検証
    1. PyTorch の出力と Core AI の出力を numpy で比較し、最大誤差が 0.01 未満であることを確認

6:16 – App integration

  • 生成したモデルをアプリに統合→実行
  • Xcode のモデルビューアーで .aimodel を検査可能
    • 入出力のテンソル形状、関数名、サイズなどを確認できる
  • CoreAI Swift フレームワークの主要型
    • AIModel.aimodel ファイルをロード
    • InferenceFunctionmodel.loadFunction(named:) で取得
    • NDArray:n 次元テンソルデータの入出力型
  • 推論の実行フロー
    • AIModel(contentsOf: modelURL) でモデルをロード
    • model.loadFunction(named: "main") で推論関数を取得
    • NDArray に入力データを書き込み
    • nextActionFunction.run(inputs:) で推論を実行し出力 NDArray を取り出す

10:48 – Profiling with Instruments

  • モデルのパフォーマンス改善
    • 例:時間経過と共に遅くなる問題の解決
  • Xcode に新設された Core AI インストゥルメントでモデルのレイテンシをプロファイル可能
  • Transformer モデルでシーケンス長が増えるにつれ、Key-Value の埋め込み再計算を実行する
    • → シーケンス長増大に応じて、推論時間が二次関数的に増大する問題
    • シーケンス各要素に計算済みの Key-Value キャッシュすることで解決

11:15 – Optimizing performance

  • KV キャッシュをステートフル入力として追加することで推論のスローダウンを解消
  • PyTorch 側の変更
    • register_buffer() でキーキャッシュ・バリューキャッシュをモジュールバッファとして登録
    • forward() でキャッシュの読み書きを実装
    • 再変換時に state_names=["keyCache", "valueCache"] を指定
  • Swift 側の変更
    • NDArray としてキャッシュバッファを確保し ModelPlayer の保存プロパティとして保持
    • InferenceFunction.MutableViews にキャッシュの MutableView を挿入して推論時に渡す
      • stateViews.insert(&keyCache, for: "keyCache") のように記述

14:13 – Additional features

  • coreai Python パッケージのリッチなオーサリング体験
    • 既存の変換・最適化 API に加え、モデルグラフの細粒度な操作が可能
  • Core AI Debugger
    • 変換済みモデルの数値デバッグ専用ツール
    • 中間テンソル出力の検査、PyTorch との差分確認
  • Xcode のデバッグゲージ
    • Xcode でアプリを実行中に Core AI の活動状況をリアルタイムにストリーミング監視

15:34 – Specialization

  • モデルの特化プロセス:Core AI がターゲットデバイス向けにモデルをコンパイルする処理
    • 初回ロード時に実行され、以後はキャッシュから即時ロード
    • かなり時間がかかる場合があるため、ユーザーフローの計画が必要
  • プログラム的なキャッシュ管理
    • AIModelCache.default でキャッシュの有無を確認
    • キャッシュがない場合にユーザーへ「AI 機能を準備中」と通知する UX パターン
    • AIModel.specialize(contentsOf:) で明示的にスペシャライゼーションを要求可能
  • SpecializationOptions でコンパイル動作をカスタマイズ
  • 特化で起こる2段階の変換
    • Compilation:コンピューティングを分割、計画し最適化
    • Excecutable generation:使用する計算ユニット向けのアーティファクト生成(デバイスとOSバージョンに紐付けられる)
  • Compliation ステップが大部分を占める → Core AI ツールチェインの事前コンパイル(Ahead-of-time, AOT compilation)で解決