GenARM: Geração Guiada por Recompensa com Modelo de Recompensa Autoregressivo para Alinhamento em Tempo de Teste

GenARM: Reward Guided Generation with Autoregressive Reward Model for Test-time Alignment

October 10, 2024

Autores: Yuancheng Xu, Udari Madhushani Sehwag, Alec Koppel, Sicheng Zhu, Bang An, Furong Huang, Sumitra Ganesh

cs.AI

Resumo

Os Modelos de Linguagem de Grande Escala (LLMs) exibem capacidades impressionantes, mas requerem alinhamento cuidadoso com as preferências humanas. Métodos tradicionais de treinamento ajustam os LLMs usando conjuntos de dados de preferências humanas, porém incorrem em custos significativos de treinamento e exigem treinamento repetido para lidar com diversas preferências de usuários. Métodos de alinhamento no momento de teste abordam isso usando modelos de recompensa (RMs) para orientar os LLMs congelados sem retrabalho. No entanto, abordagens existentes no momento de teste dependem de RMs em nível de trajetória, que são projetados para avaliar respostas completas, tornando-os inadequados para geração de texto autoregressiva que requer o cálculo de recompensas do próximo token a partir de respostas parciais. Para lidar com isso, apresentamos o GenARM, uma abordagem de alinhamento no momento de teste que aproveita o Modelo de Recompensa Autoregressivo - uma parametrização de recompensa inovadora projetada para prever recompensas do próximo token para geração autoregressiva eficiente e eficaz. Teoricamente, demonstramos que essa parametrização pode guiar de forma comprovada os LLMs congelados em direção a qualquer distribuição alcançável por RMs tradicionais dentro do framework de aprendizado por reforço regularizado por KL. Resultados experimentais mostram que o GenARM supera significativamente baselines anteriores de alinhamento no momento de teste e iguala o desempenho de métodos no momento de treinamento. Além disso, o GenARM permite um direcionamento eficiente de fraco para forte, alinhando LLMs maiores com RMs menores sem os altos custos de treinar modelos maiores. Além disso, o GenARM suporta alinhamento multiobjetivo, permitindo compensações em tempo real entre dimensões de preferência e atendendo a diversas preferências de usuários sem retrabalho.

English

Large Language Models (LLMs) exhibit impressive capabilities but require careful alignment with human preferences. Traditional training-time methods finetune LLMs using human preference datasets but incur significant training costs and require repeated training to handle diverse user preferences. Test-time alignment methods address this by using reward models (RMs) to guide frozen LLMs without retraining. However, existing test-time approaches rely on trajectory-level RMs which are designed to evaluate complete responses, making them unsuitable for autoregressive text generation that requires computing next-token rewards from partial responses. To address this, we introduce GenARM, a test-time alignment approach that leverages the Autoregressive Reward Model--a novel reward parametrization designed to predict next-token rewards for efficient and effective autoregressive generation. Theoretically, we demonstrate that this parametrization can provably guide frozen LLMs toward any distribution achievable by traditional RMs within the KL-regularized reinforcement learning framework. Experimental results show that GenARM significantly outperforms prior test-time alignment baselines and matches the performance of training-time methods. Additionally, GenARM enables efficient weak-to-strong guidance, aligning larger LLMs with smaller RMs without the high costs of training larger models. Furthermore, GenARM supports multi-objective alignment, allowing real-time trade-offs between preference dimensions and catering to diverse user preferences without retraining.