Googleは、携帯電話で0.2秒で画像を生成できる超高速拡散モデルMobileDiffusionを開発、現在最速

Stable Diffusionなどの大規模なAIモデルを携帯電話などのモバイルデバイスで実行することは、業界で追求されているホットなトピックの1つになっていますが、その中で生成速度が主な制限要因となっています。

最近、Google の論文「MobileDiffusion: モバイル デバイスでの 1 秒未満のテキストから画像への生成」では、携帯電話で最速のテキストから画像への生成が提案されており、iPhone 15 Pro ではわずか 0.2 秒しかかかりません。この論文は、UFOGenと同じチームによって執筆されました。超小型拡散モデルを構築すると同時に、サンプリングの高速化に現在人気のDiffusion GANテクノロジールートも採用しました。

論文アドレス: https://arxiv.org/abs/2311.16567

以下は、1 つのステップで生成された MobileDiffusion の結果です。

では、MobileDiffusion はどのように最適化されるのでしょうか?

まず、なぜ最適化が必要なのかという質問から始めましょう。

現在最も人気のあるテキストから画像への生成は、拡散モデルに基づいています。事前トレーニング済みモデルの強力な基本画像生成機能と、下流の微調整タスクに対する堅牢な性質に依存して、画像編集、制御可能な生成、パーソナライズされた生成、ビデオ生成などの分野で拡散モデルの並外れたパフォーマンスが確認されています。

しかし、基礎モデルとしての欠点も明らかで、主に次の 2 つの側面があります。1 つ目は、拡散モデルのパラメータ数が多いため、特にリソースが限られている場合に計算速度が遅くなることです。2 つ目は、拡散モデルではサンプリングに複数のステップが必要であり、これにより推論速度がさらに非常に遅くなることです。最も人気のある Stable Diffusion 1.5 (SD) を例にとると、その基本モデルには約 10 億のパラメータが含まれています。推論のために iPhone 15 Pro でモデルを量子化したところ、50 ステップのサンプリングに約 80 秒かかりました。このような高価なリソース要件と遅いユーザー エクスペリエンスにより、モバイル デバイスでのアプリケーション シナリオが大幅に制限されます。

上記の問題を解決するために、MobileDiffusion はポイントツーポイントの最適化を実行します。 （1）モデルサイズが大きいという問題に対処するため、その中核コンポーネントであるUNetで多数の実験と最適化を実施しました。これには、計算コストの高い畳み込みとアテンション操作を合理化して下位層に移動することや、活性化関数などのモバイルデバイス向けの操作を最適化することなどが含まれます。 （２）拡散モデルが複数段階のサンプリングを必要とする問題に対処するために、MobileDiffusionは、Progressive Distillationや現在の最先端のUFOGenなどのワンステップ推論技術を研究し、実践しています。

モデルの最適化

MobileDiffusion は、オープンソース コミュニティで最も人気のある SD 1.5 UNet に基づいて最適化されています。各最適化操作の後に、元の UNet モデルと比較したパフォーマンス損失が同時に測定されます。測定指標には、一般的に使用される 2 つのメトリック (FID と CLIP) が含まれます。

マクロデザイン

上の写真の左側は、オリジナルの UNet の設計の概略図です。基本的には Convolution と Transformer が含まれており、Transformer には Self-Attention と Cross-Attention が含まれていることがわかります。

UNet を最適化するための MobileDiffusion の中心的なアイデアは、2 つのポイントに分かれています。1)畳み込みの簡素化。ご存知のとおり、高解像度の特徴空間での畳み込みは非常に時間がかかり、パラメーターの数も膨大です。ここでは、完全畳み込みについて言及します。2)注意効率の向上。 Convolution と同様に、高い Attention では、特徴空間全体の長さの計算が必要になります。Self-Attention の複雑さは、平坦化された特徴空間の長さの 2 乗であり、Cross-Attention も空間の長さに比例します。

実験では、UNet の 16 個のトランスフォーマー全体を最も低い特徴解像度の内部レイヤーに移動し、各レイヤーで 1 つの畳み込みを削除しても、パフォーマンスに大きな影響はないことが示されています。達成された効果は次のとおりです。MobileDiffusion は、元の 22 個の畳み込みと 16 個のトランスフォーマーを 11 個の畳み込みと約 12 個のトランスフォーマーに大幅に合理化できます。これらすべての注意は低解像度の特徴マップで実行されます。効率が大幅に向上するため、40% の効率向上と 40% のパラメータ削減がもたらされます。最終モデルは、上図の右側に示されています。より多くのモデルとの比較は次のとおりです。

マイクロデザイン

ここでは、いくつかの斬新なデザインのみを紹介します。より詳しい紹介については、興味のある読者は本文をお読みください。

自己注意と相互注意の分離

従来の UNet では、Transformer には Self-Attention と Cross-Attention の両方が含まれています。MobileDiffusion は、すべての Self-Attention を最低解像度の特徴マップに配置しますが、Cross-Attention は中間層に保持します。この設計により、コンピューティング効率が向上するだけでなく、モデル出力の品質も保証されることがわかりました。

ソフトマックスをreluに微調整する

Softmax は、最適化されていないほとんどのケースでは並列化が困難であり、したがって非常に非効率的であることがよく知られています。 MobileDiffusion は、relu が各ポイントのアクティベーションであり、より効率的であるため、softmax 関数を relu に直接微調整することを提案します。驚くべきことに、わずか 10,000 ステップ程度の微調整で、モデル メトリックが改善され、画像の品質が保証されました。したがって、softmax に対する relu の利点は明らかです。

分離可能な畳み込み

MobileDiffuison のパラメータ削減の鍵は、Seprable Convolution の使用です。この技術は、MobileNet やその他の研究によって、特にモバイル デバイス上で非常に効果的であることが証明されていますが、生成モデルではほとんど使用されていません。 MobileDiffusion 実験では、Separable Convolution は、特に UNet の最内層に配置すると、パラメータの削減に非常に効果的であることがわかりました。分析により、モデルの品質に損失がないことが示されました。

サンプリングの最適化

最も一般的に使用されるサンプリング最適化方法には、それぞれ 8 ステップと 1 ステップを実現できる Progressive Distillation と UFOGen があります。モデルが極限まで合理化された後もこれらのサンプルが適用可能であることを証明するために、MobileDiffusion は両方に対して実験検証を実施しました。

サンプリング最適化前後のベースラインモデルの比較は以下のとおりです。サンプリング最適化後の8ステップモデルと1ステップモデルの指標が比較的優れていることがわかります。

実験と応用

モバイルベンチマーク

MobileDiffusionはiPhone 15 Proで最速の画像出力速度0.2秒を実現！

下流タスクのテスト

MobileDiffusion は、ControlNet/Plugin や LoRA Finetune などのダウンストリーム タスクを調査しました。下の図からわかるように、モデルとサンプリングの最適化後も、MobileDiffusion は優れたモデル微調整機能を維持しています。

要約する

MobileDiffusion は、さまざまなモデルとサンプリング最適化方法を検討し、最終的にはモバイル デバイスで 1 秒未満の画像出力機能を実現しながら、下流の微調整アプリケーションも保証できるようになりました。これは、将来的に効率的な普及モデルの設計に影響を与え、モバイルアプリケーションの事例を拡大すると考えています。