Provably avoiding over-optimization in Direct Preference Optimization without knowing the data distribution

Declaração do Problema

A Otimização de Preferência Direta (DPO) tornou-se um padrão para o pós-treinamento de modelos de linguagem de grande escala (LLMs) devido à sua simplicidade e capacidade de contornar a modelagem explícita de recompensas. No entanto, sabe-se que a DPO sofre de superotimização: à medida que o treinamento avança, o modelo pode maximizar a perda de treinamento enquanto degrada a qualidade real das respostas geradas em relação ao modelo inicial. Esse problema é exacerbado pelo fato de que a DPO padrão e o Aprendizado por Reforço a partir de Feedback Humano (RLHF) tratam a recompensa estimada como uma verdade absoluta perfeita, ignorando a incerteza.

As garantias teóricas para mitigar a superotimização geralmente dependem do pessimismo, assegurando que a política aprendida seja robusta a erros de estimação. Variantes pessimistas existentes (por exemplo, DPO regularizada por $\chi^2$, DPO+SFT) alcançam limites teóricos fortes que dependem apenas de um coeficiente de concentrabilidade de política única ($C^\star$), o que implica robustez mesmo se o conjunto de dados não cobrir todo o espaço de políticas. No entanto, essas garantias dependem criticamente do conhecimento da distribuição geradora de dados ($\pi_{data}$) ou do acesso a uma política de referência bem coberta por $\pi_{data}$. Na prática, essa suposição é frequentemente violada, como ao ajustar modelos menores em dados de preferência gerados por sistemas proprietários maiores (por exemplo, GPT-4), onde $\pi_{data}$ é desconhecido e inacessível.

A questão central de pesquisa abordada é: Podemos projetar um algoritmo semelhante à DPO que seja comprovadamente robusto à superotimização sem conhecimento da distribuição geradora de dados $\pi_{data}$?

Metodologia: PEPO

Os autores introduzem o PEPO (Otimização de Preferência baseada em Ensemble Pessimista), um algoritmo de etapa única que alcança o pessimismo por meio de um ensemble de políticas treinadas em subconjuntos de dados disjuntos, sem exigir conhecimento de $\pi_{data}$.

1. Treinamento de Ensemble:
O conjunto de dados de preferência $D$ é partitionado em $L$ subconjuntos disjuntos. Para cada subconjunto $D_\ell$, uma política leve $\tilde{\pi}_\ell$ é treinada usando um objetivo pessimista de DPO. Esse objetivo modifica a perda padrão da DPO deslocando horizontalmente a função sigmoide por um parâmetro $\lambda_\ell$, efetivamente reduzindo a margem necessária para uma probabilidade de vitória. Esse deslocamento induz uma estimativa conservadora (pessimista) da lacuna de recompensa.

2. Agregação Pessimista:
A política de saída $\pi_{out}$ é definida como a solução de um problema de RLHF "virtual" que otimiza a recompensa do pior caso através do ensemble. A recompensa implícita para o ensemble é construída tomando o mínimo sobre os $L$ modelos de suas lacunas de recompensa estimadas, ajustadas para a distribuição de dados nos outros subconjuntos (para garantir independência).
$$r^-(x, a) = \beta \min_{\ell \in [L]} \left( \log \frac{\tilde{\pi}_\ell(a|x)}{\pi_{ref}(a|x)} - \zeta_\ell(x) \right)$$
onde $\zeta_\ell(x)$ accounts pela contribuição de dados fora do $\ell$-ésimo subconjunto.

3. Amostragem Eficiente:
A solução de forma fechada para $\pi_{out}$ envolve uma constante de normalização sobre todo o espaço de respostas, que é intratável. Para contornar isso, o PEPO emprega amostragem por rejeição (Algoritmo 2). Qualquer membro do ensemble $\tilde{\pi}_\ell$ serve como distribuição de proposta porque o numerador de $\pi_{out}$ é limitado por $\tilde{\pi}_\ell$. Para aplicações práticas em grande escala, os autores também propõem uma aproximação ao nível de token que aplica a agregação pessimista em cada etapa de decodificação, evitando a necessidade de amostragem por rejeição ao nível de trajetória e reduzindo significativamente a latência de geração.

Principais Contribuições

Contribuições Teóricas:

• Complexidade de Amostra sem $\pi_{data}$: No cenário tabular, o PEPO alcança garantias de complexidade de amostra que dependem apenas do coeficiente de concentrabilidade de política única ($C^\star$). Isso contrasta com a DPO padrão, que sofre de um coeficiente de concentrabilidade de todas as políticas ($C_{all}$), e com variantes pessimistas existentes de DPO que exigem conhecimento de $\pi_{data}$.
• Tamanho Ótimo do Ensemble: A análise caracteriza o tamanho ótimo do ensemble $L$ como o menor valor suficiente para garantir o pessimismo com alta probabilidade. Limites teóricos mostram que $L$ escala logaritmicamente com os tamanhos dos espaços de estado e ação.
• Robustez: Os limites valem sem assumir que o conjunto de dados cobre toda a classe de políticas, abordando a lacuna de "$\pi_{data}$ desconhecido" em trabalhos teóricos anteriores.

Contribuições Práticas:

• Validação Empírica: O PEPO é avaliado em modelos com 7B, 8B e 34B parâmetros (Zephyr-7B, Llama-3.1-8B, Mistral-7B, Yi-34B) usando o conjunto de dados UltraFeedback.
• Mitigação da Superotimização: Os experimentos demonstram que, enquanto a DPO padrão exibe uma queda acentuada nas taxas de vitória após ganhos iniciais (superotimização), o desempenho do PEPO se estabiliza e melhora.
• Superioridade sobre Baselines: O PEPO supera consistentemente o $\chi^2$PO e o DPO+SFT em configurações onde $\pi_{data} \neq \pi_{ref}$ (por exemplo, UltraFeedback), onde as garantias teóricas dessas baselines são inválidas.
• Eficiência: A implementação ao nível de token alcança um compromisso favorável entre taxa de vitória e tempo de geração, oferecendo desempenho quase ótimo com latência significativamente menor do que a amostragem por rejeição completa.

Resultados

• Superotimização: A Figura 1 mostra que a taxa de vitória da DPO contra o modelo inicial atinge um pico e depois declina, enquanto o PEPO (com $L$ apropriado) mostra melhoria sustentada. A parada antecipada baseada na perda de validação mostrou-se pouco confiável para selecionar o melhor checkpoint de DPO.
• Comparação: Na Figura 3, o PEPO alcança taxas de vitória mais altas no AlpacaEval2 em comparação com $\chi^2$PO e SFT+DPO em todos os modelos testados. A lacuna de desempenho estreita-se apenas quando $\pi_{data} = \pi_{ref}$, confirmando que a vantagem do PEPO decorre de sua capacidade de lidar com distribuições de dados desconhecidas.
• Ablação: Os experimentos confirmam que $L$ é um hiperparâmetro crítico; um valor muito pequeno falha em induzir pessimismo suficiente, enquanto um valor muito grande reduz os dados efetivos por modelo, degradando o desempenho.

Significado e Alegações

O artigo alega apresentar o primeiro algoritmo semelhante à DPO que evita comprovadamente a superotimização sem conhecimento da distribuição geradora de dados. Isso é significativo porque preenche a lacuna entre robustez teórica (concentrabilidade de política única) e restrições práticas ($\pi_{data}$ desconhecido), um cenário comum no alinhamento de LLMs do mundo real onde os dados são gerados por modelos proprietários ou diversos.

Os autores posicionam o PEPO como uma solução prática que mantém a simplicidade da DPO (treinamento de etapa única, sem modelo de recompensa explícito) enquanto incorpora a segurança teórica da regularização pessimista. Eles reconhecem que suas garantias teóricas são atualmente estabelecidas no cenário tabular, e estender esses resultados para classes gerais de aproximação de funções sem conhecer $\pi_{data}$ permanece um desafio aberto. Além disso, o método atual produz uma política não paramétrica (via amostragem), que é ideal para a etapa final de um pipeline de treinamento, mas não para checkpoints intermediários, sugerindo uma direção para trabalhos futuros sobre saídas paramétricas.