Resumo Técnico: GDSD: Aprendizado por Reforço como Auto-Distilação de Desejador Guiado para Modelos de Linguagem de Difusão

1. Declaração do Problema

Modelos de Linguagem de Grande Escala de Difusão (dLLMs) oferecem ganhos significativos de eficiência na inferência em relação aos modelos autoregressivos (ARMs), gerando múltiplos tokens simultaneamente e quebrando a ordem estrita de geração da esquerda para a direita. No entanto, os dLLMs atualmente ficam atrás dos ARMs mais avançados em qualidade de geração, necessitando de métodos eficazes de ajuste fino, como Aprendizado por Reforço (RL).

O obstáculo central para aplicar RL a dLLMs é a intratabilidade da verossimilhança da política. As soluções existentes geralmente se enquadram em duas categorias:

1. Métodos de Verossimilhança de Trajetória: Estimam verossimilhanças exatas acumulando probabilidades de transição. Embora não viesadas, são computacionalmente caras e não alinhadas com os objetivos de pré-treinamento.
2. Métodos de Verossimilhança de Sequência (ELBO): Utilizam o Limite Inferior da Evidência (ELBO) estimado a partir de sequências mascaradas aleatoriamente como um substituto para a verossimilhança real. Estes são computacionalmente eficientes e alinhados com o pré-treinamento.

O Problema Central: O artigo argumenta que os métodos de RL baseados em ELBO sofrem de um viés de Incompatibilidade entre Treinamento e Inferência (TIM). Esse viés surge de duas fontes:

• Lacuna de Verossimilhança: O ELBO é um limite inferior com uma lacuna não desprezível em relação à verossimilhança real, introduzindo viés e variância nas razões de amostragem por importância usadas nos objetivos de Gradiente de Política (PG) ou PPO.
• Incompatibilidade da Distribuição de Amostragem: A inferência do dLLM usa remascaramento iterativo e decodificação por blocos, induzindo uma distribuição de amostragem ($\pi^{rm}_{old}$) que difere da distribuição da política de treinamento ($\hat{\pi}_{old}$) estimada via ELBO.

Essa incompatibilidade faz com que a correção de amostragem por importância seja ineficaz, levando a desempenho degradado e, em casos graves, ao colapso catastrófico do treinamento.

2. Metodologia: Auto-Distilação de Desejador Guiado (GDSD)

Os autores propõem a Auto-Distilação de Desejador Guiado (GDSD), um quadro que reformula o RL para dLLMs como um problema de auto-distilação off-policy, contornando assim a necessidade de estimativa de verossimilhança e evitando o viés TIM.

2.1 Fundamentação Teórica

Sob RL regularizado por KL reverso, a política ótima $\pi^*$ possui uma solução de forma fechada:
$$\pi^*(x_0) \propto \pi_{old}(x_0)^{1-\beta} \cdot \pi_{ref}(x_0)^\beta \cdot \exp(\psi A(x_0))$$
onde $A(x_0)$ é a função de vantagem, $\beta$ é um coeficiente de regularização e $\psi$ é um coeficiente de orientação.

Os autores mostram que essa política ótima induz uma distribuição de desmascaramento guiada (um "desejador professor" $p^*$) para o modelo de difusão mascarado:
$$p^*(x_0|x_t) = p^{ref}_{old}(x_0|x_t) \cdot \exp(\psi A(x_0) - A_t(x_t))$$
onde $p^{ref}_{old}$ é uma mistura geométrica das políticas antiga e de referência, e $A_t(x_t)$ é uma constante de normalização.

2.2 O Objetivo GDSD

Em vez de otimizar um objetivo baseado em verossimilhança, o GDSD distila diretamente o desmascarador atual $p_\theta$ a partir do professor $p^*$, correspondendo aos seus logits. O objetivo minimiza a distância $L_2$ quadrada entre os logits do aluno e do professor:
$$\mathcal{L}_{GDSD} = \mathbb{E} \left[ \left( \log p_\theta(x_0|x_t) - \log p^*(x_0|x_t) \right)^2 \right]$$

2.3 Reformulação Livre de Normalização

Um grande desafio na implementação disso é a constante de normalização intratável $Z_t$ (ou $A_t$) na distribuição do professor. Os autores propõem duas soluções para eliminar isso:

1. Correspondência Direta: Aproveitando a invariância de translação do operador Softmax, a constante pode ser ignorada se os logits do aluno corresponderem aos logits não normalizados do professor até uma constante aditiva.
2. Centralização de Logits ao Nível de Token (TLC): Para evitar deslocamentos de escala de logits descontrolados durante atualizações iterativas, os autores propõem centralizar os logits. Subtraindo a média do logits sobre o vocabulário para cada posição de token, a constante de normalização é totalmente eliminada do objetivo.

O objetivo prático final (com TLC) é:
$$\mathcal{L}_{GDSD} = \mathbb{E} \left[ \left( \log \bar{p}_\theta(x_0|x_t) - \log \bar{p}_{old}(x_0|x_t) - \psi A(x_0) \right)^2 + \beta \left( \log \bar{p}_\theta(x_0|x_t) - \log \bar{p}_{ref}(x_0|x_t) \right)^2 \right]$$
onde $\bar{p}$ denota o modelo centralizado em logits.

3. Contribuições Principais

1. Redução à Auto-Distilação: O artigo estabelece que o RL regularizado por KL reverso para dLLMs é equivalente à distilação de um desmascarador guiado a partir de um auto-professor derivado da política ótima de forma fechada.
2. Quadro GDSD: Um método inovador de auto-distilação livre de verossimilhança que elimina o viés TIM por design, pois não depende de razões de amostragem por importância ou substitutos ELBO.
3. Otimização Livre de Normalização: A introdução da Centralização de Logits ao Nível de Token (TLC) permite RL online eficiente sem o cálculo de funções de partição intratáveis.
4. Unificação Teórica: O artigo demonstra que os métodos existentes baseados em ELBO (por exemplo, wd1, DMPO, SPG, UniGRPO) podem ser vistos como instâncias específicas de distilação usando diferentes medidas de divergência (KL Direto ou KL Reverso), mas sofrem de ineficiência de dados ou viés TIM que o GDSD evita.

4. Resultados Experimentais

Os autores avaliaram o GDSD no LLaDA-8B e no Dream-7B em seis benchmarks:

• Raciocínio Matemático: GSM8K, MATH500
• Planejamento: Countdown, Sudoku
• Codificação: HumanEval, MBPP

Principais Descobertas:

• Ganhos de Desempenho: O GDSD supera consistentemente os métodos baseados em ELBO mais avançados. No modelo Dream-7B, o GDSD alcançou melhorias de precisão no teste de até +19,6% em tarefas de planejamento (Sudoku e Countdown). No LLaDA-8B, as melhorias variaram de +0,6% a +5% em benchmarks de planejamento, matemática e codificação.
• Estabilidade de Treinamento: O GDSD demonstrou dinâmicas de recompensa de treinamento significativamente mais estáveis em comparação com as bases baseadas em ELBO (por exemplo, SPG, ESPO), que exibiram oscilações ou colapso.
• Estudos de Ablação: O componente de Centralização de Logits ao Nível de Token (TLC) mostrou-se capaz de estabilizar ainda mais o treinamento e, em alguns casos (como Sudoku), ajudar o modelo a escapar de mínimos locais para alcançar recompensas finais mais altas.

5. Significado e Alegações

O artigo alega que a auto-distilação direta de desmascarador oferece uma alternativa mais estável e eficaz ao RL baseado em ELBO para dLLMs. Ao remover a dependência de substitutos de verossimilhança, o GDSD evita os vieses estruturais inerentes às abordagens de amostragem por importância. Os resultados sugerem que os ganhos de desempenho observados não são meramente devidos ao ajuste de hiperparâmetros, mas derivam de uma rota de otimização mais principista que alinha o objetivo de treinamento com a distribuição de inferência real dos dLLMs. Os autores concluem que o ELBO não é necessário para RL com dLLMs e que sua remoção leva a capacidades de alinhamento e raciocínio mais robustas.