Segmind SD-Small および SD-Tiny 知識蒸留リリース

TL;DR

Segmind は、ブロック除去知識蒸留（block-removal knowledge distillation）によって圧縮された拡散モデルである SD-Small と SD-Tiny のトレーニングコードと学習済みチェックポイントを公開しました。これらはベースモデルと同等の画像忠実度を維持しながら、パラメータ数を 35%～55% 削減し、最大 100% 高速な推論を実現しています。

知識蒸留のメソドロジー

公開されたモデルは、On Architectural Compression of Text-to-Image Diffusion Models (Shinkook et al.) で説明されている Block-Removal Knowledge Distillation 技術を使用してトレーニングされています。

  • 教師モデル: Realistic-Vision 4.0（高品質な Stable Diffusion チェックポイント）
  • 生徒アーキテクチャ: レイヤーを除去した UNet バリアント。パラメータ数を 35% (SD-Small) および 55% (SD-Tiny) 削減。
  • 損失関数（Loss composition）の構成:
    1. ターゲット画像と生成画像の潜在表現（latent representations）間の標準的な拡散損失。
    2. 生徒が生成した潜在変数と教師が生成した潜在変数を一致させる潜在レベルの損失。
    3. 教師と生徒の各 UNet ブロックの出力を一致させる特徴レベルの損失（最も重要なコンポーネント）。
  • トレーニングデータ: 画像スコア > 7.5 でフィルタリングされた LAION Art Aesthetic データセット。
  • トレーニングスケジュール: 100k ステップ（SD-Small）および 125k ステップ（SD-Tiny）で 1M 枚の画像を使用。

完全な蒸留パイプラインは [segmind/distill-sd](https://github.com/segmind/distill-sd) リポジトリで利用可能であり、学習済みチェックポイントは Hugging Face の segmind ネームスペースでホストされています。

🤗 Diffusers によるモデルの使用方法

両モデルとも、🤗 Diffusers ライブラリの DiffusionPipeline を使用して直接ロードできます。

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained(
    "segmind/small-sd", torch_dtype=torch.float16
)
prompt = "Portrait of a pretty girl"
negative_prompt = (
    "(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, "
    "cartoon, drawing, anime:1.4), text, close up, cropped, out of frame, "
    "worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, "
    "mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn "
    "face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, "
    "extra limbs, cloned face, disfigured, gross proportions, malformed limbs, "
    "missing arms, missing legs, extra arms, extra legs, fused fingers, too many "
    "fingers, long neck"
)
image = pipeline(prompt, negative_prompt=negative_prompt).images[0]
image.save("my_image.png")

モデル識別子を入れ替えることで、segmind/tiny-sd でも同じ API が動作します。

推論速度の向上

同一ハードウェアでのベンチマークにより、蒸留されたモデルは元のベースチェックポイントと比較して 最大 2 倍低いレイテンシ を示しています。これらの測定に使用された推論スクリプトは、リポジトリ (inference.py) に含まれています。

既知の制限事項

  • これらのモデルは 初期段階のリリース であり、視覚的な品質はまだプロダクショングレードの拡散モデルには及ばない可能性があります。
  • 汎用的な生成には最適化されておらず、複雑な構成のプロンプトや複数の概念の扱いに苦戦する場合があります。
  • より高い忠実度を得るためのベストプラクティスは、ドメイン固有のデータで ファインチューニングまたは LoRA を適用すること です。

ポートレートデータセットでの SD-Tiny のファインチューニング

Segmind は、Realistic-Vision 4.0 で生成された 7k 枚のポートレート画像セットを用いた SD-Tiny のファインチューニングを実証しました。トレーニングハイパーパラメータ:

  • Steps: 131,000
  • Learning rate: 1e-4
  • Batch size: 32 (gradient accumulation = 4)
  • Image resolution: 768 px
  • Mixed-precision: fp16

結果として得られたサンプルは、55% のパラメータ削減を維持しつつ、元の教師モデルの品質に近づいています。

蒸留モデルへの LoRA トレーニング

SD-Tiny に Low-Rank Adaptation (LoRA) を適用すると、モデルサイズが小さいため LoRA の収束が速くなります。抽象的な概念でトレーニングされた LoRA チェックポイントの例がリポジトリ (lora_training.py) に用意されています。

コミュニティへの招待

Segmind は、開発者がプロジェクトに貢献し、問題を報告し、ファインチューニングされたチェックポイントを共有することを奨励しています。コミュニケーションチャネルには Discord サーバーと GitHub リポジトリがあり、スターやプルリクエストを歓迎します。

Sources