Intel CPUのNNCFと🤗 Optimumを使用した安定したディフュージョンの最適化

'Intel CPUのNNCFと🤗 Optimumを使用した安定したディフュージョンの最適化' can be condensed as 'Intel CPUのNNCFと🤗 Optimumを使用したディフュージョンの最適化'.

潜在的な拡散モデルは、テキストから画像の生成問題を解決する際にゲームチェンジャーとなります。 安定した拡散は、コミュニティや産業界で広く採用されている最も有名な例の一つです。 安定した拡散モデルのアイデアはシンプルで魅力的です：ノイズベクトルから画像を複数の小さなステップで生成し、ノイズを潜在的な画像表現に洗練させます。

ただし、このようなアプローチは、全体的な推論時間を増加させ、クライアントマシンで展開された場合にユーザーエクスペリエンスの低下を引き起こします。 通常のように、強力なGPUがここで役立つことに注意することができますが、これに伴うコストも著しく増加します。 参考までに、H1’23では、8つのvCPUと64GBのRAMを備えた強力なCPU r6i.2xlargeインスタンスの価格は1時間あたり$0.504であり、同様のNVIDIA T4を搭載したg4dn.2xlargeインスタンスの価格は1時間あたり$0.75で、これは1.5倍以上です..

これにより、画像生成サービスは所有者とユーザーにとって非常に高価になります。 クライアントアプリケーションでは、GPUがまったくない場合もあります！ これにより、安定した拡散パイプラインの展開は困難な問題となります。

過去5年間、OpenVINO Toolkitは高性能推論のための多くの機能をカプセル化しました。 最初はコンピュータビジョンモデルに設計されたものですが、現在でも最先端のモデルを含む多くのコンテンポラリーモデルにおいて、最高の推論パフォーマンスを示しています。 ただし、リソース制約のあるアプリケーションに安定した拡散モデルを最適化するには、ランタイム最適化にとどまらず、さらに進んだモデル最適化機能がOpenVINO Neural Network Compression Framework（NNCF）から必要とされます。

• オープンソースAIゲームジャムを発表します 🎮
• AI音声認識をUnityで
• ファルコンはHugging Faceのエコシステムに着陸しました

このブログ記事では、安定した拡散モデルの最適化の問題を概説し、CPUなどのリソース制約のあるHWで実行される場合に、そのようなモデルのレイテンシを大幅に削減するワークフローを提案します。 特に、PyTorchと比較して5.1倍の推論高速化と4倍のモデルフットプリントの削減を達成しました。

安定した拡散の最適化

安定した拡散パイプラインでは、UNetモデルが計算上最もコストがかかります。そのため、単一のモデルの最適化によって推論速度が大幅に向上します。

しかし、このモデルに対しては、従来のモデル最適化手法であるポストトレーニングの8ビット量子化は機能しないことがわかりました。その理由は2つあります。まず、セマンティックセグメンテーション、スーパーレゾリューションなどのピクセルレベル予測モデルは、タスクの複雑さにより、モデル最適化の観点では最も複雑なものの一つであり、モデルパラメータと構造の微調整が結果を多数の方法で崩してしまいます。 2番目の理由は、モデルが数億のサンプルでトレーニングされるため、冗長性のレベルが低いことです。 そのため、研究者は最適化後も精度を保持するためにより洗練された量子化手法を使用する必要があります。 たとえば、Qualcommは安定した拡散モデルを量子化するためにレイヤーワイズの知識蒸留法（AdaRound）を使用しました。 これはつまり、量子化後にモデルの調整が必要であることを意味します。 それならば、ソースモデルがトレーニングされる方法と同じように、モデルと量子化パラメータを同時に調整できるQuantization-Aware Training（QAT）を使用するだけでないでしょうか？ そこで、私たちはNNCF、OpenVINO、Diffusersを使用してこのアプローチを試しました。

最適化ワークフロー

通常、モデルの最適化はトレーニング後に開始します。ここでは、ポケモンのデータセットに基づいてファインチューニングされたモデルから始めます。このデータセットには、ポケモンの画像とテキストの説明が含まれています。

私たちは、Diffusersから安定した拡散のテキストから画像へのファインチューニングの例を使用し、NNCFからQATを統合した次のトレーニングスクリプトを使用しました。また、実際にトレーニングされるモデルが学生として機能するこのプロセスで、ソースモデルが教師として機能するように、損失関数を変更し、知識蒸留を組み込みました。このアプローチは、トレーニングされた教師モデルをより小さい学生モデルに蒸留する従来の知識蒸留法とは異なります。私たちの場合、知識蒸留は、最適化モデルの最終的な精度を向上させるのに役立つ補助的な方法として使用されます。モデルパラメータを含む指数移動平均（EMA）メソッドを使用して、トレーニングプロセスをより安定させることもできます。モデルを4096回のイテレーションでチューニングしました。

勾配チェックポイントとVRAMではなくRAMにEMAモデルを保持するなどのいくつかのトリックを使用することで、24GBのVRAMを備えた1つのGPUで最適化プロセスを実行することができます。1つのGPUを使用して、最適化全体は1日未満で完了します！

量子化意識トレーニングを超えて

量子化だけでも、モデルのフットプリント、読み込み時間、メモリ消費量、および推論の待ち時間を大幅に改善できます。しかし、量子化の素晴らしいところは、他の最適化手法と併用することができ、累積的な高速化を実現できるということです。

最近、Facebook ResearchはVision Transformerモデル向けのToken Merging手法を導入しました。この手法の本質は、利用可能な戦略（平均、最大値など）のいずれかを使用して冗長なトークンと重要なトークンをマージすることです。これは、Transformerモデルの計算量が最も多い部分であるセルフアテンションブロックの前で行われます。したがって、トークンの次元を削減することで、セルフアテンションブロック全体の計算時間が削減されます。この手法はStable Diffusionモデルにも適用され、GPU上で高解像度の画像合成を最適化する際に有望な結果を示しました。

私たちは、Attention UNetモデルに適用する場合に、OpenVINOに準拠したToken Merging手法を修正し、8ビットの量子化と組み合わせました。これにはKnowledge Distillationなど、すべての手法が含まれます。量子化に関しては、精度を回復するためにファインチューニングが必要です。また、最適化とファインチューニングは、Pokemonsデータセットでトレーニングされたモデルから開始されます。下の図は、全体的な最適化ワークフローを示しています。

結果として得られるモデルは、クライアントやエッジCPUなどの計算リソースに制約のあるデバイスでの推論実行に非常に有益です。先に述べたように、Token Mergingを量子化と組み合わせることで、推論の待ち時間をさらに短縮することができます。

最適化されたモデルを使用したイメージ生成デモの結果。入力プロンプトは「cartoon bird」で、シードは42です。モデルは、Intel® Deep Learning Boostテクノロジーを搭載した第3世代Intel® Xeon® Scalableプロセッサを使用した「CPUアップグレード」インスタンスで、OpenVINO 2022.3で実行されています。

結果

公開された最適化ワークフローを使用して、2種類の最適化されたモデル（8ビット量子化とToken Mergingで量子化）を取得し、PyTorchのベースラインと比較しました。また、総合的な比較のために、ベースラインをバニラのOpenVINO浮動小数点（FP32）モデルに変換しました。

上の図は、イメージ生成の結果といくつかのモデルの特性を示しています。ご覧のように、OpenVINOへの変換だけでも推論の待ち時間が大幅に減少します（ 1.9倍 ）。8ビット量子化の適用により、PyTorchと比較して推論速度がさらに向上し、3.9倍 の高速化が実現されます。量子化のもう一つの利点は、モデルのフットプリントが大幅に削減され、PyTorchのチェックポイントの 0.25倍 になることで、モデルの読み込み時間も改善されます。量子化の上にToken Merging（ToME）（マージ比率0.4）を適用することで、フットプリントを同じレベルに保ちながら、5.1倍 のパフォーマンス向上が実現されます。最適化されたモデルの視覚的品質については、詳細な分析は提供していませんが、結果は非常に優れていることがわかります。

以下に、Intel CPU上で実行するために最適化された最終パイプラインでの推論の実行方法を示します：

from optimum.intel import OVStableDiffusionPipeline

# パフォーマンスのためにパイプラインをロードおよびコンパイルします。
name = "OpenVINO/stable-diffusion-pokemons-tome-quantized-aggressive"
pipe = OVStableDiffusionPipeline.from_pretrained(name, compile=False)
pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
pipe.compile()

# 画像を生成します。
prompt = "緑色のポケモンで赤い目の絵を描く"
output = pipe(prompt, num_inference_steps=50, output_type="pil").images[0]
output.save("image.png")

トレーニングおよび量子化のコードは、Hugging Face Optimum Intelライブラリで見つけることができます。最適化されたモデルとオリジナルモデルの違いを示すノートブックは、こちらで入手できます。Hugging Face HubのOpenVINO組織の下には、多くのモデルも見つけることができます。さらに、Hugging Face Spacesで実行されているデモも作成しました。このデモは、第3世代Intel Xeon Scalableプロセッサを搭載したr6id.2xlargeインスタンスで実行されています。

汎用Stable Diffusionモデルはどうですか？

Pokemon画像生成タスクで示したように、比較的少量のトレーニングリソースを使用してStable Diffusionパイプラインの高度な最適化を実現することが可能です。同時に、一般的なStable Diffusionモデルのトレーニングは高コストなタスクであることはよく知られています。ただし、十分な予算とHWリソースがあれば、説明した手法を使用して一般的なモデルを最適化し、高品質な画像を生成することが可能です。ただし、モデルの容量を大幅に削減するトークンマージング手法には注意が必要です。ここでの鉄則は、トレーニングに使用するデータセットが複雑であるほど、最適化中に使用するマージ比率を低くすることです。

この投稿を読んで楽しんだ場合、第4世代のIntel Xeon CPUのパフォーマンスを最適化するための他の補完的なアプローチについて説明したこの投稿も興味深いかもしれません。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!