Resumo Técnico: AURORA-LM

Definição do Problema

A geração de linguagem permanece um ponto fora da curva na modelagem generativa moderna. Enquanto imagens, vídeo e áudio são cada vez mais modelados em espaços latentes contínuos usando processos de difusão, a geração de texto ainda depende predominantemente de tokens discretos. As tentativas existentes de aplicar difusão contínua à linguagem enfrentam um compromisso de design fundamental:

1. Representações Herdadas: Métodos como o Diffusion-LM ou TEncDM herdam espaços de incorporação (embeddings) de embeddings de tokens ou codificadores pré-treinados. Esses espaços não são explicitamente projetados para os requisitos conjuntos de reconstrução de texto precisa e geração contínua eficaz.
2. Compromissos de Compressão: Abordagens baseadas em autoencoders (ex: LD4LG, COSMOS) mapeiam o texto para latentes contínuos, mas frequentemente comprimem essas representações para tornar a distribuição de difusão mais fácil de aprender. Essa compressão sacrifica a fidelidade ao nível do token, criando um gargalo onde a melhoria da tratabilidade generativa degrada a precisão da decodificação, e vice-versa.

O desafio central é construir uma representação de texto contínua que suporte uma decodificação fiel e precisa, enquanto se projeta um modelo generativo capaz de aprender sua distribuição sem sacrificar essa capacidade representacional.

Metodologia: Estrutura AURORA-LM

O AURORA-LM (Autoencoding Unified Representation for Continuous-Latent Diffusion Language Modeling) aborda esses desafios desacoplando a construção da representação da modelagem da distribuição em duas etapas distintas.

1. Construção de Latente de Texto Contínuo (O Autoencoder)

Em vez de herdar representações, o AURORA-LM constrói explicitamente um latente de texto de alta capacidade e decodificável usando um Codificador-Decodificador baseado em Query:

• Codificador baseado em Query: Organiza o texto discreto em uma sequência latente de alta capacidade e alinhada por prefixos. Ele utiliza $N$ queries latentes para agregar prefixos de tokens de comprimento crescente. À medida que o índice da query $i$ aumenta, o codificador atende a prefixos de tokens progressivamente mais longos ($w_{1:\lceil iL/N \rceil}$), induz um ordenamento causal onde as posições latentes se alinham com os prefixos textuais.
• Decodificador baseado em Query: Mapeia prefixos latentes de volta para logits de tokens. Crucialmente, o decodificador é restrito a atender apenas ao prefixo latente correspondente, garantindo uma correspondência causal estrita entre as sequências latente e de tokens.
• Design de Alta Capacidade: Ao contrário de trabalhos anteriores, que comprimem latentes para facilitar a geração, o AURORA-LM retém uma largura de canal alta ($D$) e uma compressão de sequência mínima ($c \approx 1$). Isso preserva a informação necessária para a recuperação exata do token, tratando o espaço latente como uma interface fiel ao texto, em vez de um prior simplificado.
• Treinamento: O autoencoder é treinado via entropia cruzada ao nível do token e, em seguida, é congelado. Isso fixa a distribuição latente alvo, permitindo que o modelo generativo seja estudado independentemente do aprendizado de representação.

2. Difusão Latente de Causalidade de Bloco (O Modelo Generativo)

Um Transformer de Difusão de Causalidade de Bloco aprende a distribuição das sequências latentes congeladas usando Flow Matching:

• Fatoração de Causalidade de Bloco: Para equilibrar a geração causal com a denoização paralela, a sequência latente é particionada em blocos. O modelo gera blocos da esquerda para a direita (condicionamento causal em blocos anteriores), mas realiza a denoização de todas as posições dentro de um bloco em paralelo.
• Predição de Latente de Largura Total: O modelo prevê o latente limpo de largura total ($x_0$) necessário pelo decodificador, em vez de uma versão comprimida.
• Gargalo de Entrada Ruidosa: Para lidar com a dificuldade de aprender uma distribuição de alta dimensão a partir de entradas ruidosas, o modelo aplica uma projeção de baixo rank apenas à via de entrada ruidosa ($W_{down} \in \mathbb{R}^{D_b \times D}$), enquanto a cabeça de saída prevê o alvo limpo de largura total. Isso permite que o modelo processe entradas corrompidas de forma eficiente sem reduzir a capacidade voltada ao decodificador.
• Calibragem de Ruído: A distribuição do nível de ruído é calibrada para a largura latente. Resultados empíricos mostram que representações latentes mais largas beneficiam-se de cronogramas de ruído que enfatizam entradas de alto ruído.
• Consistência de Trajetória Própria: Para preencher a lacuna entre o treinamento em estados ruidosos amostrados independentemente e a inferência via denoização iterativa, o modelo introduz uma perda de consistência. Isso alinha as previsões de latentes limpos em estados vizinhos ao longo da trajetória de denoização, estabilizando o processo de geração e permitindo resultados de alta qualidade com menos passos de amostragem.

3. Inferência

Durante a inferência, os blocos latentes são gerados sequencialmente da esquerda para a direita. Os latentes padronizados gerados são invertidos para o espaço do codificador original e passados ao decodificador congelado para recuperar a sequência de tokens final. A estrutura suporta tanto a geração incondicional quanto a geração condicionada por prompt (usando orientação de classificação livre ou orientação de autocondicionamento).

Principais Contribuições

1. Estrutura Unificada: O AURORA-LM separa o aprendizado de representação de texto da modelagem de distribuição generativa, construindo explicitamente uma representação contínua que o modelo de difusão aprende a gerar, em vez de herdá-la de embeddings de tokens.
2. Interface de Alta Capacidade: Retém a capacidade total de canal da representação latente para uma decodificação de token precisa, adaptando o prior de difusão para modelar a distribuição resultante complexa, em vez de simplificar a representação.
3. Flow Matching de Causalidade de Bloco: Introduz um denoiser de causalidade de bloco treinado com flow matching que gera blocos da esquerda para a direita, enquanto realiza a denoização de posições dentro dos blocos em paralelo.
4. Inovações Arquiteturais: A combinação de uma via de entrada ruidosa de baixo rank, calibragem de nível de ruído para a largura latente e consistência de trajetória própria permite que o modelo aprenda efetivamente distribuições latentes de largura total.

Resultados Experimentais

Os experimentos foram conduzidos em NPUs Ascend.

• OpenWebText (Geração Incondicional): O AURORA-LM-S (130M parâmetros) alcançou a menor Perplexidade de Geração (Gen-PPL: 23,56) e o maior MAUVE (0,890) entre as linhas de base autoregressivas, de difusão discreta e de incorporação contínua avaliadas.
• XSum (Sumarização Condicional): O AURORA-LM-S superou todas as linhas de base (incluindo ELF-B, AR e MDLM) nas métricas ROUGE-1, ROUGE-2 e ROUGE-L.
• Escalabilidade: O AURORA-LM-L (1,01B parâmetros, ~1.500 EFLOPs de computação) demonstrou ganhos de desempenho consistentes e superou um modelo de difusão latente público maior (Cola-DLM, ~1,8B parâmetros) sob um protocolo de avaliação correspondente.
• Estudos de Ablação:
  • Largura do Latente: Latentes mais largos ($D=1024$) forneceram melhor robustez à corrupção e maior desempenho generativo.
  • Gargalo de Entrada: Uma largura de gargalo moderada ($D_b=128$) para entradas ruidosas foi a ideal, permitindo a predição de largura total sem processar entradas corrompidas em sua dimensionalidade total.
  • Consistência de Trajetória: A consistência de trajetória própria melhorou significamente a qualidade da amostragem de poucos passos.

Significância e Alegações

O artigo afirma que a geração de linguagem contínua pode efetivamente unir a modelagem generativa baseada em difusão e a decodificação de tokens discretos. Os autores argumentam que o compromisso de design predominante — simplificar representações para acomodar modelos generativos — é desnecessário. Em vez disso, ao preservar um latente de texto de alta capacidade e decodificável e projetar o modelo de difusão para aprender sua distribuição diretamente (via escolhas arquiteturais específicas como fatoração de causalidade de bloco e consistência de trajetória), é possível alcançar um desempenho de estado da arte tanto em tarefas de geração de texto incondicional quanto condicional. Os resultados sugerem que espaços contínuos podem servir como uma interface geral para a linguagem, retendo fidelidade semântica e regularidade sintática, enquanto aproveitam os benefícios da modelagem de difusão.