RRM: Treinamento do Modelo de Recompensa Robusto Mitiga a Manipulação de Recompensas

RRM: Robust Reward Model Training Mitigates Reward Hacking

September 20, 2024

Autores: Tianqi Liu, Wei Xiong, Jie Ren, Lichang Chen, Junru Wu, Rishabh Joshi, Yang Gao, Jiaming Shen, Zhen Qin, Tianhe Yu, Daniel Sohn, Anastasiia Makarova, Jeremiah Liu, Yuan Liu, Bilal Piot, Abe Ittycheriah, Aviral Kumar, Mohammad Saleh

cs.AI

Resumo

Os modelos de recompensa (RMs) desempenham um papel fundamental na alinhamento de grandes modelos de linguagem (LLMs) com as preferências humanas. No entanto, o treinamento tradicional de RM, que se baseia em pares de respostas vinculadas a prompts específicos, enfrenta dificuldades para separar as preferências impulsionadas por prompts de artefatos independentes de prompts, como comprimento e formato de resposta. Neste trabalho, revelamos uma limitação fundamental dos métodos de treinamento de RM atuais, nos quais os RMs falham em distinguir efetivamente entre sinais contextuais e artefatos irrelevantes ao determinar preferências. Para lidar com isso, introduzimos um framework causal que aprende preferências independentemente desses artefatos e propomos uma nova técnica de aumento de dados projetada para eliminá-los. Experimentos extensos mostram que nossa abordagem filtra com sucesso artefatos indesejáveis, resultando em um modelo de recompensa mais robusto (RRM). Nosso RRM melhora o desempenho de um modelo de recompensa em pares treinado no Gemma-2-9b-it, no RewardBench, aumentando a precisão de 80,61% para 84,15%. Além disso, treinamos duas políticas DPO usando tanto o RM quanto o RRM, demonstrando que o RRM melhora significativamente as políticas alinhadas com DPO, aumentando os escores do MT-Bench de 7,27 para 8,31 e as taxas de vitória controladas por comprimento no AlpacaEval-2 de 33,46% para 52,49%.

English

Reward models (RMs) play a pivotal role in aligning large language models (LLMs) with human preferences. However, traditional RM training, which relies on response pairs tied to specific prompts, struggles to disentangle prompt-driven preferences from prompt-independent artifacts, such as response length and format. In this work, we expose a fundamental limitation of current RM training methods, where RMs fail to effectively distinguish between contextual signals and irrelevant artifacts when determining preferences. To address this, we introduce a causal framework that learns preferences independent of these artifacts and propose a novel data augmentation technique designed to eliminate them. Extensive experiments show that our approach successfully filters out undesirable artifacts, yielding a more robust reward model (RRM). Our RRM improves the performance of a pairwise reward model trained on Gemma-2-9b-it, on RewardBench, increasing accuracy from 80.61% to 84.15%. Additionally, we train two DPO policies using both the RM and RRM, demonstrating that the RRM significantly enhances DPO-aligned policies, improving MT-Bench scores from 7.27 to 8.31 and length-controlled win-rates in AlpacaEval-2 from 33.46% to 52.49%.