画像生成AI「Würstchen」公開：42倍圧縮で学習コストを大幅削減

Würstchenは、生成速度と計算効率の限界に挑んだオープンソースのテキスト画像生成モデルだ。SDXLよりも大幅に高速かつ省メモリで動作し、ハイスペックな計算基盤がない環境でも高解像度画像の生成を快適に行える。

42倍という前代未聞の空間圧縮

Stable Diffusionなどの一般的な潜在拡散モデルは、画像を4倍から8倍程度に空間圧縮して計算を行う。通常は16倍を超えると細部の復元が破綻し、まともな画像にならないためだ。

Würstchenはこの限界を破り、42倍という極端な空間圧縮を実現した。1024×1024や1536×1536といった高解像度画像であっても、生成の主要部分は極小の潜在空間上で処理される。

これを可能にしたのが独自の多段階設計だ。VQGANによるStage Aと、拡散オートエンコーダによるStage Bを組み合わせたDecoderが、極限まで縮めた特徴量を高精度にピクセルへ復元する。

プロンプトから画像を生成する主役は、Stage CのPriorだ。Priorは42倍に圧縮された超軽量な空間内だけで拡散処理を行うため、学習と推論の双方で計算コストが劇的に削減される。

学習コストを最大16分の1に削減

この圧縮設計は、学習コストの大幅な低減という形で成果を出している。解像度512×512のWürstchen v1の学習に要した時間はわずか9,000 GPU時間で、SD 1.4の15万時間と比べて16倍も安上がりだ。

1536×1536対応のWürstchen v2でも2万4,602 GPU時間に収まっており、低解像度のSD 1.4と比較しても依然として6倍低コストで学習が完了している。チームによれば、1024×2048といった縦長解像度への適応もスムーズに行えるという。

すでにHugging FaceのDiffusersライブラリに統合されており、普段のコードから簡単に呼び出せる。PyTorch 2.0のSDPA高速化や、メモリ不足を補うCPUオフロード機能、Apple Silicon（MPS）にも対応済みだ。 torch.compileを利用すれば、最初のコンパイルに約2分かかるものの、以降の生成速度をさらに引き上げられる。

なぜ重要か

画像生成AIの開発が高価なGPUクラスターの物量戦になりつつある中、Würstchenは潜在空間の圧縮効率を極限まで高めるというスマートな解決策を示した。

個人の開発者や中小規模のチームにとって、限られたハードウェア資源で高解像度モデルの実験やファインチューニングを行える価値は極めて大きい。

ただし、42倍という極端な圧縮を行う都合上、微細な文字や複雑なテクスチャの復元精度においてどこまで既存モデルと張り合えるかは、実際の運用を通じて見極める必要があると考える。