MicroDiffusion: 数万ドルのモデルに勝る千ドルの生成画質モデル
3つの要点
✔️ Text-to-Image拡散モデルは多くの分野で活用されているが、高額な費用と膨大な計算リソースが求められる
✔️ MicroDiffusionは、新しいマスク手法と改善されたTransformerアーキテクチャを採用し、低予算での拡散モデルを実現する
✔️ 実験結果によると、MicroDiffusionは現在の最先端モデルと比較して14分の1のコストで、同等のFIDと高品質な生成を達成した
Stretching Each Dollar: Diffusion Training from Scratch on a Micro-Budget
written by Vikash Sehwag, Xianghao Kong, Jingtao Li, Michael Spranger, Lingjuan Lyu
(Submitted on 22 Jul 2024)
Comments: 41 pages, 28 figures, 5 tables
Subjects: Computer Vision and Pattern Recognition (cs.CV); Artificial Intelligence (cs.AI); Machine Learning (cs.LG)
code：
本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。
はじめに
最新の画像生成モデルは、自然で高品質なコンテンツの作成に優れており、年間10億以上の画像を生成しています。しかし、これらのモデルをゼロから訓練するには、非常に高額な費用と膨大な時間がかかります。テキストから画像への拡散モデル（T2I）は、計算コストを一部削減しましたが、それでもかなりのリソースが必要です。
現行の最先端技術では、約18,000時間のA100 GPU時間が必要で、8台のH100 GPUを使った訓練には1ヶ月以上かかります。さらに、大規模または専有のデータセットに依存することが多く、広範な利用が難しいのが現状です。
今回の解説論文では、低コストでエンドツーエンドのテキストから画像への拡散モデルパイプラインを開発し、大規模なデータセットなしでコストを大幅に削減することを目指します。ビジョントランスフォーマーベースの潜在拡散モデルに焦点を当て、そのシンプルな設計と広範な適用性を活用しています。計算コストを削減するために、画像ごとに処理するパッチ数をランダムに入力トークンをマスキングすることで減らす方法を採用しています。既存のマスキング方法では高いマスキング比で性能が低下する課題を、本論文では克服することを目指しています。
テキストから画像への拡散モデルでの性能低下を克服するため、本論文では「遅延マスキング」戦略を提案しています。軽量なパッチミキサーでパッチを処理し、その後拡散トランスフォーマーに入力することで、マスキング比が高くても意味情報を保ちつつ、コストを抑えて信頼性の高い訓練を実現します。また、トランスフォーマーアーキテクチャの最新進展も取り入れ、大規模訓練での性能向上を図っています。
実験では、わずか1,890ドルの予算で、3,700万枚の画像と75%のマスキング比を用いて、1.16億パラメータのスパース拡散トランスフォーマーを訓練しました。その結果、COCOデータセットでのゼロショット生成において12.7のFIDを達成しました。訓練時間は単一の8×H100 GPUマシンでわずか2.6日で、現行の最先端アプローチ（37.6日、GPUコスト28,400ドル）と比べて14倍の短縮を実現しました。
提案手法
遅延マスキング
トランフォーマの計算量はシーケンスの長さに比例するため、学習コストを削減するために、図１-bのように、大きいパッチサイズによりシーケンスを削減という方法があるます。大きいパッチサイズを使用すると、画像あたりのパッチ数は二次関数的に減少しますが、画像の大きな領域を1つのパッチに積極的に圧縮するため、性能が著しく低下する可能性があります。
パッチサイズを維持しつつ、図１−ｃのように、マスクを用いてトランスフォーマーの入力層で多数のパッチを削除する方法があります。これは、畳み込みUNetでのランダムクロップ訓練に似ていますが、パッチマスキングは画像の非連続領域での訓練を可能にします。この方法は視覚と言語のドメインで広く採用されています。
マスクされたパッチからの表現学習も促すために、図１−ｄのMaskDiTは、マスクされたパッチの再構成を促す補助的な自己符号化損失を追加しています。この手法は入力画像の７５％をマスキングして、計算コストを大幅に削減できました。
しかし、高いマスキング比は、トランスフォーマーの全体的な性能を大幅に低下させます。MaskDiTを使用しても、単純なマスキングと比べてわずかな改善しか見られません。なぜなら、このアプローチでも入力層で画像パッチの大部分が削除されるためです。
本論文では、「パッチミキサー」と呼ばれる前処理モジュールを導入し、マスキング前にパッチ埋め込みを処理します。これにより、マスキングされていないパッチが画像全体に関する情報を保持し、学習効果を向上します。この方法は、既存のMaskDiT戦略と同等の計算コストでありながら、性能を改善する可能性があります。
パッチミキサーと学習損失
パッチミキサーとは、個々のパッチの埋め込みを融合できる任意のニューラルアーキテクチャを指します。トランスフォーマーモデルでは、この目的は注意機構とフィードフォワード層の組み合わせで自然に達成されます。したがって、本論文では軽量なトランスフォーマー（数層のみ）をパッチミキサーとして使用します。パッチミキサーによって処理された後に、入力シーケンストークンをマスキングします（図2e）。バイナリマスクmを仮定し、次の損失関数を用いてモデルを訓練します。
Mixture-of-experts (MoE) と Layer-wise scalingを導入したTransformerアーキテクチャ
本論文では、計算制約下でモデル性能を向上させるために、先進的なトランスフォーマーアーキテクチャの革新を取り入れています。
- Mixture-of-experts (MoE、Zhou et al., 2022): MoE層を使用することで、モデルのパラメータと表現力を拡張しつつ、訓練コストの大幅な増加を避けます。Expert選択ルーティングを用いた簡素化されたMoE層により、追加の補助損失関数なしで負荷の調整が可能です。
- Layer-wise scaling (Mehta et al., 2024): 大規模言語モデルでの性能向上が示されたアプローチで、トランスフォーマーブロックの幅（隠れ層の次元）を深さに応じて線形に増加させます。深い層には多くのパラメータを割り当て、より複雑な特徴を学習します。
全体アーキテクチャは図２に示されています。
実験
遅延マスキングとパッチミキサーの効果検証
マスキングは、多くのパッチがマスキングされると性能が低下することがあります。Zheng et al. (2024) は、MaskDiT の性能がマスキング比が50%を超えると大幅に低下することを指摘しています。本論文は、最大87.5%のマスキング比で性能を評価し、パッチミキサーを使用しない従来のナイーブマスキング法と比較しました。本論文の「遅延マスキング」は、4層のトランスフォーマーブロックパッチミキサーを使用し、これはバックボーンのトランスフォーマーのパラメータの10%未満です。両者ともAdamWオプティマイザーを同一設定で使用しました。
結果は図３にまとめられています。遅延マスキングがナイーブマスキングとMaskDiTをすべての指標で大きく上回り、マスキング比が高くなるほど性能差が広がることを示しました。例えば、75%のマスキング比では、ナイーブマスキングのFIDスコアは80、MaskDiTは16.5でしたが、提案手法のアプローチは5.03を達成し、マスキングなしのFIDスコア3.79と比較しても優れた結果を示しました。
Mixture-of-expertsと Layer-wise scalingの効果検証
Layer-wise Scaling: DiT-Tinyアーキテクチャを使用した実験で、層ごとのスケーリングと一定幅のトランスフォーマーをナイーブなマスキングで比較しました。両モデルは同じ計算負荷で、同じ期間訓練されました。層ごとのスケーリングアプローチは、すべての性能指標で一定幅モデルを一貫して上回り、マスキング訓練においてより効果的であることが示されました。
Mixture-of-Experts (MoE): MoE層を交互のブロックに持つDiT-Tiny/2トランスフォーマーをテストしました。全体的な性能はMoE層なしのベースラインモデルと類似しており、Clip-scoreはわずかに改善（28.11から28.66へ）、FIDスコアは悪化（6.92から6.98へ）しました。限定的な改善の理由は、60Kステップの訓練と、各専門家が見るサンプル数が少ないためです。
先行研究との比較
COCOデータセットでのゼロショット画像生成（表１）: 30,000枚の生成画像をキャプションから生成し、FID-30Kを用いて実画像との分布を比較しました。提案手法はFID-30Kスコア12.66を達成し、先行の低コスト訓練方法と比較して計算コストが14倍低く、専有データセットにも依存していません。また、Würstchen（Pernias et al., 2024）よりも19倍少ない計算コストで優れた性能を示しました。
詳細な画像生成比較（表２）: GenEval（Ghosh et al., 2024）を使用して、オブジェクトの位置、共起、数、色などの生成能力を評価しました。提案手法は単一オブジェクト生成でほぼ完璧な精度を示し、Stable-Diffusionバリアントと同等、Stable-Diffusion-1.5を上回る性能を発揮。Stable-DiffusionXL-turboやPixArt-αモデルと比較して、色の帰属においても優れた性能を示しました。
まとめ
今回の解説論文ではでは、拡散トランスフォーマーの訓練における計算コスト削減を目指し、パッチマスキング戦略に焦点を当てています。既存のマスキングアプローチの短所を緩和するために「遅延マスキング」戦略を提案し、全てのマスキング比で性能の大幅な改善を示しました。
特に75%の遅延マスキング比を使用し、実画像データセットと合成画像を組み合わせて大規模訓練を実施しました。最先端技術と比べてコストが大幅に低いにもかかわらず、ゼロショット画像生成性能は競争力を持つ結果を達成しました。この低コスト訓練メカニズムが、より多くの研究者による大規模拡散モデルの訓練と開発の参加を促進することが期待されています。
この記事に関するカテゴリー
