Resumo Técnico: Aprendizado por Reforço com Feedback Humano Distribucionalmente Robusto

Declaração do Problema

O Aprendizado por Reforço com Feedback Humano (RLHF) tornou-se uma técnica padrão para alinhar Grandes Modelos de Linguagem (LLMs) com as intenções humanas. No entanto, os pipelines de RLHF existentes sofrem de uma falta de robustez. Quando a distribuição de prompts encontrados durante a inferência (tarefas downstream) difere significamente da distribuição do conjunto de dados de preferência usado durante o ajuste fino, o desempenho do modelo frequentemente se deteriora. Este problema é exacerbado por fenômenos como a sobre-otimização, onde as políticas melhoram recompensas substitutas (proxy rewards) em detrimento de funções de recompensa reais, e mudanças de distribuição causadas por ambientes em mudança ou deriva nas preferências humanas. Abordagens anteriores de robustez focaram em modelos de corrupção adversariais específicos, mudanças no modelo de preferência ou métodos heurísticos como a fusão de modelos, que frequentemente falham sob mudanças distributivas leves.

Metodologia

Os autores propõem um framework de Otimização Distribucionalmente Robusta (DRO) para RLHF. O objetivo central é garantir que os modelos ajustados (tanto modelos de recompensa quanto políticas) mantenam o desempenho mesmo quando a distribuição de prompts muda dentro de uma distância limitada da distribuição de treinamento. Os autores formulam isso introduzindo um problema de otimização minimax onde o aprendiz minimiza a perda contra a pior distribuição dentro de uma bola de $\phi$-divergência (especificamente a distância de Variação Total) ao redor da distribuição empírica de treinamento.

O artigo desenvolve variantes robustas para dois pipelines primários de RLHF:

1. Estimativa de Recompensa Robusta:

   • Formulação: O modelo de recompensa $r$ é treinado para minimizar a pior expectativa de perda sobre distribuições $D$ tal que $d_{TV}(D, D_{src}) \leq \rho$.
   • Algoritmo: Um algoritmo de gradiente descendente por minibatch (Algoritmo 1) é proposto. Em cada iteração, o algoritmo amostra um minibatch e computa pesos ótimos $q^*_t$ que correspondem à mudança de distribuição de pior caso dentro da distância $\rho$. O gradiente é então reponderado por esses pesos ótimos antes de atualizar os parâmetros da recompensa.
   • Convergência: Sob a suposição de uma classe de função de recompensa linear, os autores provam que o algoritmo converge para uma solução $\epsilon$-ótima com $T = O(1/\epsilon^2)$ iterações e $n = \tilde{O}(1/\epsilon^2)$ amostras por iteração.

2. Otimização de Política Robusta:

   • Formulação: A política $\pi$ é otimizada para maximizar a função de valor regularizada por KL no pior caso.
   • Algoritmo: Os autores propõem um método de Gradiente de Política Natural (NPG) robusto (Algoritmo 2). Semelhante à estimativa de recompensa, ele computa pesos de pior caso para o minibatch. Crucialmente, ele atualiza a política usando uma matriz de informação de Fisher ponderada derivada desses pesos de pior caso, em vez da matriz de Fisher empírica padrão.
   • Convergência: Assumindo uma classe de política log-linear e condições de concentrabilidade padrão, o método alcança uma taxa de convergência linear ($T = O(\log(1/\epsilon))$), mas requer uma complexidade de amostragem mais alta de $n = \tilde{O}(1/\epsilon^4)$ por iteração para garantir a convergência da última iteração.

3. Otimização Direta de Preferência (DPO) Robusta:

   • Formulação: Uma versão distribucionalmente robusta de DPO é formulada, otimizando diretamente a política sem um modelo de recompensa explícito.
   • Algoritmo: O Algoritmo 3 aplica gradiente estocástico por minibatch com gradientes reponderados com base na mudança de distribuição de pior caso.
   • Convergência: A análise mostra uma complexidade de iteração de $T = O(1/\epsilon^2)$ e complexidade de amostragem de $n = \tilde{O}(1/\epsilon^2)$, similar à fase de estimativa de recompensa robusta.

Principais Contribuições

• Framework Teórico: O artigo estabelece as primeiras formulações distribucionalmente robustas para ambas as fases de estimativa de recompensa e otimização de política de RLHF, bem como para DPO.
• Algoritmos: Introduz algoritmos práticos baseados em minibatch para Estimativa de Recompensa Robusta, NPG Robusto e DPO Robusto. Estes algoritmos são projetados para serem minimamente invasivos aos pipelines existentes, exigindo apenas o cálculo de fatores de reponderação (resolvidos via otimização convexa) e reponderação de gradiente.
• Garantias de Convergência: Os autores fornecem provas teóricas para a convergência de seus algoritmos. Eles estabelecem limites para a complexidade de amostragem e de iteração para modelos de recompensa lineares e políticas log-lineares, demonstrando que a robustez pode ser alcançada com garantias prováveis.
• Validação Empírica: Os métodos são avaliados em tarefas fora da distribuição (OOD) usando o dataset Unified-Feedback para treinamento e RewardBench, HHH-Alignment e MT-Bench para avaliação.

Resultados Experimentais

Os autores avaliaram seus métodos utilizando os modelos base Gemma-2B-it e Mistral-7B-Instruct-v0.2.

• Modelos de Recompensa: No benchmark RewardBench, os modelos de recompensa distribucionalmente robustos (DR RM) mostraram melhoria na precisão média em comparação com os modelos Bradley-Terry (BTL) padrão, particularmente quando o parâmetro de robustez $\rho$ foi definido como 0.1 ou 0.2. Notavelmente, o desempenho no subconjunto de Raciocínio (Reasoning) melhorou significativamente (ex: +23.4% em um subconjunto com $\rho=0.1$). No entanto, o desempenho degradou quando $\rho$ foi definido muito alto (ex: $\rho \geq 0.6$), indicando um compromisso entre robustez e o ajuste à distribuição de treinamento.
• Otimização de Política (PPO): Ao usar os modelos de recompensa robustos para treinar políticas via PPO, as variantes robustas (DR PPO) geralmente igualaram ou excederam o desempenho do PPO padrão tanto em conjuntos de dados dentro da distribuição (Unified-Feedback) quanto fora da distribuição (HHH-Alignment, MT-Bench). Uma melhoria marcante foi observada em tarefas de raciocínio, com um aumento de +10.0% na precisão para $\rho=0.1$.
• Otimização Direta de Preferência (DPO): O DPO robusto (DR DPO) consistentemente melhorou o desempenho em conjuntos de dados OOD em comparação com o DPO padrão. Para o HHH-Alignment, o melhor desempenho foi alcançado com $\rho=0.8$, e para o MT-Bench com $\rho=0.5$. O DR DPO geralmente superou o DR PPO em tarefas OOD.

Significância e Alegações

O artigo afirma que a otimização distribucionalmente robusta fornece uma abordagem fundamentada para melhorar a generalização de LLMs para tarefas não vistas e mudanças de distribuição. Os autores enfatizam que seus métodos exigem alterações mínimas nos pipelines eficientes de RLHF existentes (por exemplo, aqueles que usam bibliotecas Hugging Face), enquanto oferecem garantias teóricas de convergência.

Um achado fundamental destacado pelos autores é a melhoria inesperada nas capacidades de raciocínio em conjuntos de dados OOD ao aplicar robustez de distribuição. Os autores observam que não é imediatamente claro por que o DRO deveria afetar positivamente o raciocínio, sugerindo que este fenômeno merece investigação adicional sobre as conexões entre raciocínio e robustez distribucional.

O trabalho posiciona-se como um passo em direção ao alinhamento real de LLMs com as intenções humanas em ambientes do mundo real dinâmicos, onde as distribuições de dados não são estáticas, indo além das limitações do RLHF padrão, que frequentemente falha em generalizar além dos exemplos de preferência específicos vistos durante o treinamento.