機械学習とAIモデルをAppleシリコンに統合

こんにちは Qiqi Xiaoです Core MLチームで エンジニアをしています 今日はCore ML Toolsに行われた 素晴らしいアップデートを いくつか紹介します これらのアップデートは 機械学習とAIモデルを Appleシリコンに より良く導入するのに役立ちます モデルデプロイワークフローには 重要なフェーズが3つあります ここでは 準備ステップに焦点を当て いくつかの最適化を共有し デバイス上でモデルを 最も効率的に実行するための機能が 含まれていることを確認します

このセッションではすでに機械学習モデルが あることを前提としています 事前トレーニングして微調整したものでも 一からトレーニングしたものでも構いません 後者について詳しく知りたい場合は 今年の 「Train your machine learning and AI models on Apple GPUs」を ぜひご覧ください

モデルをアプリに統合するために記述する コードについては説明しません 統合についての詳細は下の関連セッションを 確認することをおすすめします

それでは始めましょう

Core ML Toolsは オープンソースのPythonパッケージで モデルをAppleフレームワークで 使用するために 最適化および変換するための ユーティリティが含まれています PyTorchモデルをCore ML形式に変換して Appleシリコンでの実行に 最適化することができます

Appleシリコンのユニファイドメモリ CPU GPUおよびNeural Engineは デバイス上の機械学習ワークロードに 低レイテンシで効率的な演算を実現します デフォルトでは単にモデルを Core ML形式に変換し 推論に Appleのフレームワークを使用するだけで アプリはAppleシリコンのパワーを 活用できるようになります

Appleシリコンは すべてのプラットフォームを支えていますが プラットフォームはそれぞれに 独自の特性と長所があります 対象とするプラットフォームと デバイスで利用可能な ストレージ メモリおよび演算の組み合わせを 考慮する必要があります これらの属性は ユースケースのモデルのサイズや精度 レイテンシ要件と一致させる必要があります モデルの準備の一環として 選択肢を検討し 様々な最適化を適用して 最適な調整方法を見つけます 一緒にその方法を見ていきましょう まずモデル圧縮から見ていきます このセクションではさまざまな圧縮技術と ワークフローの使用方法を学びます その後はJunpeiにバトンタッチして ステートフルモデルのサポートについて 詳しく説明します またTransformerアーキテクチャの 最適化についても説明します

最後に多機能モデルのサポートも 紹介します

ではまず モデル圧縮について説明します

高性能な機械学習モデルは 以前よりはるかに大きくなっています より多くのストレージ容量 高いレイテンシ より多くの電力を必要とするため デバイスへの効率的なデプロイが 難しくなっています そこで残された選択肢は1つ モデル圧縮で モデルのサイズを縮小することです

まず 昨年導入された Weight compression技術について おさらいしましょう

1つ目の手法は パレット化と呼ばれるものです パレット化では 類似した値を持つ重みがクラスタ化され クラスタの重心値を使用して表されます クラスタの重心は ルックアップテーブルに格納されます ルックアップテーブルへのNビットの インデックスマップが圧縮された重みです この例では 重みは2ビットに圧縮されています

2つ目の手法は量子化と呼ばれるものです 量子化を実行するには Floatの重み値を取得し 整数範囲に線形にマッピングします 整数の重みは量子化パラメータと呼ばれる スケールとバイアスのペアとともに保存され 後で整数をFloatにマッピングするために 使用できます

3つ目の手法は プルーニング（刈り込み）です プルーニングを使用すると モデルの重みをスパース表現で 効率的に詰め込むことができます 重み行列から始めて 最小値を0に設定します ここで保存する必要があるのは ビットマスクとゼロ以外の値だけです

既存の手法は 多くのモデルでうまく機能しています 次にこれらの手法のいくつかを 大規模なモデルに適用し さらに圧縮できるように どのように拡張したかを見てみましょう

これはStable Diffusionモデルの例です プロンプトと呼ばれる 自然言語の記述を受け取り その記述に一致する画像を生成します

変換後 一番大きなモデルのサイズが Float16の精度で5GBを超えていることが わかります これは初期の段階で iPhoneやiPadで実行するには大きすぎます このモデルをデプロイするには 圧縮する必要があります

パレット化を試してみましょう パレット化には ビット数を選択してさまざまな圧縮率を 実現できる柔軟性があります

たとえば8ビットパレット化を適用すると モデルサイズを Float16モデルの約半分に縮小できます 私のMacでは「タキシードを着た猫 キャンバスに油絵」という プロンプトで素晴らしい画像を取得できます しかし iPhoneとiOSの統合を検討する前に このモデルを2GB以下に 抑えたいと思います 6ビットならこの要件を満たすことができ 最終的にiPadで このモデルを実行できます

4ビットではどうでしょうか もう適切な画像を得られません つまりモデルが正確ではないということです 何が起こっているのでしょうか

これをより深く理解するには まずパレット化表現を 見直す必要があります iOS 17でサポートされていたのは テンソルごとのパレット化のみで すべての値が1つのルックアップテーブルに クラスタ化されています 4ビットのパレット化を使用すると テンソル全体にマッピングする クラスタの重心は16個しかありません 粒度が低いため 大きな行列では 誤差が増える可能性があります では粒度を上げる方法について 考えてみましょう

少し戻って 機械学習の分野でよく使用される用語を いくつか説明したいと思います 重みは行列として表すことができます ここでは行列の行を使って 出力チャネルを表し 列を使用して入力チャネルを表しています

圧縮をテンソル全体に適用することも 各チャネルグループに適用することも 各チャネルに 個別に適用することもできます

各行を分割してブロックを小さくし サブチャネルレベルで 圧縮を適用することもできます

図の左から右に向かって 粒度は徐々に高くなります

パレット化の話に戻ります iOS 17では完全な重みテンソルに 1つのルックアップテーブルを 割り当てることができましたが iOS 18では 各チャネルグループに1つずつ 複数のルックアップテーブルを 保存できるようになりました これにより はるかに高い精度を達成できます

線形量子化の場合 iOS 17でチャネルごとの スケールとバイアスが許可されていましたが iOS 18ではこれらの量子化パラメータを ブロックごとに指定できます

プルーニング手法でも さらに圧縮できるようになりました

以前は 密な重みはFloat精度で ゼロ以外の値にスパース化されていました 現在float型のゼロ以外の値は パレット化または量子化を使用して さらに圧縮できるため モデルをさらに小さくすることができます 2つの異なる圧縮手法の特徴を 同時に利用できるようになり スパースパレット化とスパース量子化が 可能になりました

前述の3つの拡張機能から パレット化を例にとり Pythonコードでどのように適用できるかを 見てみましょう まず OpPalettizerConfigオブジェクトを作成し モデルの圧縮方法を記述します クラスタリングに「kmeans」モードを使い 4ビットを実行しています 1つのルックアップテーブルを 16チャネルのグループに割り当てたいので 粒度を 「per_grouped_channel」に設定し group_sizeを16に設定します

設定を定義したら それをOptimizationConfigに配置し palettize_weightsを使用して モデルを圧縮できます

ご覧のように処理は簡単です

では このモデルのパフォーマンスを 見てみましょう group_sizeが16の場合 プロンプトに 一致する猫がもう1匹表示されます 生成された画像は 以前よりもはるかに良くなりました ほとんどの精度を取り戻すことができ サイズは1.29GBから1.3GBへと ごくわずかに増加しただけです この増加は 追加のルックアップテーブルによって 余分なスペースが少し占有されるためです

まとめると iOS 18では グループ化した チャネルごとのパレット化により 複数のルックアップテーブルを持つことで 粒度を高めています Neural Engineを搭載した Appleシリコンで威力を発揮します ブロック単位の量子化に加えて 8ビットから4ビットの量子化に サポートを拡張しました このモードは特に MacのGPU用に最適化されています 最後に スパース性を他の圧縮モードと 組み合わせられるようになりました これはNeural Engineでうまく機能します

モデルの圧縮についてまとめる前に 圧縮のワークフローについて 説明したいと思います これはモデルを小さくするのに役立つ もう1つの側面です これまで圧縮に使用してきた ワークフローは トレーニング後のアプローチです まず事前トレーニング済みの Core MLモデルを使用します これは非常に便利な方法で データなしで圧縮されます

しかしこのデータフリーのワークフローでは 精度と圧縮量の間で最適なトレードオフが 得られない可能性があります 圧縮率が高くなると 精度が急速に低下します

これを回避する方法の1つは トレーニング時間の圧縮を使用することです

トレーニングデータを使って PyTorchモデルを微調整しながら 重みを圧縮できます その後Core ML形式に変換できます

これにより より正確なモデルを得ることができますが 大規模なモデルでは この処理に時間がかかり 大量のデータが必要になる 可能性があります

そこで新しいワークフローを用意しました キャリブレーションデータを使用した トレーニング後の圧縮です

これはデータフリーと微調整のアプローチの 中間のトレードオフを提供します この場合 モデル圧縮のキャリブレーションに必要な データ量は限られています

微調整はまったく必要ないため 設定が簡単になり 時間も短縮されます

今年coremltools.optimizeに導入された 新しいAPIを使用すると このワークフローを 非常に簡単に実行できます

例を使って理解を深めましょう

ここでは プルーニングを例に取ります 必要なモジュールをインポートすることで 開始できます これでprune_configを 指定できるようになります ここではtarget_sparsityを 0.4に設定しています つまり重みの値の40パーセントが プルーニングされます サンプル数は128に設定しました その後LayerwaysCompressorを使用して PyTorchモデルと 定義済みの構成を取得し プルーナーを作成します 次にcalibration_data_loaderを 処理します これはモデル用に独自に定義できます

その後calibration_data_loaderを使って pruner.compressを呼び出し スパースモデルをキャリブレーションできます

ここでスパースモデルを Core ML形式に変換できます または前述したように スパースパレット化が可能なため プルーニング後にパレット化を実行して さらに圧縮できます

スパース性のメタデータ情報を渡すために 先に取得したスパースモデルを使います ここでも必要なモジュールをインポートし パレット化用に設定します ここでも4ビットの粒度を 「per_grouped_channel」に設定し group_sizeを16に設定しました

これで PostTrainingPalettizerオブジェクトを作成し palettizer.compactを使用して sparse_palettized_modelを 取得できます このPyTorchモデルは Core ML形式に シームレスに変換することもできます

40%のスパース性を適用すると Stable Diffusionモデルのサイズは さらに1.1GBに縮小されます データフリーのトレーニング後圧縮では モデルはノイズ出力しか生成できませんが データキャリブレーションワークフローでは このモデルは別の猫の画像も 生成できます キャリブレーションベースのワークフローは うまく機能し良い画像が生成されました

まとめると Core ML Tools 8では 新しい圧縮表現を試すことができます

またキャリブレーションデータを使用する 新しい圧縮ワークフローも使用できます 今年はさらに多くの新機能がありますが このセッションでは取り上げていません 詳細についてはドキュメントをご覧ください 圧縮を楽しんでください

ではJunpeiにバトンタッチして ステートフルモデルや その他の最適化について説明します

ありがとうQiqi こんにちは Core MLチームのエンジニア Junpei Zhouです ステートフルモデル Transformerモデルの最適化 多機能サポートについて説明し Appleシリコンの機械学習モデルを 準備するお手伝いをします

まずステートフルモデルの サポートから始めましょう

機械学習モデルでは通常 入力と出力しかありません 入力はモデル内の操作によって処理され 出力が生成されます

通常 モデル内のデータフローは 特定の推論実行のみを対象としています

ただし 異なる実行間で 情報を保存できるように より長い存続期間を持つものが 必要になる場合があります この場合 通常はステートが使用されます

モデルはステートからデータを読み取り 結果をステートに書き戻すことができます

ステートの情報は永続的であり 異なる実行間で使用できます

これはアキュムレータの例です すべての履歴入力の合計を追跡します

ステートは0に初期化されます

モデルはステートと入力から値を読み取り 出力を生成します

同時に結果をステートに書き込みます

同様に別の入力が入ってくると ステートから読み取り 出力を生成すると同時に ステートを更新します

Core MLがステートをサポートする前は モデルでステートを処理するには 追加の入力と出力を明示的に 定義する必要がありました

アキュムレータの例では アキュムレータの入力フィールドと 出力フィールドが必要です

次にアキュムレータ出力を 次のアキュムレータ入力に 明示的にコピーします

これは非効率かもしれません 特にステートが大きい場合は 非効率になります

今年 Core MLは ステートフルモデルのサポートを導入します

モデルがステートテンソルを自動的に 更新するようになったため ステートテンソルを入力または出力として 定義する必要はありません 代わりにステートをインプレースで更新でき パフォーマンスの向上につながります

アキュムレータの例では PyTorchでモデルを定義する方法を 説明します

モデルのステートテンソルを示すために トーチのregister_buffer APIを 使用します

これが完了すると ステートフルなトーチモデルを ステートフルなCore MLモデルに 変換することも 簡単になります

新しく導入されたct.StateTypeを使用して ステートを指定するだけで済みます ここではregister buffer APIで 使用したのと同じステート名を指定します

次にステートをct.convertに渡して ステートフルモデルを変換します

変換されたCore MLモデルは ステートフルになります トーチモデルでは 読み取りと書き込みは 通常インプレース操作と スライス更新操作によって行われます

変換中に Core ML Toolsは 変換されたCore MLモデルで 対応する操作を生成します

ステートフルなCore MLモデルを デバイスにデプロイする方法について 詳しく知りたい場合は このセッションをご覧ください

ではTransformerの最適化について 説明します

Transformerは基礎モデルで 最も人気のある モデルアーキテクチャの1つです

いくつかのマルチヘッドアテンションブロック で構築されます これは演算負荷が高く 通常はモデルの推論中の ボトルネックになります

アテンションブロックは いくつかのMatMul操作やSoftmax操作と クエリー キー 値と呼ばれる 3つの大きなテンソルで構成されます トークンを生成する場合 現在の単語のクエリーベクトルは シーケンス内の以前のすべての単語の キーベクトルに注目し アテンションスコアに基づいて 値ベクトルを融合します

つまり各タイムスタンプでは 前のトークンのキーと 値のベクトルは 前の手順ですでに計算されています

この計算の重複を避けるため キー値キャッシュ （KVキャッシュ）を使用して 各手順で計算した キーと値のベクトルを保存し

そのあとの手順で使えるようにします

KVキャッシュは大規模な言語モデルで デコードを高速化するためによく使用される 効果的な手法です これは前に紹介した ステートフルモデルを使用するのに 最適なケースです

ステートがない場合 KVキャッシュは I/Oでのみ使用できますが KVキャッシュには通常大きなテンソルが あることを考えると これは非効率です

ステートがあれば より効率的に インプレースアップデートを行えます

このステートフルKVキャッシュは AppleシリコンでTransformerモデルを 準備するのに役立ちます

アテンションブロックの構造に戻りましょう ステートフルKVキャッシュによって アテンションに対するキーと値の 準備効率が向上する仕組みを 紹介しました 当然の帰結として

アテンションの計算を もっと効率化できるかどうかを考えます

これによりTransformerモデルに 別の最適化がもたらされます

Scaled Dot Product Attention演算 （SDPAとも呼ばれます）の 融合表現です

以前はSDPA演算を使用して トーチモデルを変換すると Core ML ToolsはSDPA操作を いくつかの小さな操作に分割しました これにはmatmulやsoftmax その他の操作が含まれます

今年は最小配備ターゲットが iOS 18に Core ML Toolsは 変換されたCore MLモデルで SDPA操作を使用します

このSDPA操作は入力をすべて 一度に受け取り アテンションをより効率的に計算します これはAppleシリコンGPUで 非常に威力を発揮します

Metalを使用したSDPA操作の 最適化について詳しくは ぜひこのセッションをご覧ください

これらの素晴らしい機能が導入されたので Core ML Toolsを使用して Mistral-7BモデルをAppleシリコン用に 準備するデモを見てみましょう Mistral-7Bは パワフルな大規模な言語モデルで 73億個のパラメータがあります 今年導入されたステートフルKVキャッシュと SDPA操作を使用して Appleシリコン上でスムーズに実行できる Core MLモデルに変換できます

このJupyter Notebookを使って Mistral-7Bモデルをトレーニング後の ブロック単位の量子化を備えた ステートフルCore MLモデルに変換する 方法を説明しましょう

コードをすっきりと理解しやすくするために ユーティリティをいくつかインポートします

まずオリジナルのMistral 7Bモデルを ステートフルにします

オリジナルのMistral-7Bモデルから始めて モデルをステートフルにして KVキャッシュを効率的に処理できるように モデルを変更する必要があります

KVキャッシュ用のバッファを登録できます これらを「keyCache」と 「valueCache」と呼びます

モデルにはほかにも変更が 加えられています たとえばフォワードパス中に KVキャッシュのインプレース更新を行います

詳細については 後でユーティリティで確認できます

フォワード関数では KVキャッシュ用のバッファは すでに登録済みなので ここで入力する必要はありません

こうしたすべての変更により KVキャッシュを永続的なステートとして持つ ステートフルMistral-7Bモデルが得られます

これでHuggingFaceから自動的に ダウンロードされた重みを使用して このモデルを初期化し モデルを評価モードに設定できます

サンプル入力を指定すると Mistral-7Bモデルの場合は input_idsとcausal_maskが必要になり その後モデルをトレースできます

このモデルを変換するには 通常どおり入力と出力を指定します ここでは入力の形状に ct.RangeDimを使用しています プロンプトの長さが実行ごとに 異なる可能性があるためです

変換されたCore MLモデルにKVキャッシュの ステートも含むようにするには 今年新しく導入された ct.StateTypeを使用して ステートを指定する必要があります

ここでは「keyCache」と 「valueCache」という名前を使用して トーチモデルにバッファを登録したときの 名前と一致するようにします

最後に 入力 出力 ステートで ct. convertを呼び出し minimum_deployment_targetを iOS 18に設定します

このステートをct. convertに渡すと 変換されたCore MLモデルは インプレースで更新できる kキャッシュとvキャッシュを備えた ステートフルモデルになります

変換後 モデルをディスクに保存して モデルのサイズを確認してみましょう

変換が完了しました モデルのディスクサイズを確認しましょう

13GBとかなり大規模なモデルです

モデル予測を実行して モデルをテストしてみましょう モデル予測では 内部で Core MLフレームワークが使用されています これは6月にシアトルで 訪れるべき場所を モデルに推奨させるためのプロンプトです

ご覧のように モデルが非常に大きいにも関わらず 変換されたステートフルCore MLモデルは 私のMacではスムーズに実行できます

このFloat16モデルは少し大きすぎます 先ほどQiqiが説明した モデル圧縮手法のいくつかを 使ってみましょう

ct.optimizeモジュールを使用すると 圧縮構成を指定できます

OpLinearQuantizerConfigでは dtypeをint4 粒度を「per_block」 block_sizeを32に指定します

またモデル内のすべての重みを 圧縮するため global_configに設定します

構成は必要に応じて 操作の種類または操作名で指定できます

その後ct.optimizeモジュールの linear_quantize_weightsツールを使用して 構成に従ってモデルを量子化します

指定した構成に基づいて 各ブロックに32個の要素がある 「int4」「per_block」の量子化になります このツールでは Core ML Toolsが モデルに含まれるすべての重みを調べ 線形量子型アルゴリズムで圧縮します

量子化が完了しました

量子化されたモデルのサイズを 確認しましょう

いいですね 量子化モデルのディスクサイズは 元のFloat16モデルよりもはるかに小さく 13GBから4GB未満に減少しました

同じプロンプトを使って この量子化モデルを実行してみましょう

素晴らしい このはるかに小さいモデルは 元のFloat16モデルよりも さらに高速に実行でき 生成される出力の品質も同様です

このデモで示したように ステートフルKVキャッシュと量子化を備えた Mistral-7Bモデルを準備するのは 簡単です 同様のワークフローに従って よく使う言語モデルを Appleシリコンに簡単に導入できます

最後に機能をもう1つ説明します 多機能Core MLモデルの サポートについてです

複数の機能で構成されるモデルを 考えてみましょう 予測にはこれらの機能を 組み合わせて実行する必要があります たとえば特徴抽出器のあとに 分類出力を生成する分類器レイヤが続く モデルがあるとします

また回帰出力を生成する 回帰レイヤもあります

この2つを組み合わせて 1つのモデルにすることができ 両方のタスクに使用できます

PyTorchでこれがどのように見えるかを 示します

分類器とリグレッサーは 同じ特徴抽出器を共有し それぞれがCore MLモデルに 変換されています

今年は多機能モデルのサポートを 導入したので Core ML Toolsを使用して この2つのモデルをマージできます

マージ中 Core ML Toolsは 重みのハッシュ値を計算して 重み共有の重複を可能な限り排除します

マージされた多機能モデルは 共有の特徴抽出器を使って 両方のタスクを実行できます

これはCore ML Toolsを使って Core MLモデルをマージする サンプルコードスニペットです

2つのモデルをct. convertで変換し 個別のmlpackagesに保存したあと

MultiFunctionDescriptorを作成して マージするモデルと マージされたモデルの 新しいfunction_nameを指定できます

次にsave_multifunctionユーティリティを 使用して マージされた多機能Core MLモデルを 作成できます

Core ML Tools Python API経由で 多機能モデルを読み込む場合 function_nameを指定して 特定の機能を読み込み 通常どおり予測を行うことができます

多機能モデルは 大規模なモデルで非常に便利です

大規模なモデルを微調整する場合 モデル全体を微調整するのは通常 コストがかかりすぎます

代わりに微調整中にアダプタを 使用する方が一般的です

アダプタは通常 パラメータが少なく小さいものですが 完全に微調整されたモデルと同等の パフォーマンスを発揮します

同じベースモデルでも フィールドや下流のタスクごとに 異なるアダプタが存在する場合があります

アダプタブロックは通常 モデルの先頭や末尾にあるだけでなく 中間レイヤとやり取りすることもできます

Core MLの新しい多機能モデルを使うと 1つのモデルで複数のアダプタを キャプチャすることができます アダプタごとに別々の機能が定義されます

ベースモデルの重みは 可能な限り共有されます

アプリに複数の機能を統合する方法の 詳細については 「Deploy machine learning and AI models on-device with Core　ML」 セッションをご覧ください

テキストを画像に変換する例では 複数のアダプタを備えた 単一のモデルを使用して 様々なスタイルの画像が生成されます

まとめましょう 新しいモデル圧縮手法とワークフローを使うと モデルのサイズを大幅に削減しながら 優れた精度を実現できます モデルでステートを使用して 複数の機能を提供するのも簡単です

大規模なTransformerモデルの Appleシリコンへの導入が これまでになく簡単になりました

詳しい情報については Core ML Toolsのドキュメントを ご覧ください ありがとうございました