RLHF in an SFT Way: From Optimal Solution to Reward-Weighted Alignment

1. O Problema

O Alinhamento de Grandes Modelos de Linguagem (LLMs) com valores humanos, geralmente realizado através de RLHF (Reinforcement Learning from Human Feedback), enfrenta desafios significativos:

• Complexidade e Custo Computacional: Métodos tradicionais baseados em amostragem online, como PPO (Proximal Policy Optimization) e GRPO (Group Relative Policy Optimization), exigem múltiplos modelos (política, referência, crítico e recompensa) operando simultaneamente, o que consome muitos recursos e é difícil de implementar.
• Instabilidade de Treinamento: Mesmo com simplificações recentes como o DPO (Direct Preference Optimization), que elimina a necessidade de um modelo de recompensa explícito durante o treinamento, problemas como overfitting e instabilidade persistem. O DPO e métodos relacionados (como A-LoL) frequentemente dependem de clipping (limitação) dos pesos de importância ou utilizam pesos negativos para respostas indesejadas, o que pode levar a paisagens de otimização não limitadas e instáveis.
• Dilema Eficiência vs. Desempenho: Existe uma lacuna entre métodos offline rápidos (como SFT/DPO) e métodos online robustos (como PPO), onde os primeiros muitas vezes não alcançam o desempenho ótimo de alinhamento, e os segundos são muito custosos.

2. Metodologia: VAR (Variational Alignment with Re-weighting)

Os autores propõem o VAR, uma nova formulação que trata o alinhamento como um problema de inferência variacional, transformando o objetivo do RLHF em uma forma de Ajuste Fino Supervisionado (SFT) ponderado por recompensa, sem necessidade de amostragem online.

Fundamentos Teóricos

• Solução Ótima Fechada: Partindo da solução ótima conhecida do RLHF (Equação de Bellman/Equação 3 no papel), que é dada por $\pi^*(y|x) \propto \pi_{ref}(y|x) \exp(r(x,y)/\beta)$.
• Minimização de Divergência KL: Em vez de otimizar a política diretamente via gradiente de política, o VAR minimiza a divergência KL direta (forward KL) entre a política ótima $\pi^*$ e a política aprendida $\pi_\theta$.
  • O objetivo torna-se: $\max_\theta \mathbb{E}_{\pi^*}[\log \pi_\theta(y|x)]$.
• Reformulação como SFT Ponderado: Ao aplicar importance sampling usando a distribuição de referência $\pi_{ref}$, o objetivo é reescrito como:
  $$\mathbb{E}_{\pi_{ref}} \left[ \frac{\pi^*(y|x)}{\pi_{ref}(y|x)} \log \pi_\theta(y|x) \right] = \mathbb{E}_{\pi_{ref}} \left[ \underbrace{\frac{\exp(r(x,y)/\beta)}{Z(x)}}_{\text{Peso Ponderado por Recompensa}} \cdot \log \pi_\theta(y|x) \right]$$
  Onde $Z(x)$ é a função de partição (normalização).

Inovações Práticas

1. Pesos Não-Negativos: A transformação exponencial da recompensa ($\exp(r/\beta)$) garante que os pesos sejam estritamente positivos. Isso elimina a instabilidade causada por pesos negativos encontrados em outros métodos, criando uma paisagem de otimização bem definida e estável.
2. Estimativa de Partição In-Batch: O cálculo de $Z(x)$ (soma sobre todo o espaço de saída) é computacionalmente proibitivo. O VAR propõe uma estimativa eficiente de Monte Carlo dentro de cada micro-batch de treinamento:
   • Em vez de amostrar novos dados online, usa-se as amostras já presentes no batch de treinamento para estimar a normalização.
   • Isso permite uma implementação escalável e prática, sem necessidade de um modelo crítico ou amostragem online.

3. Principais Contribuições

• Novo Paradigma de Alinhamento: Apresenta uma reformulação teórica do RLHF como um problema de inferência variacional sobre medidas positivas, unificando a teoria do RL com a simplicidade do SFT.
• Eliminação de Clipping e Pesos Negativos: Ao garantir pesos positivos através da exponenciação da recompensa, o método evita as armadilhas de otimização de métodos como DPO e PPO.
• Eficiência Computacional: O método é puramente offline (não requer amostragem online durante o treinamento) e elimina a necessidade de um modelo de crítico, reduzindo drasticamente o custo computacional e a complexidade de implementação.
• Estabilidade Superior: A abordagem oferece uma paisagem de otimização mais estável, permitindo treinamento mais longo e robusto sem colapso.

4. Resultados Experimentais

Os autores avaliaram o VAR em diversos benchmarks (HHA, MMLU, GSM8K, HumanEval, BBH) e modelos (Llama 3, Qwen 2.5).

• Desempenho de Alinhamento (HHA):
  • O VAR superou consistentemente o DPO em métricas de "Helpfulness" (Utilidade) e "Harmlessness" (Segurança).
  • Em média, houve um aumento de 7,16% em relação ao DPO nas métricas de utilidade.
  • O método funcionou bem tanto iniciando a partir de modelos base (pré-treinados) quanto a partir de modelos já ajustados via SFT.
• Comparação com Métodos Online (GRPO):
  • O VAR foi comparável ou superior ao GRPO (um método online simplificado) em termos de qualidade de resposta.
  • Velocidade: O VAR foi mais de 5 vezes mais rápido por época que o GRPO, demonstrando uma eficiência computacional superior.
• Benchmarks Generativos:
  • Em tarefas de raciocínio e conhecimento (MMLU, GSM8K, etc.), o VAR superou o DPO e o A-LoL, mantendo a capacidade de generalização do modelo.
• Estabilidade de Treinamento:
  • Gráficos de recompensa de validação mostraram que o VAR tem uma convergência mais suave e estável, enquanto o DPO tende a apresentar volatilidade e declínio precoce.
• Análise de Comprimento:
  • Diferente do DPO, que tende a gerar sequências excessivamente longas para maximizar recompensas, o VAR mantém comprimentos de saída mais naturais e próximos ao modelo de referência, sem sacrificar a qualidade.

5. Significado e Impacto

O trabalho VAR representa um avanço significativo na área de alinhamento de LLMs ao demonstrar que é possível alcançar (e até superar) o desempenho de métodos complexos de RLHF (como PPO) utilizando uma abordagem simples baseada em SFT ponderado.

• Ponte entre Eficiência e Desempenho: Resolve o dilema de que métodos rápidos são menos eficazes e métodos eficazes são lentos. O VAR oferece o melhor dos dois mundos.
• Acessibilidade: Ao remover a necessidade de infraestrutura complexa (múltiplos GPUs para modelos críticos, amostragem online), torna o alinhamento de alta qualidade acessível para laboratórios com recursos limitados.
• Estabilidade Teórica: A fundamentação em inferência variacional sobre medidas positivas oferece uma garantia teórica de estabilidade que falta em abordagens heurísticas atuais.

Em resumo, o VAR propõe que o "Santo Graal" do alinhamento pode ser alcançado não através de algoritmos de RL mais complexos, mas através de uma formulação matemática mais correta e eficiente do problema de otimização de políticas.