Resumo Técnico: Não Re-treine—Alinhe: Adaptando LMs Autoregressivos para LMs de Difusão via Alinhamento de Representações

Declaração do Problema

Modelos de Linguagem de Difusão (DLMs) oferecem vantagens distintas sobre os modelos autoregressivos (AR) padrão, particularmente na geração não sequencial, edição bidirecional e refinamento iterativo. No entanto, a escalabilidade dos DLMs é computacionalmente dispendiosa. Teoricamente, os DLMs devem aprender $L!$ caminhos de geração em comparação com o único caminho da esquerda para a direita dos modelos AR, frequentemente exigindo significativamente mais poder de computação.

Abordagens recentes tentam mitigar esse custo convertendo checkpoints AR pré-treinados em DLMs. Embora esses métodos consigam inicializar DLMs a partir de pesos AR, eles dependem principalmente do treinamento contínuo de remoção de ruído com modificações em objetivos, padrões de atenção ou convenções de amostragem. Crucialmente, as receitas existentes não restringem explicitamente o DLM convertido a preservar a geometria de representação interna aprendida durante o pré-treinamento AR original. Isso levanta uma questão central: é necessário tratar AR e DLM como paradigmas disjuntos que exigem o reaprendizado de representações de linguagem, ou a estrutura semântica aprendida pelo pré-treinamento AR pode ser explicitamente preservada durante a conversão?

Metodologia: REPR-ALIGN

Os autores propõem o REPR-ALIGN, um método de adaptação minimalista que trata o treinamento de DLM como um problema de alinhamento, e não como um problema de aprendizado de representações. A hipótese central é que a "parte difícil" da geração de linguagem é aprender estruturas semânticas e sintáticas (representações), que o pré-treinamento AR já dominou. A tarefa restante é meramente adaptar essas características a um mecanismo de decodificação iterativo e de qualquer ordem.

O método envolve dois modelos com arquitetura idêntica:

1. Professor: Um transformador autoregressivo (AR) pré-treinado e congelado, com atenção causal.
2. Aluno: Um modelo de linguagem de difusão mascarada (DLM) com atenção bidirecional, inicializado a partir dos mesmos pesos AR.

Procedimento de Treinamento

Durante o treinamento, o modelo aluno consome uma sequência corrompida (com tokens mascarados) e tenta prever os tokens mascarados usando um objetivo padrão de remoção de ruído mascarado. Simultaneamente, o modelo professor congelado processa a sequência limpa sob atenção causal.

Para preservar a geometria interna do modelo AR, os autores introduzem uma função de perda de alinhamento de representação por camada. Para cada camada $\ell$ na pilha do transformador, os estados ocultos do aluno ($h^{(\ell)}_D$) são alinhados com os estados ocultos do professor ($h^{(\ell)}_{AR}$) usando uma perda de similaridade cosseno.

O objetivo total de treinamento é:
$$\mathcal{L}(\theta) = \mathcal{L}_{diff}(\theta) + \lambda \mathcal{L}_{align}(\theta)$$
Onde:

• $\mathcal{L}_{diff}$ é a perda padrão de entropia cruzada de remoção de ruído mascarado.
• $\mathcal{L}_{align}$ é a distância cosseno média entre os estados ocultos do aluno e do professor em todas as camadas e posições mascaradas.
• $\lambda$ é um hiperparâmetro de ponderação (padrão 10).

Escolhas de design chave incluem:

• Sem Adaptadores: O método não requer parâmetros adicionais ou mudanças arquiteturais além da troca da máscara de atenção de causal para bidirecional.
• Professor Congelado: O professor AR permanece congelado, fornecendo um sistema de coordenadas representacional estável.
• Convenção de Deslocamento: O método adere à convenção de deslocamento padrão usada na adaptação de AR para DLM para reconciliar a indexação do próximo token com a remoção de ruído por posição.

Contribuições Principais

1. Identificação da Preservação de Representações: O artigo identifica que a preservação explícita da geometria de representação interna do modelo AR é um ingrediente ausente nas receitas atuais de conversão AR→DLM.
2. Uma Receita de Conversão Minimalista: O REPR-ALIGN introduz um método simples, sem adaptadores, que ancora o aluno DLM ao professor AR congelado via alinhamento cosseno por camada. Essa abordagem melhora a qualidade da conversão e a eficiência de amostras em várias escalas de modelos.
3. Eficiência em Dados e Computação: Os autores demonstram que a conversão AR→DLM não é inerentemente faminta por dados. Com o alinhamento de representações, o treinamento em um subconjunto minúsculo de dados (0,8B tokens) pode superar o treinamento no conjunto de dados completo (50B tokens) sob o mesmo orçamento de passos. Além disso, congelar grandes blocos (embeddings e MLPs) melhora o rendimento do treinamento em até 2$\times$ sem degradar a qualidade.
4. Validação de Escala: Os autores treinaram um "oDLM" de 4B parâmetros usando esta receita. Comparado ao Dream-7B (um DLM público forte também construído sobre inicialização AR), o oDLM alcançou melhores pontuações HumanEval e HumanEval+ pass@10, apesar de usar menos parâmetros e um orçamento de dados e computação significativamente mais leve.

Resultados Experimentais

O método foi avaliado em benchmarks de geração de código (HumanEval, HumanEval+, MBPP, MBPP+) usando checkpoints Qwen3 nas escalas de 0,6B, 1,7B e 4B.

• Velocidade de Adaptação: O REPR-ALIGN acelerou consistentemente a adaptação. No HumanEval pass@10, o modelo alinhado superou tanto o ajuste fino AR quanto o treinamento do zero durante a conversão inicial. Em 200k passos, o modelo de 0,6B viu um ganho de 6,1 pontos, e o modelo de 1,7B viu um ganho de 9,4 pontos sobre a linha de base.
• Escalonamento: Os benefícios do alinhamento cresceram com o tamanho do modelo. O modelo de 1,7B mostrou melhorias absolutas maiores do que o modelo de 0,6B, sugerindo que geometrias de pré-treinamento AR maiores são mais valiosas para preservar.
• Eficiência de Dados: O treinamento em uma subamostra aleatória de 0,8B tokens com alinhamento produziu pontuações pass@10 mais altas do que o treinamento no fluxo completo de 50B tokens no mesmo número de passos.
• Congelamento: Congelar embeddings de tokens e blocos MLP aumentou o rendimento em $\sim$2$\times$ enquanto mantinha ou melhorava ligeiramente o desempenho, apoiando a visão de que o sinal de aprendizado primário está em adaptar o mecanismo de remoção de ruído, e não em reaprender representações.
• Ablações:
  • Métrica: O alinhamento cosseno superou a distância L2, sugerindo que o sinal útil é geométrico (direção) e não métrico (magnitude).
  • Peso: $\lambda=10$ forneceu o compromisso ótimo entre ancoragem e adaptação.
  • Camadas: Alinhar todas as camadas forneceu o melhor pass@1, enquanto as camadas médias e superiores foram particularmente importantes para o pass@10.

Significado e Alegações

O artigo alega que as representações linguísticas são universais, independentemente da ordem de geração. Ao reenquadrar o treinamento de DLM como um problema de alinhamento, os autores argumentam que o alto custo do pré-treinamento de difusão deve-selargamente ao reaprendizado desnecessário de representações que o pré-treinamento AR já forneceu.

O significado deste trabalho reside em oferecer uma "técnica padrão" para treinar modelos de linguagem de difusão que:

• Reduz Computação: Alcança desempenho competitivo com uma fração do custo padrão de pré-treinamento.
• Simplifica o Treinamento: Remove a necessidade de adaptadores complexos ou ajuste fino pesado, confiando em vez disso em um objetivo de alinhamento simples.
• Democratiza o Acesso: Torna o treinamento de DLM de alta qualidade viável em regimes de baixos dados e com recursos de computação limitados.

Os autores posicionam seu trabalho como uma mudança de "re-treinamento" para "alinhamento", sugerindo que avanços futuros na geração discreta podem vir de um melhor alinhamento de mecanismos, e não de paradigmas de treinamento do zero.