Transformadores de Difusão com Autoencoders de Representação

Diffusion Transformers with Representation Autoencoders

October 13, 2025

Autores: Boyang Zheng, Nanye Ma, Shengbang Tong, Saining Xie

cs.AI

Resumo

A modelagem generativa latente, onde um autoencoder pré-treinado mapeia pixels em um espaço latente para o processo de difusão, tornou-se a estratégia padrão para Transformadores de Difusão (DiTs); no entanto, o componente do autoencoder mal evoluiu. A maioria dos DiTs continua a depender do codificador VAE original, o que introduz várias limitações: backbones desatualizados que comprometem a simplicidade arquitetônica, espaços latentes de baixa dimensionalidade que restringem a capacidade de informação e representações fracas resultantes de treinamento puramente baseado em reconstrução, limitando, em última instância, a qualidade gerativa. Neste trabalho, exploramos substituir o VAE por codificadores de representação pré-treinados (por exemplo, DINO, SigLIP, MAE) combinados com decodificadores treinados, formando o que chamamos de Autoencoders de Representação (RAEs). Esses modelos fornecem tanto reconstruções de alta qualidade quanto espaços latentes semanticamente ricos, permitindo uma arquitetura escalável baseada em transformadores. Como esses espaços latentes são tipicamente de alta dimensionalidade, um desafio chave é permitir que os transformadores de difusão operem de forma eficaz dentro deles. Analisamos as fontes dessa dificuldade, propomos soluções teoricamente motivadas e as validamos empiricamente. Nossa abordagem alcança convergência mais rápida sem perdas auxiliares de alinhamento de representação. Usando uma variante de DiT equipada com uma cabeça DDT leve e ampla, alcançamos resultados fortes na geração de imagens no ImageNet: 1.51 FID em 256x256 (sem orientação) e 1.13 tanto em 256x256 quanto em 512x512 (com orientação). O RAE oferece vantagens claras e deve ser o novo padrão para o treinamento de transformadores de difusão.

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.