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に達するという驚くべき学習効率を実証している。モデルとコードは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, DeepL