Resumo Técnico: Gradientes de Política Conscientes de Máscara para Modelos de Linguagem de Difusão

Declaração do Problema

O aprendizado por reforço (RL) tem aprimorado com sucesso as capacidades de raciocínio em Grandes Modelos de Linguagem (LLMs) autorregressivos ao otimizar gradientes de política baseados em log-verossimilhanças tratáveis. No entanto, estender esses métodos para Modelos de Linguagem de Difusão Mascarada (MDLMs) permanece um desafio. Ao contrário dos modelos autorregressivos, que geram tokens sequencialmente da esquerda para a direita, os MDLMs geram texto desmascarando iterativamente tokens de uma sequência totalmente mascarada.

A dificuldade central reside na intratabilidade da estimativa da log-verossimilhança para os MDLMs. Abordagens existentes aproximam essa verossimilhança usando variantes do Limite Inferior de Evidência (ELBO) ou modelando apenas o componente de predição de tokens da trajetória de geração. Esses métodos ignoram um aspecto crítico da geração de MDLM: a ordem na qual as posições são desmascaradas. Nos MDLMs, a trajetória de geração envolve duas decisões distintas em cada etapa: (1) prever quais tokens colocar nas posições mascaradas e (2) decidir quais dessas posições manter desmascaradas (e quais retornar ao estado de máscara). Métios anteriores que ignoram a segunda decisão falham em utilizar todo o sinal disponível no processo de geração para a otimização da política.

Metodologia

Os autores formalizam a geração de MDLM como um Processo de Decisão de Markov (MDP) de duas etapas de ação. Em cada etapa de denoising $t$, o modelo realiza:

1. Predição de Token: Predizer tokens para todas as posições atualmente mascaradas.
2. Seleção de Posição (Remascaramento): Decidir qual subconjunto de posições preditas revelar e quais retornar ao estado [MASK].

Para permitir a otimização por gradiente de política sobre este processo, os autores introduzem o Remascaramento Probabilístico. A inferência padrão de MDLM frequentemente utiliza uma estratégia determinística e gananciosa (greedy) de top-$K$ para selecionar posições com base em pontuações de confiança. Essa seleção gananciosa é não diferenciável, impedindo a otimização baseada em gradiente da ordem de desmascaramento. O método proposto substitui isso por uma variante probabilística:

• Em vez de selecionar deterministicamente as top-$K$ posições, o modelo amostra um subconjunto de posições $U_t$ de uma distribuição Plackett–Luce.
• As probabilidades para esta distribuição são derivadas diretamente dos próprios log-likelihoods de tokens do modelo (logits), escalonados por um parâmetro de temperatura $\tau$.
• Esta formulação cria uma distribuição diferenciável sobre subconjuntos de posições, permitindo que a log-probabilidade da decisão de desmascaramento seja incluída na verossimilhança da trajetória.

Decomposição do Gradiente de Política

Ao modelar a verossimilhança total da trajetória $\pi_\theta(\hat{z}|c)$ como um produto das probabilidades de predição de tokens e das probabilidades de desmascaramento, o gradiente de política naturalmente se decompõe em dois termos:
$$\nabla_\theta J(\theta) = \mathbb{E}_{\hat{z}} \left[ R(c, z_T) \sum_{t=1}^T \left( \underbrace{\nabla_\theta \log \pi_\theta(\hat{z}_t | c, z_{t-1})}_{\text{Gradiente de Token}} + \underbrace{\nabla_\theta \log p_{\text{unmask}}(U_t | \hat{z}_t, z_{t-1}, c)}_{\text{Gradiente de Máscara}} \right) \right]$$

• Termo de Token: Reforça a correção dos tokens preditos.
• Termo de Máscara: Reforça a seleção da ordem de desmascaramento.
• Eficiência: O termo de máscara é computado a partir dos logits existentes do modelo usados para a predição de tokens, exigindo nenhum parâmetro adicional, mudanças arquiteturais ou passagens de forward extras.

Os autores implementam isso usando o algoritmo Group Sequence Policy Optimization (GSPO), clipando as razões de importância para os termos de token e posição separadamente. Para lidar com o custo computacional de avaliar trajetórias completas, eles utilizam a aproximação StepMerge, agrupando etapas de denoising em segmentos para reduzir o número de passagens de forward enquanto preservam a estrutura da trajetória.

Principais Contribuições

1. Formalização de MDLM como um MDP de duas etapas: O artigo fornece um framework teórico mostrando que a geração de MDLM envolve tanto a seleção de tokens quanto a seleção de posições, e que o gradiente de política se decompõe dessa forma.
2. Remascaramento Probabilístico: A derivação de uma distribuição de desmascaramento diferenciável (Plackett–Luce) baseada nos próprios logits do modelo, permitindo a otimização da ordem de desmascaramento sem modificações arquiteturais.
3. Insight Teórico: Demonstração de que ignorar o componente de posição do gradiente de política pode resultar na perda de direções que melhoram o retorno esperado, mesmo quando as probabilidades de tokens permanecem inalteradas.
4. Desempenho de Estado da Arte: Validação empírica mostrando que a otimização conjunta de ambos os termos produz melhorias consistentes sobre métodos baseados em trajetória e ELBO existentes.

Resultados Experimentais

O método foi avaliado em benchmarks de raciocínio matemático (GSM8K, MATH500) e geração de código (HumanEval, MBPP), utilizando o LLaDA-8B-Instruct como modelo base.

• Ganhos de Desempenho: O método proposto alcançou novos resultados de estado da arte em todos os benchmarks.
  • GSM8K: 87,1% de acurácia (em comprimento de geração 128), uma melhoria de +2,5% sobre o melhor anterior (StepMerge).
  • MATH500: 53,4% de acurácia, uma melhoria de +4,0%.
  • HumanEval: 47,1% de acurácia, uma melhoria de +2,2%.
  • MBPP: 53,4% de acurácia, uma melhoria de +2,5%.
• Estudos de Ablação:
  • Tamanho do Bloco: A lacuna de desempenho entre o método proposto e os baselines aumentou conforme o tamanho do bloco de inferência aumentava (de 32 para a sequência completa), sugerindo que o benefício de otimizar a seleção de posição cresce com a complexidade do espaço de decisão de desmascaramento.
  • Eficiência: O método convergiu mais rápido em tempo de relógio (wall-clock time) comparado ao forte baseline baseado em ELBO (SPG), atingindo a acurácia final do SPG em ~15 horas versus ~18 horas, e alcançando uma acurácia final maior (81,0% vs 78,5% no GSM8K).
  • Generalização: A abordagem generalizou-se para um modelo base diferente (Dream-7B) e tarefas de planejamento (Sudoku, Countdown), superando consistentemente os baselines.
• Comparação com DCoLT: Quando comparado ao LLaDOU (que usa um cabeçote separado para seleção de posição), o método proposto alcançou acurácia comparável ou superior com custos de treinamento significativamente menores (~160 horas de GPU vs ~800), pois não requer módulos adicionais.

Significância e Alegações

O artigo afirma que a flexibilidade dos MDLMs vem acompanhada de uma estrutura de trajetória mais rica do que os modelos autorregressivos, especificamente a decisão de quando desmascarar posições. Ao tratar essa decisão como parte da política e otimizá-la via um framework probabilístico diferenciável, os autores demonstram que a otimização conjunta das predições de tokens e da seleção de posição é crucial para maximizar o desempenho de modelos de linguagem de difusão em cenários de RL.

Os autores enfatizam que sua abordagem é praticamente eficiente, pois aproveita os logits existentes do modelo sem exigir novos parâmetros ou mudanças arquiteturais, tornando-a compatível com infraestruturas de treinamento padrão de MDLM, como o StepMerge. Os resultados sugerem que o "gradiente de máscara" fornece um sinal de treinamento significativo que métodos anteriores, focados apenas em verossimilhança de tokens, falharam em capturar.