Resumo Técnico: Repensando o Raciocínio com Modelos de Linguagem de Difusão Mascarada

Problema

Os paradigmas atuais de raciocínio em Grandes Modelos de Linguagem (LLMs) dependem fortemente de modelos de Predição do Próximo Token (NTP) treinados via geração autorregressiva. Embora eficazes, os modelos NTP enfrentam limitações inerentes em tarefas de raciocínio:

1. Amostragem de Posterior Intratável: Gerar traços de raciocínio de alta qualidade condicionados a uma resposta correta conhecida (amostragem de posterior, $p(r|c,a)$) é computacionalmente intratável para modelos autorregressivos porque eles geram da esquerda para a direita e não conseguem facilmente condicionar-se a tokens futuros (a resposta) sem engenharia de prompt complexa ou modelos auxiliares.
2. Falta de Pontuação Intermediária: Avaliar a correção de traços de raciocínio parciais normalmente requer o treinamento de Modelos de Recompensa de Processo (PRMs) separados, o que adiciona sobrecarga computacional e complexidade.
3. Inferência Ineficiente: A decodificação autorregressiva padrão é sequencial, impedindo a geração paralela de tokens e limitando oportunidades de terminação precoce quando o modelo está confiante.

Modelos de Linguagem de Difusão Mascarada (MDLMs) oferecem um objetivo de treinamento alternativo onde os modelos aprendem a preencher (infill) posições aleatoriamente mascaradas em uma sequência. No entanto, as capacidades específicas dos MDLMs para o raciocínio — além da decodificação paralela — não foram totalmente exploradas ou estruturadas.

Metodologia: Raciocínio-como-Infilling

O artigo introduz o Raciocínio-como-Infilling, uma técnica de prompting que aproveita as propriedades únicas dos MDLMs para reestruturar o processo de raciocínio. Em vez de tratar o prompt como um simples prefixo, o método pré-preenche a sequência de saída com um template especificado pelo usuário que delimita explicitamente blocos de raciocínio e blocos de resposta.

O mecanismo central envolve:

1. Pré-preenchimento de Template: O contexto de entrada inclui a pergunta e um template com tokens de raciocínio mascarados seguidos por um delimitador e tokens de resposta mascarados (ex: [Bloco de Raciocínio] <Delimitador> [Bloco de Resposta]).
2. Condicionamento Explícito: Como as posições do bloco de resposta são designadas explicitamente e mascaradas, o MDLM pode computar distribuições condicionais sobre os tokens da resposta ($p(a_j | r_{unmasked}, c)$) em qualquer etapa intermediária da geração do raciocínio.

Este framework permite três capacidades primárias:

1. Saídas Antecipadas (Early Exits)

Como o modelo mantém uma distribuição sobre os tokens da resposta ao longo da geração do traço de raciocínio, o sistema pode monitorar a incerteza da região da resposta.

• Mecanismo: Os autores propõem o uso do limite superior de entropia ($H_{UB}$) do bloco de resposta como uma métrica de certeza. $H_{UB}$ é calculado como a soma das entropias marginais dos tokens de resposta dados os tokens de raciocínio não mascarados.
• Ação: Se $H_{UB}$ cair abaixo de um limiar $\gamma$, o modelo pula as etapas de raciocínio restantes e gera diretamente a resposta, acelerando significativamente a inferência.

2. Raciocínio Post-hoc (Amostragem de Posterior)

O framework permite a amostragem de traços de raciocínio condicionados a uma resposta correta, uma distribuição que é intratável para modelos NTP.

• Mecanismo: Ao pré-preencher o bloco de resposta com uma resposta correta conhecida ($a^*$), o MDLM amostra da distribuição de posterior $p_\theta(r | c, a^*)$.
• Aplicação: Esses "traços de posterior" servem como dados de raciocínio de alta qualidade para ajuste fino (fine-tuning). O artigo demonstra que o ajuste fino nesses traços gerados pelo modelo pode melhorar o desempenho mesmo quando o modelo inicialmente falhou em resolver o problema.

3. Pontuação Intermediária sem PRMs

O framework fornece um mecanismo integrado para pontuar traços de raciocínio parciais sem modelos de recompensa externos.

• Mecanismo: Dado um traço de raciocínio de referência $a^*$ e um traço de raciocínio parcial $r_{unmasked}$, a pontuação é calculada como a soma dos log-probabilidades dos tokens da resposta: $\phi_{IPRM} = \sum \log p_\theta(a^*_j | c, r_{unmasked})$.
• Vantagem: Esta pontuação correlaciona-se fortemente com a correção final e está disponível em cada passagem direta (forward pass), ao contrário dos modelos autorregressivos onde as verossimilhanças da resposta são preditores fracos até o final do processo de geração.

Resultados Principais

Aceleração de Inferência (Saídas Antecipadas)

Avaliado no GSM8k e MATH500 usando modelos LLaDA-8B e Dream-7B:

• Compromisso Velocidade-Acurácia: As saídas antecipadas sozinhas alcançaram um ganho de velocidade de 1.3× no GSM8k com uma queda de menos de 1 ponto percentual na acurácia.
• Aceleração Combinada: Quando combinados com algoritmos de decodificação paralela (ex: Block Parallel Confidence Decoding), o sistema alcançou até 4.0× de aceleração (3.5× para LLaDA, 2.5× para Dream) com degradação mínima de acurácia.
• Suporte a Comprimento Variável: O método integra-se com modelos de difusão de comprimento variável (ex: DreamOn), permitindo o redimensionamento dinâmico do orçamento de raciocínio mantendo alta acurácia.

Raciocínio Post-hoc e Ajuste Fino

• Qualidade do Traço: Quando o modelo LLaDA-8B-Instruct falhou em resolver um problema do GSM8k, a amostragem do posterior (condicionado à resposta correta) gerou traços de raciocínio que foram julgados corretos pelo GPT-4o e por um PRM especializado em ~40% dos casos. Mesmo o modelo base (sem ajuste de instrução) produziu traços corretos para >30% dessas falhas.
• Ganhos de Desempenho: O ajuste fino do modelo LLaDA-8B-Base nesses traços gerados via posterior melhorou a acurácia de teste do GSM8k em +14.9%. Este ganho igualou a melhoria (+13.4%) alcançada pelo ajuste fino em traços de raciocínio anotados por humanos caros.
• Comparação: Traços de raciocínio de posterior superaram traços gerados via prompting de dica de prefixo estilo STaR.

Pontuação Intermediária

• Correlação: As pontuações de verossimilhança da resposta dos MDLMs em etapas intermediárias mostraram uma correlação mais forte com a correção final do que as pontuações de um Modelo de Recompensa de Processo (PRM) de 7B especializado.
• Baseline Autorregressivo: Em contraste, as pontuações intermediárias de verossimilhança da resposta de modelos autorregressivos (Llama-3.1-8B, Qwen2.5-7B) foram significativamente menos preditivas de correção até fases tardias do processo de geração.

Significância e Alegações

O artigo afirma que o framework Raciocínio-como-Infilling desbloqueia um conjunto unificado de capacidades para MDLMs que não estão prontamente disponíveis para modelos NTP:

1. Eficiência: Permite aceleração significativa de inferência através de saídas antecipadas e decodificação paralela sem sacrificar a acurácia.
2. Geração de Dados: Fornece um método tratável para gerar traços de raciocínio de alta qualidade condicionados a respostas corretas, oferecendo uma alternativa escalável à anotação humana para ajuste fino.
3. Auto-Pontuação: Permite que os modelos pontuem seus próprios processos de raciocínio em etapas intermediárias usando suas distribuições internas, eliminando a necessidade de Modelos de Recompensa de Processo separados em certos cenários de pós-treinamento.

Os autores concluem que o objetivo de treinamento MDLM, quando combinado com prompting estruturado, oferece benefícios promissores para tarefas de raciocínio, particularmente em cenários que exigem recompensas verificáveis, inferência eficiente ou a geração de dados de raciocínio a partir de respostas corretas. O trabalho sugere que o comportamento "patológico" dos MDLMs (ex: decodificar respostas antes do raciocínio) pode ser reformulado como uma característica para geração de dados de pós-treinamento quando a resposta é conhecida.