Resumo Técnico: Máscaras de Divergência Preditiva para RL de LLMs

1. Definição do Problema

O Aprendizado por Reforço (RL) para Grandes Modelos de Linguagem (LLMs) é inerentemente off-policy devido a incompatibilidades entre os motores de inferência usados para os rollouts e as pilhas de treinamento, bem como pelo envelhecimento da política (policy staleness). Para estabilizar essas atualizações off-policy, métodos padrão como o Proximal Policy Optimization (PPO) empregam máscaras de região de confiança (trust-region masks). Essas máscaras baseiam-se em dois critérios para decidir se devem mascarar (zerar) o gradiente de um token:

1. Critério de Proximidade: Determina se a política se moveu demais em relação à política de comportamento (behavior policy).
2. Critério de Direção: Determina se a atualização empurra a política para mais longe da política de comportamento (direção externa) ou se a puxa de volta (corretiva).

Trabalhos recentes, especificamente o DPPO, melhoraram o critério de proximidade ao substituir a razão de importância de amostra única por uma divergência distributiva (ex: divergência KL) estimada sobre o top-K do vocabulário. No entanto, o DPPO mantém o critério de direção original do PPO, que depende exclusivamente da razão de importância do token amostrado ($r_t - 1$).

Os autores identificam uma inconsistência fundamental nesta abordagem híbrida: enquanto o critério de proximidade é agora definido por uma divergência distributiva global, o critério de direção permanece como um proxy de amostra única. Os autores observam que o sinal da mudança na razão amostrada pode discordar do sinal da mudança real na divergência. Consequentemente, um token pode ser mantido (porque a razão sugere uma atualização corretiva) mesmo quando a atualização na verdade aumenta a divergência global, ou vice-versa.

2. Metodologia

O artigo propõe a Máscara de Divergência Preditiva (Predictive Divergence Mask), um método que alinha o critério de direção com a natureza distributiva do critério de proximidade. Em vez de perguntar "A razão amostrada se afastará de uma?", o método pergunta: "O próximo passo de gradiente aumentará a divergência?".

Derivação Teórica

Para as políticas softmax usadas em LLMs autorregressivos, os autores derivam uma expressão de forma fechada para a derivada direcional da divergência de Kullback-Leibler (KL), denotada por $\dot{D}$, ao longo de um passo de gradiente de política.
$$\dot{D} = \underbrace{(\pi_k - \mu_k)}_{\text{Termo Local}} + \underbrace{\sum_{i} \pi_i (\mu_i - \pi_i)}_{\text{Termo Global}}$$
Onde:

• $\pi_k$ e $\mu_k$ são as probabilidades do token amostrado $k$ sob as políticas de treinamento e de comportamento, respectivamente.
• O Termo Local corresponde exatamente à quantidade usada no critério de direção baseado em razão padrão ($r - 1$).
• O Termo Global captura o acoplamento de normalização em todo o vocabulário. Este termo é invisível para uma razão amostrada única, mas é crucial para determinar a mudança real na divergência.

A máscara preditiva mantém uma atualização apenas se o sinal de $\hat{A}_t \cdot \dot{D}_t$ for não-positivo (indicando um movimento corretivo ou neutro), desde que o token já esteja fora da região de confiança ($D_t > \delta$).

Estimativa Prática

Motores de rollout de produção tipicamente expõem apenas as probabilidades dos tokens top-$K$, não a distribuição completa do vocabulário. Para computar $\dot{D}$ sob essa visão truncada, os autores propõem dois estimadores leves para a cauda não observada da distribuição:

1. Estimador de Cauda Agregada (Aggregated-Tail Estimator): Colapsa toda a massa de probabilidade residual em um único token agregado.
2. Estimador de Cauda Uniforme (Uniform-Tail Estimator): Distribui a massa residual uniformemente entre os tokens de vocabulário não vistos.

Ambos os estimadores requerem apenas reduções simples sobre as probabilidades do top-$K$ retidas e as massas de cauda, adicionando um custo computacional negligenciável.

3. Principais Contribuições

• Identificação de Inconsistência: O artigo destaca que misturar um critério de proximidade distributivo (divergência) com um critério de direção de amostra única (razão) cria um desalinhamento onde a máscara de direção pode falhar em rastrear a verdadeira mudança na divergência.
• Máscara de Divergência Preditiva: Uma nova regra de mascaramento que utiliza a mudança de primeira ordem da própria divergência para determinar a direção da atualização.
• Derivação de Forma Fechada: A derivação da derivada direcional para políticas softmax, separando explicitamente os componentes local (amostrado) e global (normalização).
• Estimadores de Vocabulário Truncado: Dois estimadores de baixo custo e baixa latência (cauda agregada e uniforme) que permitem que o método funcione com as restrições de visualização top-$K$ padrão de produção.

4. Resultados Experimentais

Os autores avaliaram o método no conjunto de dados DAPO-Math-17k usando modelos Qwen3 (4B, 8B e 30B) com várias configurações de precisão (FP8).

• Estabilidade de Treinamento: Em todas as escalas de modelo e configurações de precisão (incluindo testes de estresse com FP8), a Máscara de Divergência Preditiva melhorou a estabilidade do treinamento em comparação com o baseline DPPO-TopK-KL. Em contraste, o GRPO padrão com truques de clip-higher mostrou instabilidade e colapso em configurações de modelos maiores.
• Alinhamento do Critério de Direção: A análise ao nível de token revelou que o critério de direção baseado em divergência está melhor alinhado com a mudança realizada na divergência do que o critério baseado em razão.
  • Entre os tokens onde os dois critérios discordaram, o critério baseado em divergência teve uma taxa de "manutenção insegura" (unsafe-keep) menor (34,2% vs 36,9%), significando que foi menos propenso a manter atualizações que de fato aumentavam a divergência.
  • Por outro lado, ele manteve uma fração maior de atualizações que sucederam em contrair a divergência (65,1% vs 62,9%).
• Robustez: O método demonstrou melhor robustez ao hiperparâmetro de limiar da região de confiança ($\delta$), performando melhor que o baseline mesmo sob restrições mais rigorosas ($\delta = 0,05$).
• Equivalência de Estimadores: Os estimadores de cauda agregada e de cauda uniforme produziram resultados quase idênticos, sugerindo que o método não é sensível à suposição específica de modelagem da cauda.

5. Significância e Alegações

O artigo afirma que a Máscara de Divergência Preditiva resolve uma inconsistência estrutural nos métodos atuais de RL baseados em divergência para LLMs. Ao garantir que tanto o critério de proximidade quanto o de direção sejam derivados da mesma divergência distributiva, o método fornece um controle de região de confiança mais fiel.

Os autores enfatizam que, embora a melhoria nas taxas de "manutenção insegura" seja modesta (devido ao pequeno conjunto de tokens onde os critérios discordam), a consistência é teoricamente sólida e empiricamente benéfica em diversas escalas. Eles concluem que o critério de direção deve receber o mesmo tratamento distributivo que o critério de proximidade para aproveitar totalmente os benefícios das regiões de confiança baseadas em divergência. O método é apresentado como um substituto direto (drop-in replacement) para implementações existentes de DPPO que não exige novos hiperparâmetros além do limiar de divergência existente.