Escalonando Transformadores de Difusão Texto-para-Imagem com Autoencoders de Representação
Scaling Text-to-Image Diffusion Transformers with Representation Autoencoders
January 22, 2026
Autores: Shengbang Tong, Boyang Zheng, Ziteng Wang, Bingda Tang, Nanye Ma, Ellis Brown, Jihan Yang, Rob Fergus, Yann LeCun, Saining Xie
cs.AI
Resumo
Os Autoencoders de Representação (RAEs) demonstraram vantagens distintas na modelagem de difusão no ImageNet ao treinar em espaços latentes semânticos de alta dimensão. Neste trabalho, investigamos se esta estrutura pode ser dimensionada para a geração de texto para imagem (T2I) em larga escala e de forma livre. Primeiro, dimensionamos os decodificadores RAE no codificador de representação congelado (SigLIP-2) para além do ImageNet, treinando com dados da web, sintéticos e de renderização de texto, concluindo que, embora a escala melhore a fidelidade geral, a composição de dados direcionada é essencial para domínios específicos como texto. Em seguida, submetemos rigorosamente a teste as opções de design do RAE originalmente propostas para o ImageNet. Nossa análise revela que o dimensionamento simplifica a estrutura: embora o agendamento de ruído dependente da dimensão permaneça crítico, complexidades arquiteturais como cabeças de difusão amplas e decodificação com aumento de ruído oferecem benefícios negligenciáveis em escala. Com base nesta estrutura simplificada, realizamos uma comparação controlada do RAE contra o FLUX VAE, o estado da arte, através de escalas de transformadores de difusão de 0,5B a 9,8B de parâmetros. Os RAEs superam consistentemente os VAEs durante o pré-treinamento em todas as escalas de modelo. Além disso, durante o ajuste fino em conjuntos de dados de alta qualidade, os modelos baseados em VAE sofrem sobreajuste catastrófico após 64 épocas, enquanto os modelos RAE permanecem estáveis por 256 épocas e alcançam desempenho consistentemente melhor. Em todos os experimentos, os modelos de difusão baseados em RAE demonstram convergência mais rápida e melhor qualidade de geração, estabelecendo os RAEs como uma base mais simples e robusta do que os VAEs para a geração de T2I em larga escala. Adicionalmente, como tanto a compreensão quanto a geração visual podem operar em um espaço de representação compartilhado, o modelo multimodal pode raciocinar diretamente sobre os latentes gerados, abrindo novas possibilidades para modelos unificados.
English
Representation Autoencoders (RAEs) have shown distinct advantages in diffusion modeling on ImageNet by training in high-dimensional semantic latent spaces. In this work, we investigate whether this framework can scale to large-scale, freeform text-to-image (T2I) generation. We first scale RAE decoders on the frozen representation encoder (SigLIP-2) beyond ImageNet by training on web, synthetic, and text-rendering data, finding that while scale improves general fidelity, targeted data composition is essential for specific domains like text. We then rigorously stress-test the RAE design choices originally proposed for ImageNet. Our analysis reveals that scaling simplifies the framework: while dimension-dependent noise scheduling remains critical, architectural complexities such as wide diffusion heads and noise-augmented decoding offer negligible benefits at scale Building on this simplified framework, we conduct a controlled comparison of RAE against the state-of-the-art FLUX VAE across diffusion transformer scales from 0.5B to 9.8B parameters. RAEs consistently outperform VAEs during pretraining across all model scales. Further, during finetuning on high-quality datasets, VAE-based models catastrophically overfit after 64 epochs, while RAE models remain stable through 256 epochs and achieve consistently better performance. Across all experiments, RAE-based diffusion models demonstrate faster convergence and better generation quality, establishing RAEs as a simpler and stronger foundation than VAEs for large-scale T2I generation. Additionally, because both visual understanding and generation can operate in a shared representation space, the multimodal model can directly reason over generated latents, opening new possibilities for unified models.