표현 오토인코더를 활용한 확산 트랜스포머

초록

잠재 생성 모델링, 즉 사전 학습된 오토인코더가 픽셀을 확산 과정을 위한 잠재 공간으로 매핑하는 방식은 Diffusion Transformers(DiT)의 표준 전략으로 자리 잡았으나, 오토인코더 구성 요소는 거의 진화하지 못했다. 대부분의 DiT는 여전히 원래의 VAE 인코더에 의존하고 있으며, 이는 몇 가지 한계를 초래한다: 아키텍처 단순성을 저해하는 구식 백본, 정보 용량을 제한하는 저차원 잠재 공간, 순수 재구축 기반 학습으로 인한 약한 표현력, 그리고 궁극적으로 생성 품질을 제한하는 문제 등이 있다. 본 연구에서는 VAE를 사전 학습된 표현 인코더(예: DINO, SigLIP, MAE)와 학습된 디코더를 결합한 Representation Autoencoders(RAEs)로 대체하는 방안을 탐구한다. 이러한 모델은 고품질 재구축과 의미론적으로 풍부한 잠재 공간을 제공하며, 확장 가능한 트랜스포머 기반 아키텍처를 가능하게 한다. 이러한 잠재 공간이 일반적으로 고차원이기 때문에, 확산 트랜스포머가 이를 효과적으로 운영할 수 있도록 하는 것이 주요 과제이다. 우리는 이러한 어려움의 원인을 분석하고, 이론적으로 타당한 해결책을 제안하며, 이를 실증적으로 검증한다. 우리의 접근 방식은 보조 표현 정렬 손실 없이도 더 빠른 수렴을 달성한다. 경량화된 넓은 DDT 헤드를 장착한 DiT 변형을 사용하여, 우리는 ImageNet에서 강력한 이미지 생성 결과를 얻었다: 256x256에서 1.51 FID(지도 없음) 및 256x256과 512x512에서 1.13(지도 있음). RAE는 명확한 이점을 제공하며, 확산 트랜스포머 학습의 새로운 기본으로 자리 잡아야 한다.

English

Latent generative modeling, where a pretrained autoencoder maps pixels into a latent space for the diffusion process, has become the standard strategy for Diffusion Transformers (DiT); however, the autoencoder component has barely evolved. Most DiTs continue to rely on the original VAE encoder, which introduces several limitations: outdated backbones that compromise architectural simplicity, low-dimensional latent spaces that restrict information capacity, and weak representations that result from purely reconstruction-based training and ultimately limit generative quality. In this work, we explore replacing the VAE with pretrained representation encoders (e.g., DINO, SigLIP, MAE) paired with trained decoders, forming what we term Representation Autoencoders (RAEs). These models provide both high-quality reconstructions and semantically rich latent spaces, while allowing for a scalable transformer-based architecture. Since these latent spaces are typically high-dimensional, a key challenge is enabling diffusion transformers to operate effectively within them. We analyze the sources of this difficulty, propose theoretically motivated solutions, and validate them empirically. Our approach achieves faster convergence without auxiliary representation alignment losses. Using a DiT variant equipped with a lightweight, wide DDT head, we achieve strong image generation results on ImageNet: 1.51 FID at 256x256 (no guidance) and 1.13 at both 256x256 and 512x512 (with guidance). RAE offers clear advantages and should be the new default for diffusion transformer training.

표현 오토인코더를 활용한 확산 트랜스포머

Diffusion Transformers with Representation Autoencoders

초록

Support