REPA-E: Unlocking VAE for End-to-End Tuning with Latent Diffusion Transformers

In this paper we tackle a fundamental question:“Can we train latent diffusion models together with the variational auto-encoder (VAE) tokenizer in an end-to-end manner?”Traditional deep-learning wisdom dictates that end-to-end training is often preferable when possible. However, for latent diffusion transformers, it is observed that end-toend training both VAE and diffusion-model using standard diffusion-loss is ineffective, even causing a degradation in final performance. We show that while diffusion loss is ineffective, end-to-end training can be unlocked through the representation-alignment (REPA) loss - allowing both VAE and diffusion model to be jointly tuned during the training process. Despite its simplicity, the proposed training recipe (REPA-E) shows remarkable performance; speeding up diffusion model training by over $17 \times$ and $45 \times$ over REPA and vanilla training recipes, respectively. Interestingly, we observe that end-to-end tuning with REPA-E also improves the VAE itself; leading to improved latent space structure and downstream generation performance. In terms of final performance, our approach sets a new state-of-the-art; achieving FID of 1.26 and 1.83 with and without classifierfree guidance on ImageNet $256 \times 256$. Code is available at https://end2end-diffusion.qithub.io.

한 줄 요약

REPA-E는 VAE와 LDM을 함께 엔드투엔드로 훈련하여 생성 성능을 향상시키며, ImageNet 256×256에서 FID 1.26을 달성한다.

핵심 기여도

VAE와 LDM을 함께 엔드투엔드로 훈련하는 방법으로, 기존의 2단계 훈련 방식을 개선.
기존 diffusion loss 대신 representation-alignment (REPA) loss를 사용해 훈련 성능을 17×, 45× 개선.
VAE의 잠재 공간 구조를 자동 개선하여, SD-VAE, IN-VAE 등 다양한 아키텍처에서 생성 성능 향상.
ImageNet 256×256 데이터셋에서 FID 1.26 (classifier-free guidance 사용 시)를 달성하며, 기존 최고 성능을 갱신.

핵심 아이디어

기존의 LDM 훈련은 VAE와 diffusion model을 2단계로 분리하여 훈련하는 방식을 따랐다. 그러나 이는 VAE의 잠재 공간이 diffusion model의 성능 최적화를 고려하지 못하게 만들 수 있다. 본 연구는 VAE와 LDM을 함께 엔드투엔드로 훈련하는 것이 가능할지라는 질문을 제기하고, 기존 diffusion loss 대신 representation-alignment (REPA) loss를 도입한다. REPA loss는 VAE와 diffusion model의 잠재 표현을 정렬하여, 두 모델이 서로 조화롭게 학습할 수 있도록 유도한다. 이는 기존 방식에서 발생하는 성능 저하 문제를 해결하며, VAE의 잠재 공간 구조를 개선하는 부가 효과도 가져온다.

기술적 접근법

**모델 구성**: VAE와 latent diffusion model (LDM)을 함께 엔드투엔드로 훈련.
**손실 함수**: 기존 diffusion loss 대신 representation-alignment (REPA) loss를 사용.
**훈련 방식**: REPA-E는 VAE와 LDM을 동시에 훈련하며, REPA loss를 통해 잠재 표현을 정렬.
**데이터셋**: ImageNet 256×256 사용.
**성능 개선**: REPA-E는 REPA 기반 훈련 대비 17×, vanilla 훈련 대비 45× 훈련 속도 향상.

주요 결과

**ImageNet 256×256**: FID 1.26 (classifier-free guidance 사용 시), 1.83 (사용하지 않을 시) 달성.
**훈련 속도**: REPA-E는 REPA 기반 훈련 대비 17×, vanilla 훈련 대비 45× 빠른 훈련 속도를 보임.
**VAE 성능 개선**: SD-VAE의 경우 고주파 노이즈 감소, IN-VAE의 경우 과부드러운 잠재 공간 개선.
**SiT-XL 아키텍처**: 400K 훈련 스텝에서 FID 4.07 달성 (REPA 기반 4M 스텝 대비 훨씬 빠름).

의의 및 한계

REPA-E는 VAE와 LDM을 엔드투엔드로 훈련함으로써, 기존 2단계 훈련 방식의 한계를 극복하고, 생성 성능과 훈련 효율성을 동시에 향상시킨다. 특히, 다양한 VAE 아키텍처에서 잠재 공간 구조를 자동 개선하는 점이 주목할 만하다. 그러나 본 연구는 ImageNet 256×256에만 적용되었으며, 다른 데이터셋이나 더 큰 모델 아키텍처에서의 일반화 가능성은 추가 실험을 통해 검증이 필요하다. 또한, REPA loss가 왜 VAE와 LDM의 조화를 이끌어내는지에 대한 이론적 분석은 아직 미비한 상태이다.

실용적 활용

REPA-E는 이미지 생성, 생성 모델의 훈련 효율성 향상, 다양한 VAE 아키텍처의 잠재 공간 개선 등에 활용될 수 있다. 특히, 생성 모델의 훈련 시간을 대폭 줄이고, 최종 생성 품질을 향상시키는 데 유용하며, 연구 및 산업 분야에서 빠른 결과 도출이 필요한 상황에 적합하다.