Reformulate LLM Reinforcement Learning for Efficient Training under Black-box Discrepancy

Resumo Técnico: Reformulação do Aprendizado por Reforço de LLMs para Treinamento Eficiente sob Discrepância de Caixa-Preta

1. Declaração do Problema

O Aprendizado por Reforço (RL) tornou-se um paradigma crítico de pós-treinamento para aprimorar as capacidades de raciocínio de Grandes Modelos de Linguagem (LLMs). No entanto, o pós-treinamento por RL frequentemente sofre de desempenho subótimo imprevisível ou colapsos totais de treinamento. Investigações recentes atribuem essas falhas a uma onipresente discrepância entre treino e inferência (ou desajuste).

Essa discrepância surge porque as demandas industriais frequentemente forçam os praticantes a usar motores subjacentes e formatos de precisão díspares (por exemplo, diferentes kernels, estratégias de paralelização ou precisões numéricas como BF16 vs. FP8) para os exatos mesmos pesos do modelo durante o treinamento versus a inferência. Isso cria um desvio oculto entre as distribuições de probabilidade de ação da política de treinamento ($\pi_{Train}$) e da política de inferência ($\pi_{Inf}$).

As estratégias de mitigação existentes possuem limitações significativas:

• Alinhamento de Precisão Uniforme: Simplesmente reduzir a precisão do treinamento para corresponder à inferência (por exemplo, para FP16) reduz a faixa de valores representáveis e exacerba os riscos de explosão de gradiente.
• Intervenções Heurísticas: Métodos como o mascaramento de gradientes de tokens com alta discrepância apenas suprimem o impacto de dados ruidosos. Eles não abordam a causa raiz e negligenciam o objetivo final: otimizar a política para a implantação de inferência, não apenas para o processo de treinamento.
• Alinhamento de Infraestrutura: Tentar alinhar operadores de baixo nível e pilhas de hardware é trabalhoso e muitas vezes inviável devido à miríade de casos de borda.

O desafio central é permitir que a política de treinamento perceba dinamicamente e anule ativamente sua discrepância com a política de inferência sem sacrificar o objetivo primário de maximização de recompensa.

2. Metodologia: MDP com Restrição de Discrepância (DCMDP)

Os autores propõem uma mudança de paradigma do Processo de Decisão de Markov (MDP) padrão para um MDP com Restrição de Discrepância (DCMDP). Este framework trata a discrepância entre treino e inferência como um problema de otimização de caixa-preta, permitindo que o sistema lide com configurações heterogêneas de treinamento e inferência de forma transparente.

Componentes Principais:

• Métrica de Discrepância:
  Os autores avaliam empiricamente três sinais de discrepância ao nível de token: Diferença de Probabilidade ($\delta_{diff}$), Razão de Probabilidade ($\delta_{ratio}$) e Diferença de Log-Probabilidade ($\delta_{log}$). Eles identificam que a diferença absoluta de probabilidade ($\delta_{diff} = |p_{Train} - p_{Inf}|$) é o sinal de penalidade mais eficaz. Ela oferece uma regularização "plástica" que corrige o desvio relevante para a implantação sem restringir excessivamente a exploração impulsionada pela recompensa, ao contrário das penalidades baseadas em razão, que podem ser conservadoras demais para tokens de baixa probabilidade.

• Região de Tolerância de Discrepância:
  Uma descoberta empírica crítica é a existência de uma região de tolerância de discrepância. Penalizar agressivamente discrepâncias menores e negligenciáveis no início do treinamento suprime a exploração da política e reduz a eficiência do aprendizado. Por outro modo, reduzir a discrepância excessiva fora desta região melhora a consistência da otimização.

  • Mecanismo: A formulação DCMDP aplica penalidades apenas quando a discrepância excede um limiar predefinido $c$. Dentro desta região de tolerância, a política é permitida para explorar livremente para maximizar recompensas.

• Otimização de Objetivo Duplo via Relaxação Lagrangiana:
  O problema é formulado como a maximização da recompensa esperada sujeita a uma restrição na discrepância média ao nível de token:
  $$\max_{\theta} J(\theta) = \mathbb{E}_{\tau \sim \pi_{\theta}} \left[ \sum r(o_t) \right]$$
  $$\text{sujeito a } \frac{1}{|\tau|} \sum |\pi_{Train}(o_t) - \pi_{Inf}(o_t)| \leq c$$

Para resolver isso, os autores empregam Relaxação Lagrangiana. Um multiplicador aprendível $\lambda$ ajusta dinamicamente o peso da penalidade de discrepância:
* Se a discrepância violar a fronteira de tolerância ($>c$), $\lambda$ aumenta, fortalecendo a penalidade para puxar a política de volta à zona segura.
* Se a discrepância estiver dentro da fronteira, $\lambda$ decai para zero, permitindo que a política foque na maximização da recompensa.

O objetivo final (Dual-GRPO) injeta esta penalidade adaptativa diretamente na vantagem normalizada por grupo:
$$\hat{A}^{DCMDP}_{i,t} = \hat{A}_{i,t} - \lambda_k \cdot \delta_{i,t}$$

3. Principais Contribuições

1. Compensação Autônoma de Discrepância: O artigo demonstra que a incorporação de uma penalidade robusta de caixa-preta (diferença absoluta de probabilidade ao nível de token) confere à política de treinamento a capacidade de erradicar autonomamente o risco de colapso de treinamento induzido por discrepância.
2. Mudança de Paradigma para DCMDP: Os autores reformulam o RL de LLM como um MDP com Restrição de Discrepância. Esta abordagem equilibra de forma ideal os objetivos duplos de maximizar a recompensa da tarefa e minimizar a discrepância entre treino e inferência, indo além da simples escalonagem de objetivos.
3. Paradigma de Treinamento Heterogêneo: O método permite uma configuração de treinamento onde os modelos são otimizados em ambientes de alta fidelidade e intensivos em recursos (ex: precisão BF16), enquanto são explicitamente alinhados para implantações de inferência de baixo custo e recursos limitados (ex: FP8). Isso transforma a heterogeneidade de infraestrutura de um passivo em um controle de otimização gerenciável.

4. Resultados Experimentais

O método foi avaliado nos modelos Qwen3-8B (denso) e Qwen3-30B-A3B (Mistura de Especialistas - MoE) em seis benchmarks de raciocínio matemático (AIME24/25, AMC23, MATH-500, Minerva, OlympiadBench).

• Treinamento/Inferência BF16: O DC-GRPO superou consistentemente o baseline vanilla GRPO em todos os benchmarks para ambos os modelos, alcançando pontuações médias mais altas (ex: Qwen3-8B melhorou de 67.3 para 71.0). Isso confirma que impor restrições de discrepância aumenta, em vez de prejudicar, a otimização da recompensa primária.
• Treinamento Heterogêneo (BF16) + Implantação (FP8): Nesta configuração, o vanilla GRPO sofreu degradação significativa de desempenho ou colapso devido ao gap amplificado entre treino e inferência. O DC-GRPO mitigou isso com sucesso, preservando e muitas vezes excedendo o desempenho do baseline de apenas BF16.
• Estudos de Ablação: O método provou ser robusto a variações no multiplicador Lagrangiano inicial ($\lambda_0$) e na taxa de aprendizado dupla ($\eta_\lambda$). O orçamento de tolerância ($c$) mostrou-se crítico; defini-lo muito baixo causou super-regularização, enquanto defini-lo muito alto fez o método degenerar em um vanilla GRPO.

5. Significância e Alegações

O artigo alega que o DCMDP fornece um framework algorítmico unificado para mitigar a instabilidade entre treino e inferência sem exigir otimizações manuais e exaustivas das pilhas de hardware ou software subjacentes. Ao tratar a discrepância como um sinal de feedback de caixa-preta, o método permite:

• Treinamento Estável: Prevenindo os colapsos abruptos comuns no pós-treinamento de RL.
• Exploração Eficiente: Permitindo que a política explore livremente dentro de uma região de discrepância segura.
• Eficiência de Recursos: Permitindo o uso de configurações de treinamento de alta fidelidade para simular e adaptar-se a ambientes de implantação restritos, facilitando assim implantações de modelos reais mais eficientes.

Os autores reconhecem limitações, observando que seu escopo empírico é atualmente limitado às precisões BF16 e FP8 e famílias de modelos específicas, sendo necessário trabalho futuro para estender o método a outras combinações de backend e arquiteturas de modelos.