Resumo Técnico: Repensando a Regularização de Divergência em RL para LLMs

1. Declaração do Problema

O aprendizado por reforço (RL) é um pilar fundamental no pós-treinamento de grandes modelos de linguagem (LLMs) para alinhá-los às preferências humanas e aprimorar suas capacidades de raciocínio. No entanto, o treinamento de RL para LLMs modernos é inerentemente off-policy devido a discrepâncias entre os motores de treinamento e de inferência (discrepâncias numéricas) e à prática de dividir os rollouts em múltiplos mini-batches. Isso torna necessário o uso de Otimização de Política de Região de Confiança (TRPO) para garantir atualizações estáveis, restringindo a divergência entre a política atual ($\pi$) e a política de comportamento ($\mu$).

Os métodos convencionais enfrentam duas limitações primárias neste contexto:

1. Restrições Baseadas em Razão (PPO, GRPO, SPO): Estes métodos utilizam a razão de importância $r_t = \pi(y_t|s_t)/\mu(y_t|s_t)$ para definir a região de confiança. Em LLMs com vocabulários grandes e de cauda longa, a razão é um substituto ruim para o deslocamento de distribuição. Um pequeno aumento na probabilidade de um token de baixa probabilidade pode gerar uma razão enorme com um deslocamento de massa negligenciável, enquanto uma mudança moderada na razão de um token de alta probabilidade pode induzir mudanças significativas na política. Isso leva à super-restrição de tokens raros e à sub-restrição de tokens comuns.
2. Mascaramento Rígido (DPPO): Trabalhos recentes como o Divergence Proximal Policy Optimization (DPPO) abordam o problema da razão utilizando uma máscara de divergência (especificamente Binary-TV, que mede o deslocamento absoluto de probabilidade $|\pi - \mu|$). No entanto, o DPPO impõe a região de confiança via uma máscara binária rígida. Uma vez que um token cruza o limite em uma direção prejudicial, seu gradiente é descartado (definido como zero). Isso evita divergências adicionais, mas falha em fornecer um sinal corretivo para puxar a política de volta para a região de confiança, levando a uma otimização frágil próximo ao limiar.

2. Metodologia: Otimização de Política Regularizada por Divergência (DRPO)

Os autores propõem o DRPO, um método que substitui a máscara rígida do DPPO por um regularizador quadrático suave e ponderado pela vantagem sobre o deslocamento da política, preservando a mesma geometria da região de confiança Binary-TV.

Formulação Central

O DRPO deriva seu objetivo reescrevendo a restrição Binary-TV $|\pi(y_t|s_t) - \mu(y_t|s_t)| \leq \delta$ como um limite de razão adaptável ao token. Em vez de realizar o clipping, ele aplica um regularizador suave ao objetivo substituto:

$$\mathcal{L}_{DRPO}(x, \pi) = \mathbb{E}_{y \sim \mu} \left[ \sum_{t=1}^{|y|} \left( r_t \hat{A}_t - \frac{|\hat{A}_t|}{2\delta} \mu(y_t|s_t) (r_t - 1)^2 \right) \right]$$

Onde:

• $r_t$ é a razão de importância.
• $\hat{A}_t$ é a estimativa de vantagem.
• $\delta$ é o limiar da região de confiança.
• O termo $\mu(y_t|s_t)$ escala o regularizador, convertendo a penalidade de uma forma $\chi^2$ baseada em razão (como no SPO) para uma forma de deslocamento de probabilidade absoluta ($\ell_2^2$).

Comportamento do Gradiente

O gradiente do objetivo DRPO introduz um peso contínuo $w_t$ para o gradiente da política de cada token:
$$w_t = 1 - \text{sign}(\hat{A}_t(r_t - 1)) \frac{|\pi(y_t|s_t) - \mu(y_t|s_t)|}{\delta}$$

Este peso exibe três comportamentos principais:

1. Dentro da Região de Confiança: O peso é positivo, mas suavemente atenuado conforme o deslocamento se aproxima de $\delta$, permitindo que as atualizações continuem, porém com magnitude reduzida.
2. Fora da Região de Confiança (Divergindo): Quando a atualização se afasta ainda mais da política de comportamento ($|\pi - \mu| > \delta$), o peso torna-se negativo. Isso inverte o gradiente, fornecendo um sinal corretivo que puxa a política de volta para o limite, ao contrário do gradiente zero das máscaras rígidas.
3. Atualizações de Convergência: Quando a atualização retorna em direção à política de comportamento, o peso é amplificado ($>1$), incentivando um retorno suave à região de confiança.

Distinções Teóricas

• vs. SPO: O SPO utiliza um regularizador do tipo $\chi^2$ (dividindo por $\mu$), o que faz com que os pesos do gradiente cresçam sem limites na cauda longa de tokens de baixa probabilidade. O DRPO utiliza um regularizador do tipo $\ell_2^2$ (multiplicando por $\mu$), garantindo que os pesos do gradiente permaneçam limitados independentemente da probabilidade do token.
• vs. DPPO: O DRPO preserva a geometria Binary-TV do DPPO, mas substitui a máscara descontínua por um regularizador suave e diferenciável.

3. Principais Contribuições

1. Algoritmo DRPO: Um novo método de otimização de política que substitui máscaras de divergência rígidas por um regularizador quadrático suave baseado no deslocamento absoluto de probabilidade (Binary-TV).
2. Pesos de Gradiente Limitados: Demonstração teórica e empírica de que os pesos do gradiente do DRPO permanecem limitados em vocabulários de cauda longa, evitando a instabilidade causada por métodos baseados em razão (PPO, SPO), onde os pesos podem explodir para tokens de baixa probabilidade.
3. Mecanismo Corretivo: Introdução de um sinal corretivo contínuo para atualizações que excedem a região de confiança, prevenindo a "morte do gradiente" observada em abordagens de máscara rígida como o DPPO.
4. Princípios de Design: Identificação de três critérios práticos para regularizadores eficazes em RL para LLMs:
   • Induzir uma fronteira estável alinhada com o deslocamento de distribuição (não apenas razões).
   • Manter os pesos de gradiente por token limitados em vocabulários de cauda longa.
   • Fornecer sinais corretivos suaves quando a política diverge.

4. Resultados Experimentais

Os autores avaliaram o DRPO em diversos cenários, incluindo:

• Modelos: Qwen3-4B, Qwen3-30B-A3B, Qwen3.5-35B-A3B e DeepSeek-R1-Distill-Qwen-1.5B.
• Arquiteturas: Densas e Mixture-of-Experts (MoE).
• Precisão: BF16, FP8 (apenas rollout) e FP8 (end-to-end).
• Benchmarks: AIME 2024 e AIME 2025 (tarefas de raciocínio matemático).

Principais Descobertas:

• Estabilidade: O DRPO permite consistentemente um treinamento estável em todos os cenários. Em contraste, métodos baseados em razão (GRPO, SPO) sofreram frequentemente de instabilidade ou colapso, particularmente em configurações de baixa precisão (FP8).
• Desempenho: O DRPO igualou ou superou a melhor acurácia de avaliação de todos os baselines. Superou métodos de máscara rígida (DPPO) em velocidade de convergência e acurácia final.
• Ablações:
  • A remoção da ponderação de vantagem absoluta $|\hat{A}_t|$ do regularizador levou a quedas de desempenho e instabilidade, confirmando que o limite da região de confiança deve ser independente da escala de recompensa.
  • Medidas de divergência alternativas (penalidades KL, TV) tiveram desempenho inferior ao DRPO, sugerindo que a geometria específica do regularizador Binary-TV é crucial.
  • Aplicar o regularizador apenas fora da região de confiança (Mask-DRPO) resultou em desempenho comparável ao DRPO total, indicando que o principal ganho vem da regularização corretiva fora do limite.

5. Significância e Alegações

O artigo afirma que a geometria do gradiente induzida por um regularizador é mais crítica do que a medida de divergência nominal no objetivo. Enquanto trabalhos anteriores frequentemente focam em divergências padrão (KL, JS), este trabalho demonstra que:

• Restrições baseadas em razão são fundamentalmente falhas para LLMs devido aos vocabulários de cauda longa.
• Máscaras rígidas, embora melhores que as razões, introduzem fragilidade ao descartar sinais corretivos.
• DRPO oferece uma solução unificada que combina a correção geométrica de restrições baseadas em divergência (Binary-TV) com a estabilidade de otimização de regularizadores suaves.

Os autores posicionam o DRPO não apenas como um algoritmo específico, mas como uma lente empírica para projetar objetivos de otimização de política estáveis, enfatizando que regularizadores eficazes devem fornecer sinais de gradiente limitados, contínuos e corretivos para lidar com os desafios únicos do RL para LLMs.