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

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 suboptimal 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 pretrained 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× convergence speedup compared to the original DiT. Models and codes are available at https://github.com/hustvl/LightningDiT.

한 줄 요약

VA-VAE와 LightningDiT를 결합한 시스템이 ImageNet 256 생성에서 FID 1.35 달성하며 21.8× 빠른 수렴 속도를 보인다.

핵심 기여도

VA-VAE(Vision Foundation model Aligned VAE)를 제안하여 고차원 잠재 공간에서 생성 성능을 향상시키며, rFID 0.28의 재구성 능력을 유지.
VF Loss(Vision Foundation Loss) 모듈을 통해 VAE 잠재 공간을 사전 학습된 시각 기초 모델과 정렬.
LightningDiT를 통해 DiT 학습을 2.8× 더 빠르게 수렴.
ImageNet 256 생성에서 FID 1.35 달성, 64 에포크에서 FID 2.11로 21.8× 빠른 수렴.

핵심 아이디어

기존 잠재 확산 모델에서 재구성과 생성 성능 간의 최적화 딜레마는 고차원 잠재 공간 학습의 어려움에서 비롯된다. 본 연구는 이 문제를 해결하기 위해, 고정된 시각 기초 모델(Vision Foundation Model)의 잠재 공간과 정렬된 VAE(VA-VAE)를 제안한다. VA-VAE는 VF Loss를 통해 고차원 잠재 공간의 분포를 더 균일하게 만들고, 생성 성능을 향상시키는 동시에 정보 손실을 최소화한다. 이는 기존 방법에서 고차원 토크나이저를 사용하면 생성 성능이 저하되거나, 모델 크기와 학습 비용이 급증하는 문제를 해결한다.

VF Loss는 요소별(cosine similarity) 및 쌍별(거리 행렬) 유사도를 강제함으로써 전역 및 지역 구조를 모두 정규화한다. 또한 유사도 비용에 마진을 도입하여 과정정화를 방지한다. 이는 고차원 잠재 공간의 자유도를 유지하면서도 생성 성능을 향상시키는 데 기여한다.

기술적 접근법

**VA-VAE**: Vision Foundation Model과 정렬된 VAE. VF Loss를 통해 잠재 공간을 정규화.
**VF Loss**: 요소별(cosine similarity) 및 쌍별(거리 행렬) 유사도를 기반으로 잠재 공간 정규화. 마진 도입으로 과정정화 방지.
**LightningDiT**: DiT 기반의 고성능 확산 트랜스포머. 학습 전략 및 아키텍처 개선으로 2.8× 빠른 수렴.
**ImageNet 256×256 데이터셋** 사용.
**rFID 0.28**의 재구성 능력 유지.
**64 에포크에서 FID 2.11**, **최종 FID 1.35** 달성.

주요 결과

**ImageNet 256×256** 데이터셋에서 **FID 1.35** 달성 (SOTA).
**64 에포크에서 FID 2.11** (기존 DiT 대비 **21.8× 빠른 수렴**).
**VA-VAE**는 **rFID 0.28**의 재구성 능력을 유지하면서 생성 성능 향상.
**VF Loss**를 사용한 VA-VAE는 **2.8× 빠른 DiT 학습**을 가능하게 함.

의의 및 한계

본 연구는 고차원 잠재 공간에서 재구성과 생성 성능 간의 최적화 딜레마를 해결하는 새로운 접근법을 제시한다. VA-VAE와 LightningDiT의 결합은 기존 방법 대비 훨씬 빠른 수렴 속도와 높은 생성 품질을 동시에 달성함으로써, 잠재 확산 모델의 한계를 극복하는 데 기여한다. 특히, VF Loss는 기존 VAE 초기화 방식보다 효과적인 잠재 공간 정규화를 제공하며, 이는 생성 모델의 학습 안정성과 품질 향상에 기여한다.

그러나 본 연구는 특정 데이터셋(ImageNet 256)에서의 성능을 기준으로 평가되었으며, 다른 데이터셋이나 시맨틱 생성 작업에서의 일반화 가능성은 추가 연구가 필요하다. 또한, VF Loss의 마진 파라미터나 정규화 강도는 데이터 종류에 따라 조정이 필요할 수 있다.

실용적 활용

VA-VAE와 LightningDiT는 고해상도 이미지 생성, 텍스트-이미지 생성, 생성 모델의 학습 효율 향상 등에 적용 가능하다. 특히, 생성 모델의 학습 비용을 줄이면서도 높은 품질의 결과를 유지해야 하는 산업 분야(예: 콘텐츠 생성, 디지털 마케팅)에서 유용하게 활용될 수 있다.