🤖 machine learning

Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization

Este artigo introduz a Otimização de Política de Correspondência de Distribuição (DMPO), uma nova estrutura de aprendizado por reforço que aumenta significativamente as capacidades de raciocínio de modelos de linguagem de difusão ao alinhar sua distribuição de política com um alvo ótimo inclinado à recompensa, alcançando melhorias substanciais de precisão sobre as linhas de base existentes sem ajuste fino supervisionado.

Autores originais: Yuchen Zhu, Wei Guo, Jaemoo Choi, Petr Molodyk, Bo Yuan, Molei Tao, Yongxin Chen

Publicado 2026-07-14

📖 7 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Yuchen Zhu, Wei Guo, Jaemoo Choi, Petr Molodyk, Bo Yuan, Molei Tao, Yongxin Chen

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você tem um robô superinteligente que escreve histórias, resolve problemas de matemática e joga jogos. A maioria dos robôs hoje trabalha como uma pessoa lendo um livro palavra por palavra, da esquerda para a direita. Eles não podem pular adiante ou olhar para trás; eles apenas têm que adivinhar a próxima palavra com base nas anteriores. Isso é lento, como tentar preencher um quebra-cabeça gigante olhando apenas para a peça que está bem na sua frente.

Entram em cena os Modelos de Linguagem de Difusão (dLLMs). Eles são os novos fenômenos do mercado. Em vez de escrever palavra por palavra, eles começam com uma bagunça desordenada de "tokens misteriosos" (como um quebra-cabeça onde cada peça está coberta por uma névoa) e lentamente limpam essa bagunça, revelando a resposta de uma só vez ou em qualquer ordem que desejarem. Isso os torna potencialmente muito mais rápidos para pensar.

Mas aqui está o detalhe: embora esses robôs de difusão sejam rápidos, eles nem sempre são os mais espertos em tarefas de raciocínio complexas, como matemática avançada ou enigmas de lógica. Para torná-los mais inteligentes, os cientistas geralmente usam uma técnica chamada Aprendizado por Reforço (RL). Pense nisso como um videogame onde o robô ganha pontos por uma boa resposta e perde pontos por uma resposta ruim.

O Problema: A Armadilha do "Busca de Modo" (Mode-Seeking)

A maneira antiga de treinar esses robôs (usando métodos como o GRPO) era um pouco como um estudante que estuda apenas a única resposta que acredita estar correta. Se o robô encontra uma maneira de ganhar pontos que parece "segura", ele para de explorar. Ele fica preso em uma rotina, ignorando outras soluções engenhosas que podem ser tão boas quanto. No artigo, os autores chamam isso de "busca de modo" (mode-seeking). É como um robô que aprende a resolver um problema de matemática de uma única maneira específica e, se esse caminho for bloqueado, ele entra em pânico. Ele também tende a ignorar as respostas "bagunçadas", mas corretas, que parecem diferentes da primeira que encontrou.

O artigo argumenta que essa velha maneira de apenas perseguir a pontuação mais alta é falha para modelos de difusão. Sugere que, ao focar apenas no melhor caminho único, perdemos a capacidade única de exploração de múltiplos caminhos simultâneos do robô.

A Solução: DMPO (O Robô "Correspondência de Mapa")

Os autores propõem um novo método chamado Otimização de Política de Correspondência de Distribuição (DMPO).

Em vez de dizer ao robô: "Apenas encontre a pontuação mais alta", o DMPO diz: "Aqui está o mapa inteiro de todas as boas respostas que você poderia dar. Seu trabalho é aprender a corresponder a todo esse mapa".

Imagine que você está tentando ensinar um cachorro a buscar uma bola.

A Maneira Antiga: Você joga uma bola e o cachorro corre para o único lugar onde a bola caiu. Se a bola cair em um arbusto, o cachorro aprende a buscar apenas em arbustos.
O Jeito DMPO: Você mostra ao cachorro um mapa de todos os lugares possíveis onde a bola poderia cair (arbustos, grama, areia, água) e diz: "Aprenda a buscar em qualquer um desses lugares, ponderado pela qualidade do local". O cachorro aprende a ser flexível e a explorar todo o quintal, não apenas um canto.

Em termos técnicos, o DMPO usa um truque matemático especial chamado Entropia Cruzada de Denoising Ponderada (WDCE). Isso permite que o robô aprenda com suas tentativas passadas (mesmo aquelas que não foram perfeitas) sem precisar regenerar tudo do zero a cada vez. É como ter um "buffer de replay" onde o robô pode estudar seus jogos antigos repetidamente, aprendendo com os movimentos bons e ruins sem se confundir.

O Ingrediente Secreto: O "Peso de Referência" (Weight Baseline)

Os autores encontraram um problema complicado ao treinar com pequenos grupos de exemplos (pequenos "tamanhos de lote" ou batch sizes). Às vezes, o robô ficava confuso e começava a recompensar respostas ruins apenas porque eram as únicas que ele via.

Para corrigir isso, eles inventaram um truque inteligente chamado Subtração de Linha de Base de Peso (Weight Baseline Subtraction).
Pense nisso como um professor corrigindo uma prova. Se um aluno acerta uma questão, ele ganha uma estrela dourada. Mas se o professor apenas viu aquela questão, ele pode pensar que tudo o que o aluno fez foi ótimo. A "linha de base" é como um "padrão de média". O professor subtrai a pontuação "média" da pontuação do aluno.

Se o aluno fez melhor que a média, ele ganha uma grande estrela dourada.
Se ele fez pior que a média, ele recebe uma "penalidade" (um peso negativo), mesmo que tecnicamente tenha obtido alguns pontos.

Isso garante que o robô não fique animado demais com respostas medíocres e continue buscando as verdadeiramente excelentes.

Os Resultados: O Quanto Melhorou?

Os autores testaram este novo método em alguns benchmarks de raciocínio muito difíceis. Eles aplicaram o DMPO em modelos pré-treinados (como LLaDA-Instruct e Dream-Instruct) de duas maneiras diferentes:

Aplicação Direta: Aplicaram o DMPO diretamente aos modelos base sem qualquer "lição de casa" prévia (Aprendizado Supervisionado ou SFT) em conjuntos de dados de raciocínio. Essa abordagem "tipo R1-Zero" foi usada para demonstrar claramente o potencial bruto do DMPO.
Aplicação Aprimorada: Eles também aplicaram o DMPO a modelos que já haviam passado pelo SFT para mostrar que ele funciona como um upgrade poderoso para modelos que já foram treinados.

Os resultados foram bastante impressionantes em ambos os setups:

Em um enigma matemático chamado GSM8K, o modelo aplicado diretamente ao base pré-treinado (sem SFT) mostrou melhorias massivas. Especificamente, o DMPO alcançou até 39,63 pontos percentuais de melhoria de precisão sobre as bases de RL não-DMPO anteriores, e uma melhora estrondosa de 67,97 pontos percentuais sobre o próprio modelo base.
Em um enigma de lógica chamado Sudoku, eles viram ganhos enormes, com uma versão de seu modelo saltando de uma taxa de sucesso de 16,41% para 80,86% (uma diferença de mais de 64 pontos!).
Crucialmente, o DMPO continuou a entregar ganhos de desempenho significativos mesmo quando aplicado a modelos que já haviam passado pelo SFT, provando que é um método robusto que funciona quer você comece do zero ou construa sobre treinamentos existentes.

O artigo sugere que o DMPO não é apenas um pequeno ajuste, mas uma mudança fundamental. Ele permite que o robô seja off-policy, o que significa que ele pode aprender com dados antigos de forma eficiente, e forward-only (apenas para frente), o que significa que não precisa de cálculos de retropropagação caros que atrasam outros robôs.

A Conclusão

Os autores estão confiantes de que este método funciona bem para os modelos específicos que testaram (como LLaDA e Dream) nessas tarefas de raciocínio específicas. Eles mediram essas melhorias através de experimentos rigorosos em conjuntos de dados padrão, mostrando que o DMPO pode impulsionar o desempenho, seja aplicado a um modelo que acabou de ser pré-treinado ou a um que já passou por ajuste fino supervisionado. No entanto, eles admitem que ainda não testaram isso em todos os tipos possíveis de modelos ou tarefas, então, embora os resultados sejam fortes, o pleno potencial ainda está sendo explorado.

Em resumo, o DMPO ensina os robôs de difusão a parar de fixar em uma única resposta "perfeita" e, em vez disso, aprender a apreciar todo o cenário de boas soluções, tornando-os solucionadores de problemas mais inteligentes, rápidos e criativos.

Resumo Técnico: Aprimorando o Raciocínio para dLLMs de Difusão via Otimização de Correspondência de Distribuição

Declaração do Problema

Os Modelos de Linguagem de Difusão (dLLMs) oferecem uma alternativa promissora aos modelos de linguagem autorregressivos (AR-LLMs) ao permitir a geração em qualquer ordem e potencialmente um maior rendimento de inferência. No entanto, escalar os dLLMs para igualar as capacidades de raciocínio dos AR-LLMs continua sendo um desafio significativo. Os métodos existentes de ajuste fino por Aprendizado por Reforço (RL), como a Otimização de Gradiente de Política Relativa de Grupo (GRPO), são inadequados para dLLMs devido a dois fatores principais:

Descompasso Computacional: Os dLLMs dependem de processamento bidirecional, tornando a estimativa de log-probabilidades de sequência (necessária para métodos padrão de gradiente de política) computacionalmente dispares e frequentemente intratáveis em comparação com os modelos AR.
Limitações Algorítmicas: As abordagens atuais de RL para dLLMs frequentemente dependem de aproximações heurísticas (ex: decomposição de campo médio) que ignoram correlações entre posições de tokens. Além disso, elas tipicamente focam na maximização de recompensa, o que tende a otimizar para um único modo de recompensa. Esse comportamento de "busca de modo" (mode-seeking) pode levar ao reward hacking e à falha em explorar os diversos caminhos de raciocínio de alta qualidade que os dLLMs são teoricamente capazes de gerar.

Metodologia: Otimização de Correspondência de Distribuição de Política (DMPO)

O artigo propõe a Otimização de Correspondência de Distribuição de Política (DMPO), um framework de ajuste fino de RL fundamentado em Controle Ótimo Estocástico (SOC). Diferente dos métodos tradicionais que maximizam a recompensa esperada, a DMPO visa fazer com que a distribuição da política do modelo corresponda a uma distribuição alvo inclinada pela recompensa.

Estrutura Central

O objetivo é aprender uma política $\pi_\theta$ tal que a distribuição de sequência induzida $p_\theta(o|q)$ corresponda à distribuição ideal $p^*(o|q)$ , definida como:
$p^*(o|q) = \frac{1}{Z(q)} p_{ref}(o|q) e^{r(q,o)/\alpha}$
onde $p_{ref}$ é a política pré-treinada, $r$ é a função de recompensa e $\alpha$ é um parâmetro de temperatura.

Para alcançar isso, os autores formulam o problema como a minimização da divergência Kullback-Leibler Direta (Forward KL) (equivalente à Entropia Cruzada) entre o alvo e a política atual:
$\min_\theta \mathbb{E}_{q \sim D} \mathbb{E}_{p^*(o|q)} \left[ \log \frac{p^*(o|q)}{p_\theta(o|q)} \right]$
Este objetivo de "cobertura de massa" (mass-covering) incentiva o modelo a explorar trajetórias de alta recompensa diversas, em vez de colapsar para um único modo.

Componentes Técnicos Principais

Entropia Cruzada de Denoising Ponderada (WDCE):
Como $p^*$ é intratável (devido à função de partição desconhecida $Z(q)$ e à incapacidade de amostrar diretamente dela), a DMPO emprega Amostragem de Importância. Ela trata as amostras de uma política de referência $p_v$ (uma cópia do modelo com parada de gradiente) como amostras ponderadas de $p^*$ .
A função de perda resultante, WDCE, aproveita o processo único de ruído direto (forward/noising) dos dLLMs. Ela computa a perda usando apenas as amostras limpas finais e seus pesos de importância associados, descartando as trajetórias intermediárias de rollout. Isso torna a WDCE um objetivo off-policy e apenas de ida (forward-only).
Subtração de Baseline de Peso:
Um desafio de implementação crítico identificado é a instabilidade da WDCE ao usar pequenos tamanhos de lote de treinamento. Em lotes pequenos, pesos positivos para respostas "ruins" (que ainda possuem probabilidade não nula sob $p^*$ ) podem dominar, levando a atualizações de gradiente incorretas.
Para resolver isso, os autores propõem a Subtração de Baseline de Peso, onde um peso de baseline $w_{base}$ é subtraído do peso de importância:
$w_{real} = w - w_{base}$
Três estratégias para $w_{base}$ são propostas:
- Baseline de Peso de Grupo: Um valor constante (1) para incentivar a convergência.
- Baseline de Peso Individual: Baseado em uma inclinação de recompensa negativa.
- Baseline de Peso do Modelo: Baseado nos próprios pesos de importância do modelo atual.
  Esta técnica efetivamente injeta gradientes negativos para amostras de baixa recompensa, evitando que o modelo promova respostas ruins em cenários de baixo lote.
Otimização Discriminativa Direta Ponderada (WDDO):
Como alternativa à entropia cruzada, o artigo explora um objetivo inspirado em GAN que incorpora naturalmente gradientes negativos para amostras ruins, aumentando ainda mais a correspondência de distribuição.

Contribuições Principais

Novo Framework de RL: A proposta da DMPO, um framework de correspondência de distribuição projetado especificamente para dLLMs, mudando o paradigma da maximização de recompensa para a correspondência de toda a distribuição de política inclinada pela recompensa.
Fundamentação Teórica: Derivação da perda WDCE usando amostragem de importância e Controle Ótimo Estocástico, fornecendo um método teoricamente sólido para treinamento off-policy sem exigir estimativas caras de log-probabilidade de sequência.
Soluções Práticas: Identificação do modo de falha de pequeno lote em objetivos de KL direta e a introdução da Subtração de Baseline de Peso e WDDO para estabilizar o treinamento e melhorar a eficiência de amostragem.
Eficiência: Demonstração de que a WDCE permite o treinamento apenas via processo direto (forward), permitindo o reuso de rollouts (off-policy) e compatibilidade com técnicas de inferência rápida (ex: cache de KV), acelerando significativamente o treinamento em comparação com métodos baseados em trajetória.

Resultados Experimentais

Os autores avaliaram a DMPO nos modelos LLaDA-8B-Instruct e Dream-7B em quatro benchmarks de raciocínio: GSM8K, MATH500, Countdown e Sudoku.

Ganhos de Desempenho: A DMPO superou consistentemente os modelos base e as linhas de base de RL existentes (incluindo diffu-GRPO, cGRPO e d1).
- No Sudoku, a DMPO-LLaDA alcançou uma melhoria de 39,63 pontos percentuais sobre a melhor linha de base de RL não-DMPO e uma melhoria de 67,97 pontos percentuais sobre o modelo base.
- No Countdown, foram observadas melhorias de até 16,40 pontos percentuais sobre bases fortes.
Sem Necessidade de SFT: O método alcançou esses resultados sem Ajuste Fino Supervisionado (SFT) em conjuntos de dados curados, demonstrando sua capacidade de escalar capacidades de raciocínio diretamente de modelos pré-treinados.
Eficiência de Treinamento: Devido à sua natureza off-policy e computação apenas de ida (forward-only), a DMPO alcançou 2–3x mais eficiência de amostragem e até 8x mais velocidade de amostragem de rollout (quando combinada com inferência de d-LLM rápida) comparado a métodos do tipo GRPO. Ela atingiu os níveis de recompensa do d1 com apenas 31% do orçamento de computação de treinamento.
Estabilidade: O método permaneceu estável mesmo com dados altamente defasados (até 24x de reuso do mesmo lote de rollout), uma propriedade atribuída à estabilidade dos pesos de importância de nível de sequência usados na DMPO.

Significância e Alegações

O artigo afirma que a DMPO representa um passo significativo no escalonamento de dLLMs para tarefas de raciocínio. Ao se afastar do paradigma de maximização de recompensa e dos gargalos computacionais de RL baseado em trajetória, a DMPO aproveita totalmente as características únicas dos modelos de difusão:

Diversidade: O objetivo de KL direta evita o colapso de modo, permitindo que os dLLMs explorem diversos camros de raciocínio.
Eficiência: A natureza off-policy e apenas de ida (forward-only) da WDCE alinha-se com as capacidades de alto rendimento de inferência dos dLLMs, oferecendo um caminho para o ajuste fino de RL escalável que é computacionalmente viável para grandes modelos.
Generalização: O framework é apresentado como uma abordagem principista que pode ser aplicada a vários backbones de dLLM, como evidenciado pelo seu sucesso nas arquiteturas LLaDA e Dream.

Os autores concluem que, embora o trabalho atual foque em modelos e benchmarks específicos, o framework de correspondência de distribuição abre novas direções para melhorar modelos generativos e projetar técnicas de baseline de peso mais eficazes para pesquisas futuras.