Resumo Técnico: Otimização de Política Robusta para Prevenir Esquecimento Catastrófico

Declaração do Problema

Modelos de Linguagem de Grande Escala (LLMs) são tipicamente treinados por meio de pipelines de múltiplos estágios, começando com Aprendizado por Reforço a partir de Feedback Humano (RLHF) para alinhar os modelos com segurança e utilidade, seguidos por ajuste fino para tarefas específicas a jusante (por exemplo, raciocínio matemático, programação). Uma vulnerabilidade crítica neste paradigma é o esquecimento catastrófico: mesmo atualizações menores a jusante podem degradar comportamentos previamente aprendidos, como barreiras de segurança.

Pesquisas anteriores focaram amplamente em estratégias de mitigação no momento do ajuste fino a jusante (por exemplo, reensaio, ajuste fino eficiente em parâmetros, regularização ou mesclagem de modelos) para preservar capacidades após o fato. No entanto, os autores argumentam que esses métodos são insuficientes porque dependem de procedimentos específicos a jusante. Se um modelo for posteriormente submetido a um ajuste fino "ordinário" que se desvie desses protocolos específicos, os comportamentos previamente aprendidos permanecem frágeis. O desafio central é garantir que a própria política base seja robusta a adaptações futuras, independentemente da tarefa ou algoritmo a jusante utilizado.

Metodologia: Otimização de Política Robusta para Ajuste Fino (FRPO)

Os autores propõem a Otimização de Política Robusta para Ajuste Fino (FRPO), um framework projetado para otimizar a política base quanto à estabilidade contra mudanças futuras a jusante.

1. Formulação do Problema: Robustez Max-Min

Os autores formalizam o problema como uma otimização max-min. Em vez de maximizar a recompensa esperada apenas na política atual $\pi_\theta$, o objetivo busca maximizar a recompensa mínima alcançável por qualquer política a jusante $Q$ que esteja dentro de uma vizinhança limitada por KL de $\pi_\theta$.

O problema de otimização é definido como:
$$\max_{\pi_\theta} \left( \inf_{Q: \mathbb{E}[\text{KL}(Q \parallel \pi_\theta)] \le \rho} \mathbb{E}_{x,y \sim Q}[r(x,y)] - \beta \text{KL}(\pi_\theta \parallel \pi_{\text{ref}}) \right)$$

Aqui, $\rho$ representa o raio da vizinhança KL, refletindo restrições padrão no ajuste fino a jusante (por exemplo, regiões de confiança, taxas de aprendizado pequenas ou restrições LoRA). A restrição é imposta em expectativa sobre a distribuição de prompts, permitindo mudanças significativas em prompts relevantes para a tarefa, mantendo a proximidade média em relação à política base.

2. Derivação Teórica

Usando a teoria da dualidade para otimização robusta distribucional (especificamente sob uma restrição de f-divergência), os autores derivam uma forma dual do problema de minimização interna.

• Objetivo de Risco Entrópico: A formulação dual revela que o objetivo robusto é equivalente a otimizar um objetivo de risco entrópico. A minimização interna sobre $Q$ transforma o problema em maximizar:
  $$-\lambda \log \mathbb{E}_{y \sim \pi_\theta} \left[ e^{-r(x,y)/\lambda} \right]$$
  onde $\lambda$ é um hiperparâmetro ajustável que controla a sensibilidade ao risco.
• Interpretação: Este objetivo penaliza trajetórias de baixa recompensa mais severamente do que a maximização padrão de recompensa esperada. Ele efetivamente busca regiões "planas" na paisagem de recompensa dentro do espaço de políticas, garantindo que pequenas perturbações (ajuste fino a jusante) não levem a quedas bruscas de desempenho.
• Caso Limite: À medida que $\lambda \to \infty$, o objetivo recupera o objetivo padrão GRPO (Otimização de Política Relativa em Grupo), que maximiza a recompensa esperada.

3. Implementação Algorítmica

Os autores adaptam o algoritmo GRPO para otimizar este novo objetivo sem introduzir sobrecarga computacional extra (por exemplo, nenhum crítico aprendido ou passagens forward adicionais).

• Amostragem por Importância: A expectativa sob a nova política é estimada usando amostras de uma política defasada $\pi_{\text{old}}$ via amostragem por importância.
• Recorte e Linhas de Base: Para garantir estabilidade, o algoritmo emprega recorte estilo PPO nas razões de importância. Crucialmente, os autores derivam um termo de linha de base para cancelar a alta variância no gradiente da função de log-partição quando $\lambda$ é grande.
• Redução de Viés: A estimativa de Monte Carlo da função de log-partição introduz um viés de $O(1/G)$ (onde $G$ é o tamanho do grupo). Os autores empregam uma técnica jackknife (estimação leave-one-out) para reduzir esse viés para $O(1/G^2)$, garantindo convergência estável mesmo com tamanhos de grupo pequenos.

Contribuições Principais

1. Framework RLHF Robusto: O artigo introduz uma formulação max-min que considera todas as mudanças a jusante dentro de uma bola KL. Demonstra que essa formulação é matematicamente equivalente à otimização de um objetivo de risco entrópico, fornecendo uma base teórica para a planicidade da recompensa no espaço de políticas.
2. Algoritmo FRPO: Uma substituição prática e direta para o GRPO que não requer computação adicional. Integra-se perfeitamente com pipelines RLHF existentes, inclui uma linha de base derivada para estabilidade e recupera o GRPO padrão quando $\lambda \to \infty$.
3. Validação Empírica: O método é avaliado em diversos cenários, incluindo alinhamento de segurança e raciocínio matemático, demonstrando preservação superior de capacidades anteriores em comparação com ajuste fino padrão e outras bases de robustez.

Resultados Experimentais

Os autores avaliam o FRPO em dois cenários principais: Treinamento de Segurança e Raciocínio Matemático.

1. Treinamento de Segurança (Barreiras de Segurança)

• Configuração: Modelos (Mistral-v0.1 e Qwen2.5) foram treinados em um conjunto de dados de segurança usando FRPO com $\lambda$ variável, e depois ajustados finamente em tarefas a jusante (Alpaca para seguimento de instruções, GSM8K para matemática e UltraFeedback para RL).
• Descobertas:
  • Modelos treinados com FRPO ($\lambda \approx 0.2$) mantiveram taxas de recusa de segurança significativamente mais altas e pontuações de nocividade mais baixas (StrongREJECT) após o ajuste fino em comparação com o GRPO padrão.
  • O método preservou efetivamente as barreiras de segurança mesmo sob configurações agressivas de ajuste fino (SFT de parâmetros completos, várias taxas de aprendizado, LoRA).
  • Análise de Trade-off: Embora um $\lambda$ menor tenha preservado melhor a segurança, um $\lambda$ maior melhorou ligeiramente a utilidade à custa da segurança. O $\lambda$ ótimo foi encontrado para ser dependente da tarefa, mas geralmente baixo para segurança.
  • Desempenho a Jusante: Crucialmente, o FRPO não comprometeu o desempenho da tarefa a jusante (por exemplo, precisão no GSM8K, pontuações MMLU) em comparação com o GRPO, confirmando que a retenção de segurança não foi devido a uma falha em aprender a nova tarefa.

2. Aprendizado Contínuo (Matemática para Código)

• Configuração: Modelos foram primeiro treinados em raciocínio matemático (conjunto de dados MATH) e depois ajustados finamente em geração de código (OpenCodeInstruct).
• Descobertas:
  • FRPO com $\lambda = 2.0$ (otimizado para a maior variância nas recompensas de matemática) preservou a precisão matemática significativamente melhor do que o GRPO.
  • Especificamente, o FRPO manteve 22% de precisão a mais no benchmark MATH500 após o ajuste fino de código em comparação com o GRPO, enquanto alcançou ganhos comparáveis em benchmarks de codificação (MBPP+).

3. Estudos de Ablação

• Comparação com SAM: Os autores compararam o FRPO com Minimização Sensível à Acurácia (SAM). Embora o SAM (operando no espaço de parâmetros) tenha oferecido alguma melhoria sobre o GRPO, o FRPO (operando no espaço de política/KL) foi superior. Os autores observam que a bola L2 do SAM no espaço de parâmetros não captura totalmente as mudanças de política a jusante, enquanto a bola KL restringe diretamente as mudanças na distribuição de saída.
• Papel de $\lambda$: Experimentos confirmaram que $\lambda$ controla o trade-off entre recompensa média e variância. O $\lambda$ ótimo correlaciona-se com o desvio padrão das recompensas e o orçamento KL $\rho$.

Significado e Alegações

O artigo alega que a robustez ao ajuste fino a jusante deve ser incorporada diretamente na política base durante o RLHF, em vez de depender de intervenções posteriores. Ao otimizar para estabilidade de recompensa dentro de uma vizinhança limitada por KL, o FRPO identifica políticas que residem em regiões de "recompensa plana", tornando-as inerentemente resistentes ao esquecimento catastrófico.

Os autores posicionam o FRPO como uma solução de propósito geral que:

• Funciona em diversos domínios (segurança, matemática, codificação).
• Não requer custo computacional adicional em comparação com o GRPO padrão.
• Fornece uma abordagem principista e matematicamente fundamentada para aprendizado contínuo em LLMs que é agnóstica ao protocolo específico de ajuste fino a jusante.

O trabalho sugere que o princípio de otimizar para robustez de adaptação futura pode se estender além do RLHF para pré-treinamento e ajuste fino supervisionado, oferecendo uma nova direção para pesquisa em alinhamento e aprendizado contínuo.