GenARM: Generación Guiada por Recompensa con Modelo de Recompensa Autoregresivo para Alineación en Tiempo de Prueba

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

Resumen

Los Modelos de Lenguaje de Gran Tamaño (LLMs) exhiben capacidades impresionantes pero requieren una alineación cuidadosa con las preferencias humanas. Los métodos tradicionales de entrenamiento ajustan los LLMs utilizando conjuntos de datos de preferencias humanas, pero incurren en costos significativos de entrenamiento y requieren entrenamientos repetidos para manejar diversas preferencias de usuario. Los métodos de alineación en tiempo de prueba abordan esto utilizando modelos de recompensa (RMs) para guiar los LLMs congelados sin necesidad de volver a entrenarlos. Sin embargo, los enfoques existentes en tiempo de prueba se basan en RMs a nivel de trayectoria, diseñados para evaluar respuestas completas, lo que los hace inadecuados para la generación de texto autoregresivo que requiere calcular recompensas de siguiente token a partir de respuestas parciales. Para abordar esto, presentamos GenARM, un enfoque de alineación en tiempo de prueba que aprovecha el Modelo de Recompensa Autoregresivo, una parametrización de recompensa novedosa diseñada para predecir recompensas de siguiente token de manera eficiente y efectiva para la generación autoregresiva. Teóricamente, demostramos que esta parametrización puede guiar de manera demostrable a los LLMs congelados hacia cualquier distribución alcanzable por RMs tradicionales dentro del marco de aprendizaje por refuerzo regularizado por KL. Los resultados experimentales muestran que GenARM supera significativamente a los baselines de alineación en tiempo de prueba anteriores y se equipara al rendimiento de los métodos en tiempo de entrenamiento. Además, GenARM permite una guía eficiente de débil a fuerte, alineando LLMs más grandes con RMs más pequeños sin los altos costos de entrenar modelos más grandes. Además, GenARM admite la alineación multiobjetivo, permitiendo compensaciones en tiempo real entre dimensiones de preferencia y atendiendo a diversas preferencias de usuario sin necesidad de volver a entrenar.

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.