Resumo Técnico: Aprendendo com o Próprio Futuro (d-OPSD)

1. Definição do Problema

Embora a Auto-distilação On-policy (OPSD) tenha se mostrado eficaz para o pós-treinamento de Grandes Modelos de Linguagem Autoregressivos (AR-LLMs), sua aplicação a Grandes Modelos de Linguagem de Difusão (dLLMs) permanece inexplorada. Os frameworks de OPSD existentes são inerentemente projetados para modelos AR, criando dois desajustes fundamentais quando aplicados a dLLMs:

1. Construção de Professor Subotimizada: A OPSD baseada em AR injeta informações privilegiadas (ex: soluções de referência) como um prefixo ao prompt, baseando-se na condicionamento da esquerda para a direita $p(\text{suffix}|\text{prefix})$. Isso conflita com os dLLMs, que geram sequências de forma não autoregressiva e suportam geração de ordem arbitrária, incluindo a capacidade de condicionar em sufixos $p(\text{prefix}|\text{suffix})$. Além disso, soluções de referência estáticas não se alinham com a natureza "on-policy" da auto-distilação.
2. Nível de Supervisão Incompatível: Modelos AR utilizam supervisão de divergência ao nível de token baseada na predição do próximo token. Em contraste, os dLLMs operam via um processo de denoising iterativo onde múltiplos tokens são preditos simultaneamente em cada etapa, e apenas um subconjunto é revelado enquanto outros permanecem mascarados. A supervisão ao nível de token é, portanto, incompatível com as transições de estado passo a passo dos dLLMs.

Consequentemente, há uma necessidade de um framework de OPSD customizado que aproveite a mecânica generativa única dos dLLMs para permitir a auto-melhoria eficiente sem depender de modelos professores externos mais fortes.

2. Metodologia: d-OPSD

Os autores propõem o d-OPSD (diffusion On-Policy Self-distillation), o primeiro framework de OPSD projetado especificamente para dLLMs. O método opera em um único modelo $p_\theta$ atuando tanto como aluno quanto como professor, utilizando duas inovações principais:

A. Construção de Auto-Professor via "Auto-Futuro"

Em vez de anexar soluções de referência estáticas como um prefixo, o d-OPSD aproveita a capacidade de condicionamento bidirecional dos dLLMs.

• Mecanismo: O aluno amostra uma trajetória completa $Y = \{y_T, \dots, y_0\}$ de um estado totalmente mascarado até uma resposta limpa $y_0$.
• Entrada do Professor: O professor é construído revelando parcialmente tokens da própria resposta final gerada pelo aluno como uma condição de sufixo. Especificamente, em cada etapa de denoising $t$, a entrada do professor $y_{\text{teacher},t}$ consiste na sequência ruidosa atual com um subconjunto de índices $S_t$ (selecionados da resposta final) revelados.
• Conceito: Isso permite que o aluno aprenda com sua própria "experiência futura" (a resposta autogerada) condicionada como um sufixo, aderindo estritamente ao paradigma on-policy onde todos os dados são gerados pelo próprio aluno.

B. Supervisão de Divergência ao Nível de Passo

Reconhecendo que os dLLMs decodificam via predição do próximo passo em vez da predição do próximo token, os autores deslocam o objetivo de supervisão do nível de token para o nível de passo.

• Mecanismo: Em cada etapa de denoising $t$, o modelo calcula as distribuições para todos os tokens. A supervisão foca exclusivamente nos top-$k$ tokens mais confiantes entre as posições atualmente mascaradas (o subconjço $K_t$ que dita a transição de estado de $t$ para $t-1$).
• Objetivo: A função de perda minimiza a divergência de Kullback-Leibler (KL) entre a distribuição do aluno e a distribuição do professor sobre este subconjunto específico $K_t$ em cada etapa:
  $$L_t = \frac{1}{|K_t|} \sum_{i \in K_t} D_{KL}(P_{\text{student},t}^i \parallel P_{\text{teacher},t}^i)$$
  O objetivo total agrega esta perda de nível de passo ao longo de toda a trajetória de $T$ passos.

3. Principais Contribuições

1. Identificação de Subotimalidade: O artigo identifica que as formulações de OPSD existentes centradas em AR (condicionamento de prefixo e supervisão de nível de token) estão fundamentalmente desalinhadas com os dLLMs.
2. Novo Framework (d-OPSD): Os autores introduzem o primeiro framework de OPSD para dLLMs, apresentando:
   • Uma construção de auto-professor baseada em sufixo que utiliza respostas autogeradas como informações privilegiadas, permitindo o aprendizado a partir da "experiência de auto-futuro".
   • Supervisão de divergência ao nível de passo que se alinha com a mecânica de denoising iterativo de modelos de difusão.
3. Validação Empírica: Experimentos extensos demonstram que o d-OPSD permite que um único modelo impulsione sua própria auto-melhoria, superando as linhas de base de Aprendizado por Reforço com Recompensas Verificáveis (RLVR) e Ajuste Fino Supervisionado (SFT).

4. Resultados Experimentais

Experimentos foram conduzidos em quatro benchmarks de raciocínio: GSM8K, MATH500 (raciocínio matemático) e Countdown, Sudoku (planejamento). O modelo base utilizado foi o LLaDA-8B-Instruct.

• Desempenho de Raciocínio: O d-OPSD supera consistentemente ou iguala as linhas de base de RLVR (ex: diffu-GRPO) e SFT em todas as tarefas e comprimentos de sequência. Por exemplo, no GSM8K (256 tokens), o d-OPSD alcançou 81,0% de acurácia comparado a 79,8% para o diffu-GRPO e 78,8% para o SFT.
• Eficiência de Amostragem: O d-OPSD demonstra uma eficiência de amostragem vastamente superior, convergindo em aproximadamente 10% das etapas de otimização exigidas pelo RLVR. Por exemplo, no GSM8K, o d-OPSD exigiu apenas 425 passos comparado aos 7.700 do diffu-GRPO.
• Comparação com OPSD Estilo AR: Um estudo de ablação substituindo o professor baseado em sufixo por um professor baseado em prefixo estilo AR mostrou que a construção de sufixo proposta supera significativamente a contraparte AR. A análise das métricas de "Overlap Top-K" revelou que a abordagem estilo AR falha em introduzir novos conhecimentos (alto overlap), enquanto a abordagem de sufixo do d-OPSD fornece uma gama adequada de novos padrões de pensamento para o aluno aprender.
• Descobertas de Ablação:
  • Objetivo de Divergência: A KL reversa supera a KL direta (Forward KL).
  • Razão de Retenção: Uma razão de retenção ($\rho_{\text{teacher}}$) de 0,25 produziu resultados ótimos, embora mesmo 0,10 tenha proporcionado ganhos significativos.
  • Seleção de Top-k: Derivar o subconjunto top-$k$ a partir da distribuição do professor rendeu melhor desempenho do que a partir da distribuição do aluno.
  • Estabilidade: O clipping de KL pontual foi considerado essencial para evitar o colapso da política após o pico de desempenho, um modo de falha observado tanto no RLVR quanto no d-OPSD sem clipping.

5. Significância e Alegações

O artigo afirma abrir um caminho promissor para o pós-treinamento de dLLMs ao estabelecer que a auto-distilação pode ser efetivamente adaptada para modelos de difusão. A significância reside em:

• Eficiência: Fornecer um método que alcança desempenho de raciocínio de estado da arte com custo computacional drasticamente reduzido (eficiência de amostragem) em comparação ao RLVR.
• Alinhamento Teórico: Unir com sucesso a lacuna entre o aprendizado on-policy e a mecânica de difusão ao reestruturar a informação privilegiada como sufixos de "auto-futuro" e a supervisão como divergência de nível de passo.
• Auto-Melhoria: Demonstrar que um único modelo pode servir como seu próprio professor sem a necessidade de um modelo externo mais forte ou conjuntos de dados de referência estáticos, desde que a trajetória de geração seja corretamente utilizada.

Os autores observam que, embora eficaz, a abordagem compartilha a suscetibilidade do RLVR ao colapso da política após atingir o desempenho máximo, sugerindo a necessidade de trabalhos futuros em técnicas de estabilização.