TextLDM: Modelagem de Linguagem com Difusão Latente Contínua

Resumo

Transformadores de Difusão (DiT) treinados com correspondência de fluxo em um espaço latente de VAE unificaram a geração visual em imagens e vídeos. Um passo natural em direção a uma arquitetura única tanto para geração (síntese visual) quanto para compreensão (geração de texto) é aplicar esse framework à modelagem de linguagem. Propomos o TextLDM, que transfere a receita de difusão visual latente para a geração de texto com modificação arquitetural mínima. Um VAE baseado em Transformer mapeia tokens discretos para latentes contínuos, aprimorado pelo Alinhamento de Representações (REPA) com um modelo de linguagem pré-treinado congelado para produzir representações eficazes para remoção de ruído condicional. Um DiT padrão então realiza correspondência de fluxo nesse espaço latente, idêntico em arquitetura à sua contraparte visual. O desafio central que abordamos é obter representações contínuas de texto de alta qualidade: descobrimos que a fidelidade de reconstrução sozinha é insuficiente, e que alinhar características latentes com um modelo de linguagem pré-treinado via REPA é crucial para a qualidade da geração downstream. Treinado do zero no OpenWebText2, o TextLDM supera substancialmente modelos anteriores de linguagem por difusão e iguala o GPT-2 sob as mesmas configurações. Nossos resultados estabelecem que a receita do DiT visual se transfere eficazmente para a linguagem, dando um passo concreto em direção a arquiteturas de difusão unificadas para geração e compreensão multimodal.

English

Diffusion Transformers (DiT) trained with flow matching in a VAE latent space have unified visual generation across images and videos. A natural next step toward a single architecture for both generation (visual synthesis) and understanding (text generation) is to apply this framework to language modeling. We propose TextLDM, which transfers the visual latent diffusion recipe to text generation with minimal architectural modification. A Transformer-based VAE maps discrete tokens to continuous latents, enhanced by Representation Alignment (REPA) with a frozen pretrained language model to produce representations effective for conditional denoising. A standard DiT then performs flow matching in this latent space, identical in architecture to its visual counterpart. The central challenge we address is obtaining high-quality continuous text representations: we find that reconstruction fidelity alone is insufficient, and that aligning latent features with a pretrained language model via REPA is critical for downstream generation quality. Trained from scratch on OpenWebText2, TextLDM substantially outperforms prior diffusion language models and matches GPT-2 under the same settings. Our results establish that the visual DiT recipe transfers effectively to language, taking a concrete step toward unified diffusion architectures for multimodal generation and understanding.