Resumo Técnico: Ler o Traço, Direcionar o Caminho (CAPR)

Declaração do Problema

Modelos de Linguagem de Difusão (dLLMs) geram texto através da remoção iterativa de ruído de uma sequência totalmente mascarada ao longo de passos paralelos. Diferente dos modelos autorregressivos, os dLLMs expõem uma distribuição preditiva para cada posição mascarada em cada etapa, criando um rico "traço de denoising" que revela quais tokens se tornam confiantes e estáveis versus aqueles que permanecem instáveis ou oscilantes.

Os métodos atuais de pós-treinamento por Aprendizado por Reforço (RL) para dLLMs falham em aproveitar esse traço de forma eficaz, caindo em dois extremos:

1. Rollouts Planos (Flat Rollouts): Atribuem uma recompensa de resultado única para toda a trajetória. Embora computacionalmente baratos, eles não fornecem sinal sobre quais etapas de denoising específicas ou blocos de tokens foram críticos para o resultado final.
2. Rollouts em Árvore (Tree Rollouts): Criam ramificações de continuações alternativas para comparar recompensas e gerar vantagens em nível de segmento. Embora forneçam uma granularidade mais fina, o custo computacional cresce rapidamente com o fator de ramificação e a profundidade, tornando-os frequentemente proibitivamente caros.

A questão central abordada por este trabalho é: Pode o próprio traço de denoising fornecer sinais de supervisão semelhantes a uma árvore sem o custo computacional de uma expansão total da árvore?

Metodologia: CAPR

Os autores introduzem o CAPR (Cached-Amortized Path Refinement), um algoritmo de RL baseado em blocos projetado para extrair vantagens por bloco a partir do traço de denoising. O CAPR opera através de três mecanismos primários:

1. Cache & Steer (Armazenar e Direcionar)

O CAPR registra as etapas iniciais de denoising para construir um estado de caminho ($h_j$) compacto para cada posição $j$. Este estado resume:

• Confiança no token atual de topo.
• Entropia da distribuição de probabilidade.
• Estabilidade do token de topo ao longo dos passos.

Este estado é atualizado via um mecanismo do tipo GRU onde a alta entropia redefine a memória (esquecendo evidências obsoletas) e a estabilidade a preserva. O estado de caminho serve a duas funções:

• Direcionamento (Steering): Ele enviesa os logits para o próximo passo de reversão. Posições estáveis e confiantes recebem um impulso positivo, enquanto posições oscilantes são atenuadas.
• Controle de Desmascaramento (Unmasking Control): Ele influencia quais posições mascaradas são reveladas (desmascaradas) com base em sua certeza, garantindo que o modelo foque em compromissos estáveis.

2. Branch & Prune (Ramificar e Podar)

Para gerar sinais de treinamento sem a expansão total da árvore, o CAPR realiza um único evento de ramificação no ponto médio da trajetória de denoising ($f_b = 0.5$).

• Cache: O modelo armazena em cache o prefixo (logits, tokens preditos e estado de caminho) até o ponto de ramificação.
• Ramificação (Forking): Ele gera um pequeno número de continuações irmãs ($G_b$, ex: 2) aplicando perturbações Gumbel às probabilidades do vocabulário. Esses irmãos reutilizam o prefixo em cache, evitando passagens de forward redundantes.
• Poda (Pruning): Nos limites dos blocos, as ramificações são podadas com base na qualidade do seu estado de caminho (estabilidade média). Apenas as ramificações mais promissoras continuam, impedindo que o espaço de busca exploda em uma árvore completa.

3. Block Critic (Crítico de Bloco)

Uma cabeça de valor leve (Block Critic) lê o estado de caminho em cada limite de bloco.

• Redistribuição de Recompensa: Após o término do rollout, a recompensa de resultado final é redistribuída entre os blocos com base no número de tokens revelados em cada bloco e nas características do estado de caminho.
• Cálculo de Vantagem: Usando TD/GAE (Generalized Advantage Estimation), o crítico converte a recompensa esparsa final em uma vantagem por bloco ($\hat{A}_b$).
• Atualização PPO: Estas vantagens de bloco reponderam a vantagem padrão de nível de rollout usada no PPO. Isso permite que o modelo aprenda quais blocos específicos contribuíram mais para o sucesso, recuperando efetivamente a granularidade da busca em árvore sem o custo da árvore.

Objetivo de Treinamento: O loss total combina o loss de PPO clipado (usando vantagens ponderadas por bloco), divergência KL (para permanecer próximo ao modelo de referência), um regularizador $\chi^2$ (para suprimir picos raros de log-razão), perda de autodestilação (usando rollouts resolvidos como alvos supervisionados) e perda da cabeça de valor.

Contribuições Principais

1. RL Baseado em Traço de Denoising: O artigo propõe uma nova visão onde a confiança e a estabilidade ao longo do caminho de denoising são convertidas em vantagens por bloco, eliminando a necessidade de expansão total de árvore.
2. Algoritmo CAPR: Uma implementação prática que armazena em cache saídas iniciais, amortiza rollouts irmãos, direciona passos de denoising via estados de caminho e utiliza um crítico de bloco para gerar pesos granulares de PPO.
3. Desempenho de Estado da Arte: O CAPR alcança novos resultados de estado da arte para dLLMs ajustados por RL em quatro benchmarks de raciocínio (Sudoku 4x4, Countdown, GSM8K, Math500) em backbones densos e Mixture-of-Experts (MoE) LLaDA com orçamentos de 256 e 512 tokens.

Resultados Experimentais

Os autores avaliaram o CAPR contra baselines incluindo Diffu-GRPO, VRPO, wd1, SAPO, d2-stepMerge, TraceRL e o d-TreeRPO estruturado em árvore.

• Acurácia: No orçamento de 256 tokens, o CAPR superou ou empatou com todos os baselines em todas as tarefas para LLaDA-8B-Instruct e LLaDA-MoE-7B-A1B-Instruct.
  • Sudoku: O CAPR alcançou 93,0% (LLaDA-8B) e 84,8% (LLaDA-MoE), igualando ou excedendo o baseline de árvore mais forte (d-TreeRPO: 92,9% e 78,3%).
  • Countdown: Ganhos significativos foram observados, com o CAPR atingindo 72,7% (8B) e 68,8% (MoE), superando o d-TreeRPO.
• Eficiência: O CAPR reduz significativamente o custo computacional em comparação com rollouts em árvore.
  • Custo de Geração: O CAPR requer aproximadamente 0,75× o custo de rollouts planos e 0,6× o custo de rollouts em árvore (em configurações padrão).
  • Tempo de Execução (Wall-Clock Time): No Sudoku, o CAPR levou 165 segundos por etapa de otimização comparado a 596 segundos para o d-TreeRPO, alcançando acurácia comparável com menos de um terço do computação por etapa.
• Estudos de Ablação:
  • O Block Critic proporcionou o maior ganho intermediário, alterando substancialmente a curva de aprendizado.
  • Branch & Prune adicionou diversidade útil, mas foi menos eficaz isoladamente.
  • Self-Distillation foi crucial para estabilizar o treinamento em grupos de variância zero (rollouts resolvidos) e reduzir picos de gradiente.

Significância e Alegações

O artigo afirma que o CAPR oferece um "meio termo" prático para o RL de dLLM. Ele recupera com sucesso a granularidade da busca em árvore (comparação local de alternativas) enquanto evita a explosão computacional da expansão total da árvore.

• Eficiência vs. Granularidade: O CAPR demonstra que o próprio traço de denoising contém informações suficientes para atribuir recompensas ao nível de bloco, removendo a necessidade de rollouts caros baseados em árvore.
• Generalização: O método transfere-se efetivamente de backbones densos para esparsos (MoE) sem necessidade de ajuste específico da tarefa.
• Praticidade: Ao reduzir o custo por etapa para cerca de 0,6× dos rollouts em árvore enquanto iguala ou supera sua acurácia, o CAPR torna o treinamento de RL de alta granularidade viável para modelos maiores e contextos longos.

Os autores observam limitações, reconhecendo que o método é atualmente avaliado em tarefas de raciocínio automaticamente verificáveis e pode exigir adaptação para seguimento de instruções de código aberto ou domínios com recompensas ruidosas e não locais. Além disso, o crítico de bloco é aprendido online, o que pode apresentar desafios de confiabilidade para cadeias de raciocínio muito longas.