aMUSEd: 効率的なテキストから画像への生成

Hugging Face は、aMUSEd という効率的な非拡散テキストから画像へのモデルを導入しました。Google の MUSE のオープン再現として、aMUSEd は Masked Image Modeling (MIM) を活用し、従来の潜在拡散モデルに対するより高速で解釈しやすい代替手段を提供します。

Masked Image Modeling アーキテクチャ

aMUSEd は Masked Image Modeling (MIM) アプローチを利用し、潜在拡散と比較して推論ステップを減らし、解釈性を高めます。モデルアーキテクチャは主に 3 つのコンポーネントから構成されます：トークン化のための VQGAN、プロンプト埋め込みのための CLIP-L/14 テキストエンコーダ、マスクされたパッチを予測する U-ViT モデルです。

学習プロセス

学習時にモデルは以下の手順に従います:

Tokenization: 入力画像は VQGAN を使用して画像トークンに変換されます。
Masking: 画像トークンはコサインマスキングスケジュールに従ってマスクされます。
Prediction: マスクされたトークンは CLIP-L/14 テキストエンコーダからのプロンプト埋め込みで条件付けされ、U-ViT モデルに渡されてマスクされたパッチを予測します。

推論プロセス

推論時にモデルは反復的なプロセスで画像を生成します:

Embedding: 入力プロンプトは CLIP-L/14 テキストエンコーダで埋め込まれます。
Iterative Prediction: プロセスはランダムにマスクされたトークンから開始します。U-ViT モデルがマスクされたトークンを予測し、最も確信度の高い予測のうち一定割合（マスクスケジュールとステップ数 N によって決定）だけが保持されます。残りのトークンは再度マスクされ、U-ViT モデルに戻されます。
Decoding: 最終出力は VQGAN デコーダを通して最終画像が生成されます。

MUSE との主な違い

Single-Stage Prediction: 最終マスクパッチを予測するための二段階アプローチを使用しません。
Text Conditioning: テキスト埋め込みの計算に T5 の代わりに CLIP L/14 を使用します。
Micro-conditioning: SDXL に倣い、画像サイズやクロッピングなどの追加条件付けを組み込んでいます。

能力とパフォーマンス

aMUSEd は効率性と速度を重視して設計されています。テキストエンコーダと VQ-GAN を含めて約 8 億パラメータで、非常に高いパフォーマンスを発揮します。

推論速度とレイテンシ

A100 GPU 上で実施したベンチマークにおいて、aMUSEd は他のモデルに比べて推論レイテンシが大幅に低いことを示しています。この効率性により、デバイス上での利用に適した候補となります。

ゼロショットインペインティング

事前学習目的により、aMUSEd はゼロショット画像インペインティングが可能で、SDXL などのモデルにはない機能です。

スタイル転送

MUSE の機能に倣い、aMUSEd は単一画像を用いたスタイル転送が可能で、パーソナライズされたスタイル特化型画像生成の可能性を提供します。

ファインチューニングとアクセシビリティ

aMUSEd は OpenRAIL ライセンスでリリースされており、商用利用に適しています。Hugging Face はカスタムデータセットでのファインチューニング用トレーニングスクリプトを提供しています。

ファインチューニングのハードウェア要件

Standard Fine-Tuning: 8 ビット Adam オプティマイザと float16 精度を使用すると、ファインチューニングに必要な GPU VRAM は 11GB 未満です。
LoRA: Low-Rank Adaptation (LoRA) を使用すると、メモリ要件はさらに 7GB の VRAM に削減されます。

制限事項

aMUSEd は生の画像品質に関して最先端ではありませんが、コミュニティが画像生成における MIM の可能性を探求することを促すためにリリースされました。著者らは MIM フレームワークのいくつかの利点を強調しています。具体的には、推論効率、デバイス上で使用できる小型モデル、そして高価なファインチューニングなしでタスク転送が可能である点です。