d2: Improving Reasoning in Diffusion Language Models via Trajectory Likelihood Estimation

O artigo apresenta o d2, um novo framework de aprendizado por reforço para modelos de linguagem de difusão mascarada que utiliza estimadores de verossimilhança de trajetória especializados (d2-AnyOrder e d2-StepMerge) para aumentar significativamente as capacidades de raciocínio em benchmarks lógicos e matemáticos, alcançando um novo estado da arte em desempenho.

O essencial

Imagine que você tem um artista muito talentoso que pode pintar uma obra-prima começando com uma tela em branco e adicionando detalhes gradualmente até que a imagem esteja completa. É assim que os Modelos de Linguagem de Difusão (DLMs) funcionam. Ao contrário dos escritores de IA tradicionais que constroem frases palavra por palavra (como empilhar tijolos), esses modelos começam com uma bagunça de "ruído" (como estática em uma TV antiga) e lentamente "denoizam" esse ruído, revelando uma história ou solução coerente passo a passo.

O artigo apresenta um novo método de treinamento chamado d2 para ensinar esses artistas a resolver enigmas difíceis (como problemas matemáticos ou jogos de lógica) de forma muito melhor. Aqui está o detalhamento de como ele funciona, usando analogias simples.

O Problema: O Crítico "Cego"

Para ensinar uma IA a raciocinar, pesquisadores usam Aprendizado por Reforço (RL). Pense nisso como um jogo onde a IA tenta resolver um enigma, e um "Crítico" (um sistema de recompensa) dá uma pontuação a ela. Se a IA obtiver uma boa pontuação, ela aprende a fazer aquilo novamente.

No entanto, há um detalhe com os modelos de Difusão. Para aprender de forma eficaz, a IA precisa saber exatamente o quão provável foi a geração daquela sequência específica de palavras.

• Para a IA tradicional (Autoregressiva): É como escrever uma carta. Você sabe exatamente o que escreveu primeiro, segundo e terceiro. Calcular a "probabilidade" é fácil e rápido.
• Para a IA de Difusão: É como observar uma escultura emergindo de um bloco de pedra. A IA faz milhares de pequenos ajustes. Para saber a probabilidade exata da estátua final, você teria que, teoricamente, rebobinar o vídeo e reproduzir cada um dos quadros do processo de escultura. Fazer isso para cada etapa de treinamento é computacionalmente impossível (muito lento e caro).

Como não conseguiam calcular essa "pontuação" com precisão, os métodos anteriores estavam apenas supondo, o que levava a habilidades de raciocínio ruins.

A Solução: d2 (O Marcador de Pontos Inteligente)

Os autores criaram o d2, um framework que atua como um marcador de pontos super eficiente. Ele descobre a "probabilidade" da jornada da IA sem ter que reproduzir o filme inteiro todas as vezes. Eles construíram duas ferramentas diferentes para dois tipos diferentes de artistas:

Ferramenta 1: d2-AnyOrder (O "Espelho Mágico")

Alguns modelos de difusão são especiais; eles podem revelar a imagem final em qualquer ordem. Você poderia pintar os olhos primeiro, depois o nariz, depois o fundo, ou vice-versa, e ainda assim funciona.

• A Analogia: Imagine que você tem um espelho mágico. Em vez de assistir à pintura acontecer passo a passo, você pode olhar para a pintura terminada e ver instantaneamente o que o artista teria dito em cada etapa, tudo de uma vez.
• O Resultado: Esta ferramenta, chamada d2-AnyOrder, calcula a pontuação perfeita em um único olhar (uma passagem de computador). É incrivelmente precisa e faz a IA aprender muito rápido.

Ferramenta 2: d2-StepMerge (A "Estratégia de Fragmentação")

A maioria dos modelos de difusão populares (como o famoso LLaDA) não pode pintar em qualquer ordem; eles precisam seguir uma sequência específica. Eles não possuem o "espelho mágico".

• A Analogia: Como não podemos ver o filme inteiro de uma vez, dividimos o filme em fragmentos. Em vez de assistir a cada quadro individualmente, assistimos ao filme em segmentos de 10 segundos. Estimamos a pontuação para cada segmento e os somamos.
• O Resultado: Esta ferramenta, chamada d2-StepMerge, é uma aproximação. Não é perfeita, mas os autores provaram matematicamente que o erro é pequeno e controlável. Ela troca um pouco de precisão por um ganho massivo de velocidade, tornando-a prática para modelos padrão.

Os Resultados: De Novato a Grande Mestre

Os pesquisadores testaram essas ferramentas em alguns dos enigmas mais difíceis para a IA:

1. Jogos de Lógica: Sudoku e Countdown (um jogo matemático).
2. Problemas Matemáticos: GSM8K e MATH500 (benchmarks padrão para raciocínio matemático).

O Desfecho:

• Sem nenhum "guia de estudo" extra (Aprendizado Supervisionado/Fine-Tuning): O framework d2 ensinou os modelos a raciocinar melhor do que qualquer método anterior.
• Vencendo a Competição: Nos testes de Sudoku e Countdown, os modelos treinados com d2 alcançaram pontuações próximas a 92% e 56% respectivamente, esmagando os melhores métodos anteriores que estavam travados em torno de 22% e 42%.
• Maestria Matemática: Em benchmarks matemáticos complexos, o d2 estabeleceu um novo recorde de "Estado da Arte", provando que modelos de difusão podem ser tão bons quanto os modelos tradicionais de geração de texto no raciocínio, mas com os benefícios adicionais de serem mais rápidos e controláveis.

Resumo

O artigo diz: "Encontramos uma maneira de ensinar modelos de IA de Difusão a pensar claramente."

• Eles perceberam que a antiga maneira de avaliar o trabalho da IA era muito lenta e imprecisa.
• Eles inventaram o d2, um novo sistema de avaliação que é perfeitamente preciso (para modelos especiais) ou inteligente e eficiente (para modelos padrão).
• Usando este novo sistema, os modelos de IA aprenderam a resolver enigmas de lógica e matemática significativamente melhor do que antes, sem precisar ser pré-ensinados as respostas.

Em suma, o d2 dá aos modelos de difusão a capacidade de "olhar para trás" em seu próprio processo de pensamento de forma precisa, permitindo que aprendam com seus erros e se tornem verdadeiros especialistas em raciocínio.

Resumo técnico

Resumo Técnico: d2: Melhorando o Raciocínio em Modelos de Linguagem de Difusão via Estimativa de Verossimilhança de Trajetória

Declaração do Problema

Modelos de Linguagem de Difusão (DLMs), particularmente os DLMs de máscara (masked DLMs), emergiram como alternativas competitivas aos modelos autorregressivos (AR) devido à sua controlabilidade e capacidades de geração paralela. No entanto, aumentar suas habilidades de raciocínio via Aprendizado por Reforço (RL) permanece um desafio. Embora o RL tenha se tornado o padrão para induzir o raciocínio em modelos AR, aplicá-lo a DLMs não é trivial. A dificuldade central reside na formulação do gradiente de política: ao contrário dos modelos AR, onde as verossimilhanças fatorizam de forma limpa entre as posições dos tokens, a verossimilhança exata de uma trajetória de difusão é computacionalmente intratável. Calcular essas verossimilhanças de forma ingênua requer $T$ passagens de ida (onde $T$ é o número de passos de difusão), tornando métodos de gradiente de política padrão, como o Group Relative Policy Optimization (GRPO), computacionalmente proibitivos. As aproximações existentes frequentemente introduzem um viés significativo, levando a atualizações de política subótimas.

Metodologia

Os autores apresentam o d2, um framework de raciocínio adaptado para masked DLMs. O framework centra-se em um novo algoritmo de gradiente de política derivado do objetivo GRPO, que exige explicitamente a estimativa precisa das verossimilhanças das trajetórias de amostragem. Para abordar o gargalo computacional, o artigo propõe uma família de estimadores adaptados a distintas classes de modelos:

1. Fundamentação Teórica

Os autores derivam um objetivo do tipo GRPO para masked DLMs através da marginalização da verossimilhança sobre os latentes temporais e do uso de amostragem por importância. A decomposição de gradiente resultante (Teorema 3.1) depende da razão das verossimilhanças de trajetória entre a política atual $\pi_\theta$ e uma política obsoleta $\pi_{old}$. O desafio principal é avaliar $\pi(x_{0:T})$ de forma eficiente.

2. Os Estimadores

O framework d2 oferece dois estimadores específicos baseados na arquitetura do modelo:

• d2-AnyOrder (Estimador Exato):

  • Alvo: DLMs que suportam decodificação de qualquer ordem (AO-dLLMs). Esses modelos permitem que os tokens sejam decodificados em qualquer ordem, desde que a máscara de atenção respeite restrições de causalidade específicas (máscaras independentes e causalidade de ordem).
  • Mecanismo: Ele constrói uma sequência de comprimento $2L$ (concatenando os tokens limpos e os tokens de máscara) com codificações posicionais e máscaras de atenção específicas. Isso permite que o modelo compute as verossimilhanças condicionais exatas de todos os tokens em uma única passagem de ida (single forward pass).
  • Limitação: Masked DLMs padrão (como o LLaDA) não suportam nativamente a decodificação de qualquer ordem; seus objetivos de treinamento não garantem as propriedades de atenção necessárias para que este estimador seja não enviesado.

• d2-StepMerge (Estimador Aproximado):

  • Alvo: Modelos de difusão de máscara (MDMs) padrão que não suportam decodificação de qualquer ordem.
  • Mecanismo: Inspirado pelo likelihood composto por blocos, este estimador aproxima a verossimilhança da trajetória dividindo a trajetória de $T$ passos em $N$ segmentos temporais contíguos. Em vez de avaliar cada passo, ele avalia a verossimilhança dos tokens não mascarados dentro de cada segmento usando uma única passagem de ida por segmento.
  • Compromisso (Trade-off): Ele reduz o número de passagens de ida de $T$ para $N$. O artigo fornece um limite superior teórico (Teorema 4.1) sobre o erro de aproximação (divergência KL), mostrando que o erro diminui monotonicamente à medida que $N$ aumenta, quantificando um compromisso entre computação e viés.

Contribuições Principais

1. Framework d2: Um framework de pós-treinamento de RL principiado para masked DLMs que deriva um objetivo GRPO dependente da estimativa precisa da verossimilhança da trajetória.
2. d2-AnyOrder: Um estimador de verossimilhança de passagem única e exato para AO-dLLMs, demonstrando que a estimativa de verossimilhança não enviesada supera significativamente os baselines existentes em cenários de RL.
3. d2-StepMerge: Um estimador determinístico de múltiplas passagens para MDMs padrão com erro de aproximação analiticamente limitado, oferecendo um compromisso eficiência-precisão superior em comparação com métodos anteriores como o diffu-GRPO.
4. Validação Empírica: Experimentos extensos mostrando que o d2 atinge o estado da arte em benchmarks de raciocínio sem depender de ajuste fino supervisionado de cadeia de pensamento (CoT).

Resultados Experimentais

Os autores avaliaram o d2 em benchmarks de raciocínio lógico (Countdown, Sudoku) e raciocínio matemático (GSM8K, MATH500), bem como tarefas de codificação e controle de toxicidade.

• Desempenho em AO-dLLMs: Quando aplicado a modelos que suportam decodificação de qualquer ordem (ex: Eso-LM, Set Diffusion e um LLaDA customizado de Any-Order), o d2-AnyOrder superou significativamente os baselines como DDPO e diffu-GRPO. Por exemplo, no GSM8K com uma base Qwen3-1.7B, o d2-AnyOrder alcançou 67% de acurácia comparado a 63% do diffu-GRPO.
• Desempenho em MDMs Padrão: Aplicado ao LLaDA-8B-Instruct (um masked DLM padrão), o d2-StepMerge (com $N=16$) estabeleceu novos resultados de estado da arte:
  • Sudoku: 91.9% (vs. 22.1% para d1 com SFT).
  • Countdown: 56.6% (vs. 42.2% para d1).
  • GSM8K: 85.0% (vs. 82.1% para d1).
  • MATH500: 41.6% (vs. 39.0% para wd1).
    Notavelmente, esses resultados foram alcançados sem ajuste fino supervisionado em dados extras de CoT, enquanto alguns baselines dependiam de tais dados.
• Generalização: O método generalizou-se para outras arquiteturas (Dream 7B) e tarefas (HumanEval, MBPP), mostrando melhor generalização de comprimento em comparação ao diffu-GRPO.
• Ablação: Experimentos confirmaram que verossimilhanças exatas (d2-AnyOrder) fornecem sinais de maior fidelidade do que aproximações, e que o d2-StepMerge com um $N$ ótimo (ex: 16) equilibra melhor o custo computacional e a precisão do que valores extremos.

Significância e Alegações

O artigo afirma que a estimativa precisa da verossimilhança da trajetória é central para o sucesso do RL para raciocínio em DLMs. Os autores argumentam que métodos anteriores falharam em atingir o desempenho ideal porque dependiam de estimativas de verossimilhança enviesadas ou computacionalmente caras.

• Contribuição Teórica: O trabalho fornece a primeira derivação rigorosa de um objetivo GRPO para masked DLMs e estabelece limites teóricos sobre o erro de aproximação de estratégias de fusão de passos (step-merging).
• Impacto Prático: O d2 demonstra que modelos de difusão podem alcançar capacidades de raciocínio competitivas com, ou superiores a, modelos autorregressivos quando treinados adequadamente com RL, mesmo sem dados extensos de CoT supervisionado.
• Perspectiva Futura: Os autores observam que, embora o d2-AnyOrder seja superior, sua aplicabilidade é atualmente limitada a modelos treinados especificamente para decodificação de qualquer ordem. Eles veem o desenvolvimento de own-order dLLMs de propósito geral mais fortes como um próximo passo natural para aproveitar totalmente o estimador exato.

O artigo conclui que o d2 estabelece um novo estado da arte para DLMs em tarefas de raciocínio lógico e matemático, validando a eficácia de suas estratégias de estimativa de verossimilhança para desbloquear o potencial de raciocínio de modelos de linguagem baseados em difusão.