Transformadores de Difusión con Autoencoders de Representación

Diffusion Transformers with Representation Autoencoders

October 13, 2025

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

cs.AI

Resumen

El modelado generativo latente, donde un autocodificador preentrenado mapea píxeles en un espacio latente para el proceso de difusión, se ha convertido en la estrategia estándar para los Transformadores de Difusión (DiT); sin embargo, el componente del autocodificador apenas ha evolucionado. La mayoría de los DiT siguen dependiendo del codificador VAE original, lo que introduce varias limitaciones: arquitecturas obsoletas que comprometen la simplicidad estructural, espacios latentes de baja dimensión que restringen la capacidad de información y representaciones débiles que resultan de un entrenamiento basado únicamente en la reconstrucción, lo que finalmente limita la calidad generativa. En este trabajo, exploramos reemplazar el VAE con codificadores de representación preentrenados (por ejemplo, DINO, SigLIP, MAE) combinados con decodificadores entrenados, formando lo que denominamos Autocodificadores de Representación (RAE). Estos modelos proporcionan tanto reconstrucciones de alta calidad como espacios latentes semánticamente ricos, al mismo tiempo que permiten una arquitectura escalable basada en transformadores. Dado que estos espacios latentes suelen ser de alta dimensión, un desafío clave es permitir que los transformadores de difusión operen eficazmente dentro de ellos. Analizamos las fuentes de esta dificultad, proponemos soluciones teóricamente fundamentadas y las validamos empíricamente. Nuestro enfoque logra una convergencia más rápida sin pérdidas auxiliares de alineación de representaciones. Utilizando una variante de DiT equipada con una cabeza DDT ligera y amplia, obtenemos resultados sólidos en la generación de imágenes en ImageNet: 1.51 FID a 256x256 (sin guía) y 1.13 tanto a 256x256 como a 512x512 (con guía). RAE ofrece ventajas claras y debería ser el nuevo estándar para el entrenamiento de transformadores de difusión.

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.