REPA-E: 잠재 확산 트랜스포머를 통한 종단 간 튜닝을 위한 VAE의 잠재력 해제

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

April 14, 2025

저자: Xingjian Leng, Jaskirat Singh, Yunzhong Hou, Zhenchang Xing, Saining Xie, Liang Zheng

cs.AI

초록

본 논문에서 우리는 근본적인 질문에 대해 다룹니다: "잠재 확산 모델(Latent Diffusion Model)을 변분 오토인코더(VAE) 토크나이저와 함께 종단 간(end-to-end) 방식으로 학습시킬 수 있을까?" 전통적인 딥러닝 지혜에 따르면, 가능한 경우 종단 간 학습이 종종 선호됩니다. 그러나 잠재 확산 트랜스포머(Latent Diffusion Transformer)의 경우, 표준 확산 손실을 사용하여 VAE와 확산 모델을 동시에 종단 간 학습시키는 것은 비효율적이며, 최종 성능 저하를 초래하는 것으로 관찰되었습니다. 우리는 확산 손실이 비효율적이지만, 표현 정렬(REPA) 손실을 통해 종단 간 학습이 가능해짐을 보여줍니다. 이를 통해 학습 과정에서 VAE와 확산 모델을 함께 조정할 수 있습니다. 제안된 학습 방법(REPA-E)은 단순함에도 불구하고 놀라운 성능을 보입니다. REPA와 기본 학습 방법에 비해 각각 17배와 45배 이상 빠른 확산 모델 학습 속도를 달성합니다. 흥미롭게도, REPA-E를 통한 종단 간 조정은 VAE 자체도 개선시키는 것으로 관찰되었습니다. 이는 잠재 공간 구조와 하류 생성 성능의 향상으로 이어집니다. 최종 성능 측면에서, 우리의 접근 방식은 새로운 최첨단 기술을 제시합니다. ImageNet 256 x 256에서 분류자 없는 지도(classifier-free guidance)를 사용한 경우와 사용하지 않은 경우 각각 1.26과 1.83의 FID(Fréchet Inception Distance)를 달성했습니다. 코드는 https://end2end-diffusion.github.io에서 확인할 수 있습니다.

English

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-to-end 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 17x and 45x 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 classifier-free guidance on ImageNet 256 x 256. Code is available at https://end2end-diffusion.github.io.