Resumo Técnico: Orientação de Recompensa Consciente de Entropia para Alinhamento de Modelos de Linguagem de Difusão

Declaração do Problema

A orientação de recompensa (ou amostragem posterior) provou ser altamente eficaz para modelos de difusão contínuos, onde feedback diferenciável de um modelo de recompensa permite o refinamento iterativo dos passos de geração. No entanto, aplicar esse paradigma a Modelos de Linguagem Grandes de Difusão Discreta (dLLMs) apresenta um desafio fundamental: os dLLMs geram tokens discretos, que são não diferenciáveis.

As abordagens existentes tentam preencher essa lacuna, mas enfrentam uma compensação crítica entre confiabilidade do modelo de recompensa e precisão do gradiente:

1. Relaxamento Contínuo (Expectativa): Alimentar o modelo de recompensa com embeddings "suaves" (valores esperados das probabilidades dos tokens) permite o fluxo de gradiente. No entanto, os modelos de recompensa são tipicamente treinados em tokens rígidos; assim, entradas suaves e fora do vocabulário levam a sinais de recompensa pouco confiáveis e desempenho degradado.
2. Estimador Straight-Through (STE): Métodos como o APS (Rout et al., 2025c) avaliam a recompensa em tokens rígidos amostrados (garantindo confiabilidade), mas propagam gradientes como se a entrada fosse suave. Isso cria uma incompatibilidade entre onde a recompensa é avaliada e onde os gradientes são aplicados, levando a erros de aproximação, particularmente quando a entropia preditiva do modelo é alta (ou seja, durante os primeiros passos de remoção de ruído).

Metodologia: EntRGi

Os autores introduzem o EntRGi (Orientação de Recompensa Consciente de Entropia), um mecanismo que interpola dinamicamente entre relaxações de tokens contínuos e tokens rígidos amostrados em uma base token a token. A inovação central é usar a entropia preditiva por token do próprio modelo de difusão para determinar o peso de interpolação.

O Mecanismo

Em cada posição mascarada $l$ durante o processo de remoção de ruído, o EntRGi constrói um embedding de entrada $e_l^{EntRGi}$ para o modelo de recompensa como uma combinação convexa:
$$e_l^{EntRGi} = (1 - w_l)e_l^{soft} + w_l e_l^{hard}$$
Onde:

• $e_l^{soft}$ é o embedding suave esperado derivado da distribuição de probabilidade do modelo.
• $e_l^{hard}$ é o embedding de um token rígido amostrado.
• $w_l$ é um peso adaptativo derivado da entropia do token $H(q_l)$, especificamente $w_l = H(q_l) / \log K$.

Lógica de Adaptação:

• Baixa Entropia (Alta Confiança): Quando o modelo está confiante (baixa entropia), $w_l \to 0$. A entrada é predominantemente o embedding suave, permitindo um fluxo de gradiente limpo sem erro de aproximação significativo, pois os embeddings suave e rígido estão próximos.
• Alta Entropia (Alta Incerteza): Quando o modelo está incerto (alta entropia), $w_l \to 1$. A entrada desloca-se para o embedding do token rígido. Isso garante que o modelo de recompensa receba entradas realistas e dentro da distribuição, preservando a confiabilidade da recompensa, embora o erro de aproximação do gradiente seja maior.

Essa abordagem resolve a tensão confiando em embeddings suaves quando o modelo está confiante e revertendo para tokens discretos quando a incerteza é alta.

Pós-Treinamento: RGRL

O artigo também propõe o RGRL (Aprendizado por Reforço Guiado por Recompensa), uma receita de pós-treinamento que utiliza o EntRGi (ou APS) para gerar conclusões guiadas por recompensa. Essas conclusões são usadas para ajustar finamente o dLLM via auto-distilação, criando um loop denso de feedback recompensa-gradiente, em vez de depender exclusivamente de recompensas escalares usadas em métodos de RL padrão como o diffu-GRPO.

Contribuições Principais

1. EntRGi: Um mecanismo inovador de orientação de recompensa consciente de entropia para dLLMs de difusão discreta que equilibra adaptativamente a precisão do gradiente e a confiabilidade do modelo de recompensa.
2. Superioridade na Adaptação em Tempo de Inferência: Validação empírica mostrando que o EntRGi supera o estado da arte anterior (APS) e outras linhas de base (Expectativa, Best-of-N) em múltiplos benchmarks, particularmente em temperaturas de amostragem mais altas onde a entropia é maior.
3. Pós-Treinamento RGRL: Uma nova receita de treinamento que utiliza amostras guiadas por recompensa para ajuste fino. Os autores demonstram que o RGRL-EntRGi supera significativamente o diffu-GRPO (um algoritmo de RL amplamente adotado para dLLMs), alcançando até uma melhoria relativa de 70% nas pontuações de recompensa em conjuntos de dados específicos.
4. Robustez a Incompatibilidade de Tokenizadores: O método permanece eficaz mesmo quando o dLLM e o modelo de recompensa possuem tokenizadores incompatíveis (por exemplo, LLaDA vs. Qwen/Llama), um cenário comum do mundo real onde métodos baseados em gradiente existentes frequentemente lutam.

Resultados Experimentais

Os autores avaliaram o EntRGi em modelos de 7 bilhões de parâmetros (Dream-v0-Instruct-7B e LLaDA-8B-Instruct) em cinco conjuntos de dados multi-habilidade e quatro modelos de recompensa.

• Desempenho em Tempo de Inferência: O EntRGi alcançou consistentemente pontuações de recompensa Top@1 e Avg@4 mais altas em comparação com o APS e a Expectativa. Notavelmente, enquanto o desempenho do APS degradou-se significativamente em temperaturas mais altas ($\tau=0.7$), o EntRGi manteve ou melhorou o desempenho, demonstrando sua capacidade de lidar efetivamente com estados de alta entropia.
• Eficiência de Pós-Treinamento: O RGRL-EntRGi mostrou convergência mais rápida e pontuações de recompensa finais mais altas em comparação com o diffu-GRPO. Os ganhos foram mais pronunciados em conjuntos de dados com recompensas iniciais mais baixas (por exemplo, WildChat-IF), sugerindo que o método é particularmente eficaz para tarefas de alinhamento difíceis.
• Estudos de Ablação: A análise confirmou que o ponderamento baseado em entropia é superior a heurísticas alternativas (como distância L2-norm) e que o método reduz o erro de aproximação durante os primeiros passos de remoção de ruído, onde a entropia é mais alta.

Significado

O artigo afirma que o EntRGi fornece uma solução principial para o conflito "discreto vs. diferenciável" em modelos de linguagem de difusão. Ao aproveitar a incerteza do modelo como um sinal para regular a incorporação de feedback, o EntRGi permite:

1. Direcionamento Mais Forte em Tempo de Inferência: Permitir que dLLMs otimizem objetivos específicos de tarefa sem retreinamento, com confiabilidade comparável a modelos de difusão contínua.
2. Alternativas Viáveis de Pós-Treinamento: Desbloquear o feedback denso recompensa-gradiente como uma alternativa mais eficaz ao ajuste fino supervisionado e ao RL de recompensa escalar para dLLMs.

Os autores posicionam este trabalho como um passo necessário para tornar a adaptação guiada por recompensa e o pós-treinamento robustos e escaláveis para a classe emergente de modelos de linguagem de difusão discreta.