Patch the Distribution Mismatch: RL Rewriting Agent for Stable Off-Policy SFT

Título: Patch the Distribution Mismatch: RL Rewriting Agent for Stable Off-Policy SFT

Autores: Jiacheng Wang, Ping Jian, Zhen Yang, et al. (Beijing Institute of Technology)

1. O Problema: Esquecimento Catastrófico e Desalinhamento de Distribuição

O artigo aborda um desafio fundamental no ajuste fino supervisionado (SFT) de Grandes Modelos de Linguagem (LLMs): o esquecimento catastrófico. Quando um modelo pré-treinado é ajustado para uma tarefa de domínio específico (downstream), ele frequentemente perde capacidades gerais adquiridas anteriormente.

Os autores identificam que a raiz desse problema reside na natureza "off-policy" do SFT downstream:

• Descompasso de Distribuição: Os dados de demonstração de especialistas (downstream) muitas vezes pertencem a uma distribuição diferente da distribuição de geração natural do modelo base (que é estilo "Resposta a Perguntas" ou QA, condicionada apenas à entrada $x$).
• Instabilidade de Otimização: Quando o modelo tenta aprender dados que são de baixa probabilidade sob sua própria política atual, os gradientes podem ser amplificados (devido a fatores de reamostragem inversa), levando a atualizações instáveis e a um desvio excessivo da política original.
• Limitações das Abordagens Atuais: Métodos existentes de "reescrita de dados" (data rewriting) tentam mitigar isso reescrevendo as demonstrações para se alinharem melhor ao modelo. No entanto, eles geralmente:
  1. Usam prompts fixos e templates rígidos, o que leva a uma colapso de diversidade (as reescritas tornam-se homogêneas e artificiais).
  2. Alinham-se a uma distribuição condicional induzida pelo prompt de reescrita ($\pi_0(\cdot | x, y^*, x_{prompt})$), e não à distribuição natural de geração QA do modelo ($\pi_0(\cdot | x)$), falhando em fechar completamente a lacuna de distribuição.

2. Metodologia: Agente de Reescrita Baseado em RL

Para resolver essas limitações, os autores propõem tratar a reescrita de dados como um problema de aprendizado por política e utilizam Aprendizado por Reforço (RL) para otimizar essa política.

Arquitetura do Sistema

O framework opera em duas etapas principais:

1. Treinamento do Agente de Reescrita (Stage I):

   • Um agente de reescrita ($R_\phi$) é treinado sobre um modelo base congelado ($\pi_0$) usando adaptadores LoRA (Low-Rank Adaptation) para eficiência de parâmetros.
   • O agente recebe uma entrada $x$, uma solução de especialista $y^*$ e um prompt de reescrita, gerando uma reescrita $\tilde{y}$.
   • Otimização: Utiliza o algoritmo GRPO (Group Relative Policy Optimization), que elimina a necessidade de uma rede de valor separada, estimando a linha de base a partir de grupos de amostras.

2. Construção do Dataset e SFT Downstream (Stage II):

   • O agente treinado gera um novo dataset reescrito.
   • Aplica-se um pipeline de Geração-Verificação-Retorno (Generate-Verify-Fallback): se a reescrita passar em uma verificação de consistência de tarefa, ela é usada; caso contrário, a demonstração original é mantida.
   • O modelo final é ajustado (SFT) neste dataset reescrito.

Objetivo Unificado e Função de Recompensa

O núcleo da inovação é a função de recompensa unificada que otimiza três objetivos simultaneamente, mas com restrições rigorosas:

1. Consistência da Tarefa (Hard Gate):

   • Uma verificação binária ($r_{task}$) garante que a reescrita mantenha a resposta correta e a validade do raciocínio.
   • Mecanismo de Portão Rígido: As recompensas de alinhamento e diversidade só são aplicadas se a reescrita passar neste teste. Se falhar, a recompensa é zero. Isso evita que o agente aprenda a gerar textos "bonitos" mas incorretos.

2. Alinhamento Distribucional (QA-Style Alignment):

   • Recompensa ($r_{dist}$) baseada na verossimilhança (likelihood) da reescrita sob a distribuição natural do modelo base $\pi_0(\cdot | x)$ (apenas condicionado à entrada, sem o prompt de reescrita).
   • O objetivo é fazer com que as reescritas pareçam respostas naturais que o modelo geraria espontaneamente, reduzindo o choque de distribuição.

3. Diversidade (Anti-Colapso):

   • Recompensa ($r_{div}$) que penaliza a homogeneidade. Calcula a distância semântica entre múltiplas reescritas viáveis para a mesma entrada.
   • Incentiva o agente a explorar diferentes formas válidas de expressar a solução, evitando o colapso de modo típico de templates fixos.

3. Contribuições Principais

• Insight sobre Descompasso de Distribuição: Identificação de que métodos de reescrita atuais falham em alinhar-se à distribuição de geração QA natural do modelo, limitando sua eficácia contra o esquecimento.
• Agente de Reescrita via RL: Proposta de um agente aprendido (em vez de baseado em prompts fixos) que otimiza explicitamente o alinhamento distribucional e a diversidade.
• Objetivo Unificado com Portão Rígido: Uma estratégia de recompensa que garante a correção da tarefa antes de otimizar a qualidade estilística e a diversidade, estabilizando o treinamento.
• Eficiência: Uso de LoRA e GRPO para treinar o agente de reescrita de forma leve sobre modelos base congelados.

4. Resultados Experimentais

Os experimentos foram realizados em modelos de diferentes escalas (Llama-3.2-1B, 3B e Mistral-7B) focando em raciocínio matemático como tarefa downstream.

• Desempenho em Tarefa Downstream: O método proposto alcançou ganhos em matemática comparáveis ou superiores ao SFT padrão (Vanilla SFT). Por exemplo, no Mistral-7B, obteve um aumento de 55.23% (vs. 54.19% do SFT padrão).
• Redução do Esquecimento Catastrófico: O método demonstrou uma retenção significativamente superior em benchmarks de domínio geral (MMLU, MMLU-Pro, AGIEval) após o ajuste fino.
  • Enquanto o SFT padrão causou uma queda de 19.17% no desempenho geral (Mistral-7B), o método proposto reduziu essa queda para apenas 7.35%.
  • Em média, a redução do esquecimento foi de 12.34% em relação às abordagens baseadas em SFT padrão.
• Análise de Perda de Treinamento: As curvas de perda mostram que o SFT realizado sobre dados reescritos pelo agente começa com uma perda menor e converge para um platô mais estável, indicando que os dados são mais fáceis de ajustar sob o objetivo de verossimilhança padrão.
• Ablação: A remoção da recompensa de alinhamento distribucional causou a maior queda na retenção geral, confirmando que alinhar-se à distribuição de geração natural é crucial para mitigar o esquecimento.

5. Significado e Conclusão

O trabalho oferece uma solução elegante e eficaz para o dilema estabilidade-plasticidade no ajuste fino de LLMs. Ao invés de restringir as atualizações do modelo (como fazem métodos de "trust-region"), o método atua na fonte dos dados, "consertando" (patching) a distribuição dos dados de treinamento para que sejam compatíveis com a política atual do modelo.

Isso permite que o modelo aprenda novas habilidades de domínio específico sem "esquecer" suas capacidades gerais, superando as limitações de métodos anteriores baseados em templates rígidos. A abordagem demonstra que o aprendizado por reforço pode ser aplicado de forma eficaz para otimizar a qualidade dos dados de treinamento, não apenas a política do modelo final.