最近、カリフォルニア大学アーバイン校などの研究者たちが、遅延マスキング、MoE、階層的拡張などの戦略を用いて、拡散モデルの訓練コストを1890ドルまで削減しました。これは、以前の最も安価な方法（Wuerstchen）の28400ドルよりも一桁低く、Stable Diffusionのようなモデルのコストはさらに高いです。
研究者たちは、この膨大な費用を削減するためにさまざまな方法を試みました。例えば、元の拡散モデルではノイズから画像まで約1000ステップ必要でしたが、現在では20ステップ程度、あるいはそれ以下に減少しています。拡散モデルの基本モジュールがUnet（CNN）からDiT（Transformer）に徐々に置き換えられた後、Transformerの特性に基づいた最適化も追随しました。例えば、量子化、Attentionの冗長な計算のスキップ、パイプラインなどです。
この研究では、11.6億パラメータの拡散モデルを一から訓練するのに、わずか1890ドルしかかかりません。これは最先端技術と比較して一桁の改善であり、一般の人々にも事前訓練の可能性を示しています。さらに重要なのは、コスト削減技術がモデルのパフォーマンスに影響を与えていないことです。11.6億のパラメータで非常に良好な結果が得られました。視覚的な印象だけでなく、モデルのデータ指標も優れており、例えばFIDスコアはStable Diffusion 1.5やDALL·E 2に非常に近いです。
コスト削減の秘訣は主に以下を含みます：
-
遅延マスキング戦略：マスク前にパッチミキサーを使用して前処理を行い、破棄されたパッチの情報を生き残ったパッチに埋め込むことで、高マスクによるパフォーマンス低下を大幅に減少させます。
-
ファインチューニング：事前訓練（マスク）後に小規模なファインチューニング（アンマスク）を行い、マスク使用による不適切な生成アーティファクトを軽減します。
-
MoEと階層的拡張：専門家選択ルーティングに基づく簡略化されたMoE層を使用し、訓練コストを大幅に増加させることなくモデルのパラメータと表現能力を増加させます。同時に、Transformerブロックの幅を線形に増加させる階層的スケーリング方法も考慮しています。
実験設定では、著者らはDiTの2つの変種：DiT-Tiny/2とDiT-Xl/2を使用し、パッチサイズは2です。すべてのモデルは、コサイン学習率減衰と高い重み減衰を持つAdamW最適化器で訓練されました。モデルのフロントエンドには、Stable-Diffusion-XLモデルの4チャンネル変分自己エンコーダー（VAE）を使用して画像特徴を抽出し、さらに最新の16チャンネルVAEの大規模訓練（コスト削減版）でのパフォーマンスもテストしました。
著者らは、すべての拡散モデルの統一訓練設定としてEDMフレームワークを使用し、FIDとCLIPスコアを用いて画像生成モデルのパフォーマンスを測定しました。テキストエンコーダーには最も一般的なCLIPモデルを選択しました。
訓練データセットには、3つの実際の画像データセット（Conceptual Captions、Segment Anything、TextCaps）を使用し、2200万の画像テキストペアが含まれています。SA1Bは実際のキャプションを提供していないため、ここではLLaVAモデルで生成した合成キャプションを使用しています。