Diffusers との Stable Diffusion 3 Medium 統合
Hugging Face は、Stability AI の 2B パラメータモデルである Stable Diffusion 3 Medium を diffusers ライブラリに統合しました。このリリースにより、ユーザーは Multimodal Diffusion Transformer (MMDiT) と整流フローマッチング訓練目的と呼ばれる新しいアーキテクチャ手法を使用して、高品質なテキストから画像への合成および画像から画像への合成を実行できます。
Stable Diffusion 3 のアーキテクチャ的イノベーション
Stable Diffusion 3 (SD3) は、テキストと画像データ間の双方向情報フローを利用することで、従来のテキストから画像へのアーキテクチャとは異なります。
マルチモーダル拡散トランスフォーマー (MMDiT)
SD3 は、CLIP L/14、OpenCLIP bigG/14、T5-v1.1-XXL の 3 つのテキストエンコーダと共に新しい MMDiT モデルを採用しています。また、Stable Diffusion XL に見られるものと類似した 16 チャネルの AutoEncoder も使用します。
従来のモデルが固定されたテキスト表現でクロスアテンションを使用していたのとは異なり、MMDiT ブロックはテキスト入力とピクセル潜在表現を埋め込みシーケンスとして処理します。これらのシーケンスは別々の重みセットを用いて共通の次元に埋め込まれ、連結され、モジュレートされたアテンションと MLP を通過します。これにより、注意操作の間に両モダリティが相互に影響し合うことが可能になります。
整流フローマッチング
SD3 は条件付きフローマッチング目的で訓練されます。この手法は、前方ノイズ付加プロセスをデータ分布とノイズ分布を直線で結ぶ「整流フロー」と定義します。
これをサポートするために、Hugging Face は FlowMatchEulerDiscreteScheduler を導入しました。このスケジューラはオイラー法ステップを実装し、解像度依存のタイムステップスケジュールシフト用に shift パラメータを含みます。2B モデルでは、より高解像度でのノイズスケーリングを適切に処理するために shift=3.0 が推奨されます。
メモリ最適化手法
SD3 は大規模な T5-XXL テキストエンコーダ（4.7B パラメータ）を使用しているため、VRAM が 24GB 未満の GPU でモデルを実行することは困難です。Hugging Face はメモリ使用量を削減するためのいくつかの最適化戦略を提供しています：
- モデルオフロード:
pipe.enable_model_cpu_offload()を使用すると、使用していないときにモデルコンポーネントが CPU に移動し、VRAM 使用量が削減されますが、レイテンシが増加します。 - T5 エンコーダの除外: 推論時に T5-XXL エンコーダを削除 (
text_encoder_3=None) すると、メモリ要件が大幅に削減され、性能のわずかな低下が生じます。 - 8 ビット量子化:
bitsandbytesライブラリを介して T5-XXL モデルを 8 ビット精度でロードすることで、メモリ消費がさらに低減します。
メモリベンチマーク比較
fp16 精度と PyTorch 2.3 を使用して A100 GPU（80GB VRAM）で実施したベンチマークは、以下のトレードオフを示しています：
| 手法 | 推論時間 (秒) | メモリ (GB) |
|---|---|---|
| Default | 4.762 | 18.765 |
| Offloading | 32.765 | 12.0645 |
| Offloading + no T5 | 19.110 | 4.266 |
| 8bit T5 | 4.932 | 10.586 |
パフォーマンスとファインチューニング
torch.compile() を用いた推論高速化
ユーザーは torch.compile() を使用して VAE とトランスフォーマーコンポーネントの計算グラフを最適化することで、顕著な速度向上を実現できます。80GB の A100 1 台でのベンチマークでは、平均推論時間が 0.585 秒に短縮され、イーガー実行に比べて 4 倍の高速化を示しました。
DreamBooth と LoRA トレーニング
Hugging Face は、Low-Rank Adaptation (LoRA) を活用した SD3 用 DreamBooth ファインチューニングスクリプトをリリースしました。このスクリプトは効率的なモデルカスタマイズを可能にし、整流フローに基づくトレーニングパイプライン実装の参考となります。