SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models

1. O Problema: Alinhamento de RL em Modelos de Difusão Discretos

Os Modelos de Linguagem de Difusão (dLLMs), como o MDLM (Masked Diffusion Language Model), emergiram como uma alternativa eficiente aos modelos autoregressivos (AR) devido à sua capacidade de decodificar múltiplos tokens em paralelo, reduzindo a latência de inferência. No entanto, alinhar esses modelos com preferências humanas ou recompensas específicas de tarefas através de Aprendizado por Reforço (RL) apresenta um desafio fundamental:

• Intratabilidade da Log-Verossimilhança: Ao contrário dos modelos autoregressivos, onde a probabilidade de uma sequência é o produto das probabilidades condicionais (facilmente computável), a log-verossimilhança exata ($\log \pi_\theta(x|c)$) em modelos de difusão é computacionalmente intratável.
• Vieses nas Aproximações Atuais: Trabalhos anteriores tentaram contornar isso usando o Limite Inferior da Evidência (ELBO) como um substituto para a log-verossimilhança no gradiente de política.
  • O Erro Crítico: O ELBO é apenas um limite inferior ($ELBO \leq \log \pi_\theta$). Maximizar o ELBO para recompensas positivas é válido, mas minimizar o ELBO para recompensas negativas (ou vantagens negativas) não garante a minimização da verdadeira log-verossimilhança. Isso impede que o modelo aprenda efetivamente com feedback negativo (penalizando saídas ruins) e introduz um viés significativo no gradiente da política, limitando o desempenho em tarefas complexas de raciocínio.

2. Metodologia: Sandwiched Policy Gradient (SPG)

Os autores propõem o SPG (Sandwiched Policy Gradient), um novo algoritmo de RL projetado para resolver o problema da estimativa de log-verossimilhança em dLLMs. A ideia central é "sanduichar" a log-verossimilhança intratável utilizando limites superiores e inferiores baseados na recompensa.

A. O Objeto de Otimização "Sanduíche"

O algoritmo divide as amostras geradas em dois grupos baseados em sua vantagem (diferença entre a recompensa da amostra e a média do grupo):

1. Amostras com Vantagem Positiva ($A_j \geq 0$): O algoritmo maximiza o ELBO (limite inferior). Isso é consistente com a maximização da log-verossimilhança real.
2. Amostras com Vantagem Negativa ($A_j < 0$): Em vez de minimizar o ELBO (o que é ineficaz), o algoritmo minimiza um Limite Superior da Evidência (EUBO - Evidence Upper Bound).
   • Isso garante que, ao penalizar saídas ruins, o modelo esteja efetivamente reduzindo a probabilidade real dessas saídas, corrigindo o viés das abordagens anteriores.

A função objetivo combinada é:
$$J_{SPG}(\theta) = \mathbb{E} \left[ \frac{1}{g} \sum_{j=1}^g \left( \mathbb{1}_{A_j \geq 0} \cdot A_j \mathcal{L}_{ELBO} + \mathbb{1}_{A_j < 0} \cdot A_j \mathcal{L}_{EUBO} \right) \right]$$

B. Limite Superior da Evidência (EUBO)

O EUBO é derivado com base no Limite Variacional de Rényi. Para uma sequência $x$ e um hiperparâmetro $\beta \geq 1$, o limite superior é dado por:
$$\mathcal{L}_{EUBO}(x; \theta) = \frac{1}{\beta} \sum_{i=1}^n \log \mathbb{E}_{t, z_t} \left[ w(t) \cdot \mathbb{1}(z_{t,i} = m) \cdot \pi_\theta^\beta(x_i | z_t) \right] + C(T)$$

• Diferente do ELBO, o logaritmo está fora da expectativa, o que torna a estimativa via Monte Carlo enviesada, mas teoricamente um limite superior válido.

C. Estratégia de Mascaramento em Blocos (Block-Wise Masking)

Para estimar os limites (ELBO e EUBO) de forma estável e eficiente durante o rollout de RL:

• Em vez de mascarar tokens aleatoriamente (como em métodos anteriores), o SPG utiliza uma estratégia block-wise.
• A sequência é dividida em blocos. Um bloco aleatório é selecionado para mascaramento, mantendo os blocos anteriores limpos e os subsequentes totalmente mascarados.
• Isso alinha melhor a distribuição de dados entre o processo de geração (inferência) e a otimização, aumentando a estabilidade do treinamento.

D. Mistura de Limites (Mixture Approach)

Para equilibrar a estabilidade e a precisão, especialmente para traços de vantagem negativa, o SPG utiliza uma mistura ponderada do EUBO e do ELBO:
$$\mathcal{L}_{Mix} = \omega \cdot \mathcal{L}_{EUBO} + (1-\omega) \cdot \mathcal{L}_{ELBO}$$

• Teoria: A Proposição 1 demonstra que essa mistura reduz estritamente a variância do gradiente em comparação com o uso isolado de qualquer um dos limites.
• Prática: O termo EUBO atua como um sinal de correção forte para amostras ruins, enquanto o ELBO fornece estabilidade. A mistura evita gradientes que desaparecem ou explodem.

3. Contribuições Principais

1. Algoritmo SPG: Um novo método de gradiente de política que elimina o viés na otimização de dLLMs ao utilizar limites variacionais "sanduichados" (inferior para recompensas positivas, superior para negativas).
2. Derivação do EUBO: Uma formulação teórica rigorosa do Limite Superior da Evidência para modelos de difusão mascarados, permitindo a penalização correta de saídas indesejadas.
3. Técnica de Mascaramento em Blocos: Uma estratégia prática que melhora a estabilidade da estimativa de Monte Carlo e a eficiência da otimização, alinhando-se com estratégias de inferência semiautoregressivas modernas.
4. Redução de Variância: Demonstração teórica e empírica de que a mistura de limites reduz a variância do gradiente, levando a um treinamento mais estável.

4. Resultados Experimentais

O SPG foi avaliado em quatro benchmarks de raciocínio matemático e lógico, além de tarefas de codificação, utilizando o modelo base LLaDA-8B-Instruct.

Desempenho em Raciocínio (Comparado ao SOTA anterior):

• GSM8K (Matemática): +3.6% de acurácia (atingindo 86.1%).
• MATH500 (Matemática Avançada): +2.6% de acurácia (atingindo 40.0%).
• Countdown (Raciocínio Lógico): +18.4% de acurácia (atingindo 70.7%).
• Sudoku (Raciocínio Lógico): +27.0% de acurácia (atingindo 94.0%).

Desempenho em Codificação:

• HumanEval: +1.9% de melhoria.
• MBPP: +4.7% de melhoria.

Análises Adicionais:

• Convergência: O SPG demonstrou convergência mais rápida e níveis de recompensa mais altos em comparação com baselines como D1, WD1 e UniGRPO.
• Robustez: O modelo treinado com SPG manteve desempenho superior sob diversas estratégias de inferência (diferentes tamanhos de bloco, decodificação aleatória vs. baseada em confiança), demonstrando forte generalização.
• Eficiência: O custo computacional por atualização de gradiente é comparável aos métodos existentes, pois o custo dominante permanece no rollout de geração, não na estimativa de Monte Carlo.

5. Significado e Impacto

Este trabalho é significativo por resolver uma barreira teórica fundamental na aplicação de RL a modelos de linguagem baseados em difusão.

• Correção de Viés: Ao demonstrar que a minimização do ELBO é insuficiente para penalizar saídas ruins, o SPG oferece uma solução teoricamente fundamentada para o alinhamento de dLLMs.
• Avanço de Estado da Arte: Os resultados mostram que dLLMs, quando adequadamente otimizados com RL, podem superar ou igualar o desempenho dos modelos autoregressivos em tarefas de raciocínio complexo, mantendo a vantagem de inferência paralela.
• Generalidade: A abordagem não é limitada a um único modelo, sugerindo que o uso de limites variacionais "sanduichados" pode ser um padrão para o futuro treinamento por reforço de modelos generativos não autoregressivos.

Em resumo, o SPG transforma os modelos de difusão de linguagem em candidatos viáveis e de alto desempenho para aplicações que exigem raciocínio complexo e alinhamento com preferências humanas, superando as limitações de otimização que anteriormente restringiam seu uso em RL.