iPhone、iPad、MacでのCore MLを使った高速なStable Diffusion

TL;DR

AppleとHugging Faceは、iPhone、iPad、MacでStable Diffusionをより高速に、かつメモリオーバーヘッドを低減して実行できるように、新しいCore MLの最適化、具体的には6ビットパレタイゼーションと更新された注意層の実装を導入しました。これらの改善により、精度の損失を最小限に抑えながらモデルの重みを大幅に圧縮し、デバイス上での生成AIをより効率的に行えるようになります。

Core ML Optimizations and 6-Bit Palettization

Core MLは、AppleデバイスのCPU、GPU、Neural Engineを活用して、デバイス上で機械学習モデルを実行します。Stable Diffusionのような大規模モデルの効率を向上させるため、Appleはcoremltools.optimizeサブモジュールを介して6ビットパレタイゼーションと呼ばれる新しい圧縮技術を導入しました。

How Palettization Works

パレタイゼーションは、モデルの重みを16ビット浮動小数点表現から1パラメータあたり6ビットに圧縮する量子化の一形態です。これはコンピュータグラフィックスにおけるカラーパレットと同様に機能します：固定された数の色（パレット）が定義され、画像の色はそのパレット内で最も近い利用可能な色のインデックスに置き換えられます。

Memory and Execution Efficiency

以前のCore MLのバージョンでは、ディスクからロード時に圧縮された重みが解凍され、メモリ使用量は未圧縮モデルのサイズと同じ状態が続いていました。更新されたフレームワークでは、推論が層から層へと進行する際に、パレタイズされた重みをオン・ザ・フライで変換します。このアプローチは、未圧縮の計算負荷を増やすのではなく、実行における主要なボトルネックであるメモリ転送の量を減らすことで高速化されます。

Updated Stable Diffusion Implementation

Appleのml-stable-diffusionリポジトリ（diffusersライブラリに基づく）は、これらの新しい最適化を組み込むために更新されました：

  • 量子化サポート: --quantize-nbitsフラグを使用して、モデルを8、6、4、または2ビットに量子化できるようになりました。6ビット量子化は、精度とパフォーマンスの最適なバランスを取るために推奨されます。
  • 注意層の最適化: SPLIT_EINSUM_V2と呼ばれる新しいメソッドは、クエリーシーケンスを512のチャンクに分割し、大きな中間テンソルの生成を避けます。この最適化により、Neural Engineでのパフォーマンスが10%から30%向上する可能性があります。

Available Optimized Models

以下の公式Stable Diffusionモデルの6ビットパレタイズドバージョンが、Hugging Face Hubで今利用可能になりました：

モデル 未圧縮 6ビットパレタイズド
Stable Diffusion 1.4 Core ML float16 Core ML 6ビットパレタイズド
Stable Diffusion 1.5 Core ML float16 Core ML 6ビットパレタイズド
Stable Diffusion 2 base Core ML float16 Core ML 6ビットパレタイズド
Stable Diffusion 2.1 base Core ML float16 Core ML 6ビットパレタイズド

Deployment and Custom Model Conversion

System Requirements

6ビットモデルを利用するには、デバイスはiOS/iPadOS 17または**macOS 14 (Sonoma)**の開発バージョンを実行している必要があります。これらには必要な更新されたCore MLフレームワークが含まれています。

Converting Custom Models

開発者は、ml-stable-diffusionリポジトリを使用してファインチューン済みまたはDreamboothモデルを変換できます。プロセスは以下の通りです：

  1. coremltools 7.0 ベータと Xcode 15.0 ベータを使用します。
  2. --quantize-nbits 6 フラグを付けて torch2coreml 変換スクリプトを実行します。
  3. macOSでは通常より良いORIGINAL注意実装と、iOSでは通常より高速なSPLIT_EINSUM_V2の間から選択します。

Advanced Compression: Training-Time Quantization

6ビットパレタイゼーションは事後トレーニング圧縮の一例ですが、coremltoolsはトレーニングタイム圧縮も導入しています。これにより、開発者は重みの圧縮を同時に行いながらモデルをファインチューンすることができます。重みのクラスタリングに微分可能なアルゴリズムを使用することで、トレーニング中に量子化テーブルを最適化し損失を最小限に抑えることができ、これにより事後トレーニング法よりも品質の劣化が少ない4ビットまたは2ビット圧縮が可能になる可能性があります。

Sources