Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models

要約

Transformer アーキテクチャを備えた潜在拡散モデルは、高忠実度の画像の生成に優れています。
ただし、最近の研究では、この 2 段階の設計における最適化のジレンマが明らかになりました。ビジュアル トークナイザーでトークンごとの特徴次元を増やすと再構成の品質が向上しますが、同等の生成パフォーマンスを達成するには、大幅に大規模な拡散モデルとより多くのトレーニング反復が必要になります。
その結果、既存のシステムは、トークナイザー内の情報損失により視覚的なアーティファクトが生成されたり、高価な計算コストにより完全に収束できなかったりして、次善のソリューションに落ち着くことがよくあります。
私たちは、このジレンマは、制約のない高次元の潜在空間を学習することの本質的な困難に起因すると主張します。
これに対処するために、ビジュアルトークナイザーをトレーニングするときに、潜在空間を事前トレーニングされたビジョン基盤モデルと調整することを提案します。
私たちが提案する VA-VAE (Vision Foundation Model Aligned variational AutoEncoder) は、潜在拡散モデルの再構成生成フロンティアを大幅に拡張し、高次元潜在空間における拡散変換器 (DiT) のより高速な収束を可能にします。
VA-VAE の可能性を最大限に活用するために、LightningDiT と呼ばれる、改善されたトレーニング戦略とアーキテクチャ設計を備えた強化された DiT ベースラインを構築します。
この統合システムは、ImageNet 256×256 世代で FID スコア 1.35 の最先端 (SOTA) パフォーマンスを達成するとともに、わずか 64 エポックで FID スコア 2.11 に達することで顕著なトレーニング効率を実証し、これは 21 倍を超える収束速度の向上を示しています。
オリジナルのDiTと比較してください。
モデルとコードは https://github.com/hustvl/LightningDiT から入手できます。

要約(オリジナル)

Latent diffusion models with Transformer architectures excel at generating high-fidelity images. However, recent studies reveal an optimization dilemma in this two-stage design: while increasing the per-token feature dimension in visual tokenizers improves reconstruction quality, it requires substantially larger diffusion models and more training iterations to achieve comparable generation performance. Consequently, existing systems often settle for sub-optimal solutions, either producing visual artifacts due to information loss within tokenizers or failing to converge fully due to expensive computation costs. We argue that this dilemma stems from the inherent difficulty in learning unconstrained high-dimensional latent spaces. To address this, we propose aligning the latent space with pre-trained vision foundation models when training the visual tokenizers. Our proposed VA-VAE (Vision foundation model Aligned Variational AutoEncoder) significantly expands the reconstruction-generation frontier of latent diffusion models, enabling faster convergence of Diffusion Transformers (DiT) in high-dimensional latent spaces. To exploit the full potential of VA-VAE, we build an enhanced DiT baseline with improved training strategies and architecture designs, termed LightningDiT. The integrated system achieves state-of-the-art (SOTA) performance on ImageNet 256×256 generation with an FID score of 1.35 while demonstrating remarkable training efficiency by reaching an FID score of 2.11 in just 64 epochs–representing an over 21 times convergence speedup compared to the original DiT. Models and codes are available at: https://github.com/hustvl/LightningDiT.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google