Resumo Técnico: PreDiff-LM

Declaração do Problema

Modelos de linguagem de difusão discreta com mascaramento oferecem capacidades de geração bidirecional e preenchimento (infilling), abordando as limitações dos modelos autorregressivos (AR), que são estritamente da esquerda para a direita. No entanto, existe uma lacuna de desempenho significativa entre modelos de difusão treinados do zero e modelos AR de tamanho comparável. Embora trabalhos recentes tenham estabelecido que inicializar modelos de difusão com pesos AR pré-treinados (ex: DiffuGPT, Dream) é uma estratégia viável, um desafio específico permanece sem solução: como adaptar o mecanismo de atenção ao reutilizar pesos AR causais para denoising bidirecional.

Os pesos AR são aprendidos sob máscaras de atenção estritamente causais. Simplesmente substituir essas máscaras por máscaras bidirecionais uniformes para toda a sequência interrompe as representações aprendidas do prompt de condicionamento, enquanto manter a causalidade total impede que os tokens alvo utilizem o contexto bidirecional necessário. O artigo investiga se um padrão de atenção assimétrico pode reconciliar esses requisitos conflitantes.

Metodologia: PreDiff-LM

O framework proposto PreDiff-LM (Pretrained Discrete Masked Diffusion Language Modeling) introduz um mecanismo de atenção híbrido e uma receita de treinamento específica para adaptar transformers causais pré-treinados para difusão discreta.

1. Atenção Híbrida Causal-Bidirecional

A inovação central é uma máscara de atenção 4D assimétrica ($M$) que particiona a entrada em um prefixo de prompt ($p$) e uma região alvo ($x_t$):

• Tokens do Prompt: Mantêm o padrão de atenção causal original ($j \le i$). Isso preserva o regime de representação interna no qual os pesos AR foram pré-treinados, garantindo que o prompt de condicionamento seja processado como o modelo originalmente aprendeu.
• Tokens Alvo: São permitidos a atenção bidirecional total. Eles podem atender a todo o prompt e a todos os outros tokens alvo (contexto à esquerda e à direita).
• Restrição: Informações de tokens alvo corrompidos não podem fluir de volta para atualizar as representações do prompt ($i \in prompt, j \in target \implies M_{ij}=0$).

2. Componentes de Treinamento e Inferência

Para suportar a atenção híbrida, o modelo incorpora vários componentes padrão, porém críticos, para a adaptação de difusão:

• Auto-Condicionamento (Self-Conditioning): Para mitigar inconsistências entre tokens simultaneamente mascarados, o modelo alimenta suas próprias previsões suaves de uma passagem anterior de volta para a rede como um resíduo de embedding. Isso é crucial para altas taxas de máscara.
• Embedding de Taxa de Máscara por Tempo: O nível de ruído (taxa de máscara $t$) é injetado via codificação senoidal e um MLP, permitindo que a rede se condicione ao estágio de denoising.
• Desmascaramento Consciente de Confiança (CAU - Confidence-Aware Unmasking): Durante a inferência, os tokens são desmascarados iterativamente com base em pontuações de confiança. O número de tokens desmascarados segue um cronograma de cosseno, começando com poucos tokens (alta incerteza) e acelerando conforme o contexto se acumula.

3. Objetivo

O modelo minimiza a entropia cruzada esperada sobre as posições mascaradas, consistente com os objetivos de difusão de máscara discreta (ex: MDLM), mas treinado em pesos inicializados de um modelo AR pré-treinado (GPT-2 ou Llama-3.2).

Principais Contribuições

1. Formulação de Adaptação de Atenção: O artigo isola a adaptação de atenção como um componente distinto da transferência de AR para difusão. Ele introduz a máscara híbrida causal-bidirecional, demonstrando que preservar a causalidade do lado do prompt enquanto permite a bidirecionalidade do lado do alvo é superior à atenção bidirecional uniforme.
2. Componibilidade de Adaptações: O estudo mostra que a adaptação ao nível de atenção e a adaptação ao nível de objetivo (estilo DiffuGPT) são complementares. Combinar ambas produz melhores resultados (26.9 PPL) do que qualquer uma delas isoladamente.
3. Avaliação Rigorosa: Os autores fornecem controles AR ajustados com computação equivalente, métricas sensíveis à repetição (Rep-4, Self-BLEU), tarefas de zero-shot de domínio inferior e uma auditoria de preferência humana cega. Isso estabelece uma linha de base clara para o que os modelos de difusão podem alcançar em relação aos modelos AR sem exagerar a paridade.
4. Análise de Eficiência: O trabalho distingue entre eficiência de treinamento (onde a difusão se beneficia grandemente da inicialização AR) e latência de inferência (onde o decoding AR otimizado permanece mais rápido nos pontos operacionais padrão).

Resultados Experimentais

Experimentos foram conduzidos no WikiText-103 usando backbones GPT-2 (133M a 1.56B parâmetros) e backbones Llama-3.2.

• Perplexidade (PPL): Em um cenário de configuração correspondente (GPT-2 Medium, 90K passos), o PreDiff-LM com atenção híbrida alcançou 28.7 PPL, uma melhoria significativa em relação ao controle de atenção bidirecional uniforme (34.1 PPL) e outros baselines de difusão como DiffuGPT (42.8 PPL) e MDLM (51.2 PPL).
  • Nota: Um modelo AR ajustado com configuração correspondente permaneceu mais forte com 18.9 PPL.
• Eficiência de Treinamento: A inicialização pré-treinada reduziu drasticamente os passos necessários para atingir PPL < 50. O PreDiff-LM atingiu esse limiar em 8K passos, comparado a aproximadamente 350K passos para o MDLM treinado do zero.
• Qualidade de Geração: O PreDiff-LM mostrou melhorias na qualidade distributiva (MAUVE: 0.78 vs. 0.71 para o controle), redução de repetição (Rep-4: 0.012) e melhor desempenho em tarefas de zero-shot (LAMBADA, HellaSwag, etc.) em comparação com outros baselines de difusão, embora ainda atrás do controle AR.
• Velocidade de Inferência: Com 32 passos de denoising (a configuração para os principais resultados de qualidade), o PreDiff-LM foi mais lento (450 tokens/s) que o decoding AR otimizado (1.120 tokens/s). No entanto, com 8 passos, ele alcançou um throughput competitivo (1.800 tokens/s) com menor qualidade (35.4 PPL), ofereando um equilíbrio ajustável entre velocidade e qualidade.
• Preferência Humana: Em uma avaliação cega, o PreDiff-LM foi preferido em relação ao MDLM (64.2%) e SEDD (72.5%), mas venceu apenas 34.2% contra o controle do GPT-2 Medium AR.

Significância e Alegações

O artigo posiciona a atenção híbrida como um mecanismo complementar para adaptar backbones causais pré-treinados para difusão. Ele afirma explicitamente:

• Efetividade: A atenção híbrida melhora significamente a qualidade dos modelos de difusão inicializados a partir de pesos AR em comparação com a atenção bidirecional uniforme.
• Eficiência: A inicialização pré-treinada torna o processo de adaptação altamente eficiente, exigindo ordens de magnitude a menos de passos de treinamento do que o treinamento do zero.
• Modéstia sobre Paridade: Os autores não alegam que os modelos de difusão superaram os modelos AR em qualidade ou velocidade de inferência. Eles reconhecem que modelos AR ajustados com configurações equivalentes permanecem superiores em PPL e que o decoding AR otimado é mais rápido.
• Identificação de Lacunas: O trabalho destaca as lacunas restantes em qualidade correspondente, eficiência de decoding e robustez de contexto longo que precisam ser abordadas para que os modelos de difusão possam competir diretamente com sistemas AR otimizados.

Em resumo, o PreDiff-LM demonstra que a adaptação cuidadosa da atenção é crítica para transferir o conhecimento causal de AR para a difusão discreta, permitindo uma geração bidirecional de alta qualidade enquanto mantém os benefícios das representações pré-treinadas, mesmo que ainda não tenha fechado totalmente a lacuna com os modelos autorregressivos de última geração.