Диффузионные трансформаторы с автоэнкодерами представлений

Diffusion Transformers with Representation Autoencoders

October 13, 2025

Авторы: Boyang Zheng, Nanye Ma, Shengbang Tong, Saining Xie

cs.AI

Аннотация

Латентное генеративное моделирование, при котором предобученный автокодировщик отображает пиксели в латентное пространство для процесса диффузии, стало стандартной стратегией для диффузионных трансформеров (DiT). Однако компонент автокодировщика практически не эволюционировал. Большинство DiT продолжают полагаться на оригинальный VAE-кодировщик, который вводит несколько ограничений: устаревшие архитектурные основы, компрометирующие простоту архитектуры, низкоразмерные латентные пространства, ограничивающие информационную емкость, и слабые представления, возникающие из-за чисто реконструкционного обучения, что в конечном итоге ограничивает качество генерации. В данной работе мы исследуем замену VAE на предобученные кодировщики представлений (например, DINO, SigLIP, MAE), объединенные с обученными декодерами, формируя то, что мы называем Автокодировщиками Представлений (RAE). Эти модели обеспечивают как высококачественную реконструкцию, так и семантически богатые латентные пространства, позволяя использовать масштабируемую архитектуру на основе трансформеров. Поскольку эти латентные пространства обычно высокоразмерны, ключевой задачей является обеспечение эффективной работы диффузионных трансформеров в них. Мы анализируем источники этой сложности, предлагаем теоретически обоснованные решения и подтверждаем их эмпирически. Наш подход достигает более быстрой сходимости без вспомогательных потерь на выравнивание представлений. Используя вариант DiT, оснащенный легковесной и широкой DDT-головой, мы достигаем сильных результатов генерации изображений на ImageNet: 1.51 FID при разрешении 256x256 (без управления) и 1.13 как при 256x256, так и при 512x512 (с управлением). 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.