PRX Part 3: 24時間でテキストから画像へのモデルをトレーニング

Hugging Face と Photoroom は、32 台の H200 GPU と約 1,500 ドルの計算予算で、テキストから画像への実用的なモデルをわずか 24 時間でスクラッチから学習できることを実証しました。最新のアーキテクチャとトレーニング最適化を組み合わせることで、これを達成しました。

ピクセル空間でのトレーニングと X‑Prediction

このモデルは x‑prediction 方式を採用しており、ピクセル空間で直接トレーニングできるため、Variational Autoencoder (VAE) が不要です。シーケンス長を計算上管理しやすくするため、初期トークン投影層ではパッチサイズ 32 と 256 次元のボトルネックを使用しました。

  • 512px 解像度ではシーケンス長は 256 トークン
  • 1024px ではシーケンス長が 1,024 トークンに増加

トレーニングは 512px から開始し、最終的に 1024px でファインチューニングしました。

視覚品質向上のための知覚損失

ピクセルを直接予測することで、古典的なコンピュータビジョンの知覚損失を利用でき、収束速度と視覚品質が向上します。標準の flow matching 目的に加えて、以下の 2 つの補助損失を実装しました。

  • LPIPS：重み 0.1 で低レベルの知覚的類似性を捉える。
  • DINO ベースの知覚損失：DINOv2 を使用し、重み 0.01 でより強い意味的シグナルを提供。

これらの損失は、パッチ単位の特徴ではなく、すべてのノイズレベルにわたるプールされた全画像に対して適用し、実証的に優れた結果を得ました。

TREAD による効率的なトークンルーティング

ステップごとの計算コスト削減のため、チームは TREAD（Token Routing for Efficient Architecture‑agnostic Diffusion Training）を採用しました。この手法は、トランスフォーマーブロックの連続したチャンクをバイパスするトークンの一部をランダムに選択し、後で再注入します。

  • トークンの 50% を第 2 ブロックから最終ブロック手前までルーティング。
  • バニラの Classifier‑Free Guidance (CFG) での品質低下を防ぐため、密な条件予測とルーティングされた条件予測を組み合わせた自己ガイダンス方式を実装。

表現アラインメントとオプティマイザ

表現アラインメントは REPA を用い、教師モデルとして DINOv3 を使用しました。アラインメント損失は 8 番目のトランスフォーマーブロックで 1 回だけ適用し、重みは 0.5 としました。TREAD のルーティングと整合させるため、損失はルーティングされていないトークンのみに計算されます。

最適化には、2 次元パラメータ（行列）に Muon、それ以外のパラメータ（バイアス、ノルム、埋め込み）に Adam を使用しました。

Optimizer Group Target Parameters Key Hyperparameters
Muon 2D parameters lr=1e-4, momentum=0.95, nesterov=true, ns_steps=5
Adam Non-2D parameters lr=1e-4, betas=(0.9, 0.95), eps=1e-8

トレーニングデータとスケジュール

以下の 3 つの合成データセットでトレーニングを実施しました。

  1. Flux generated（1.7M 画像）
  2. FLUX-Reason-6M（6M 画像）
  3. midjourney-v6-llava（1M 画像）— 一貫性のため Gemini 1.5 で再キャプション。

トレーニングスケジュールは、512px でバッチサイズ 1024 の 100k ステップ、その後 1024px でバッチサイズ 512 の 20k ステップ（REPA なし）です。

結果と示唆

最終モデルは一部テクスチャの乱れや稀な解剖学的エラーが見られるものの、"明らかに実用的" であり、プロンプトへの追従性と一貫した美学が高く評価されています。著者らは、残る課題はレシピの構造的欠陥ではなく、トレーニング不足とデータ多様性の制限によるものと結論付けています。

Photoroom は、トレーニングコードと実験フレームワークを PRX GitHub リポジトリ でオープンソース化しています。

Sources