Segmentación de texto y aprendizaje de sus recompensas para mejorar RLHF en modelos de lenguaje.

Segmenting Text and Learning Their Rewards for Improved RLHF in Language Model

January 6, 2025

Autores: Yueqin Yin, Shentao Yang, Yujia Xie, Ziyi Yang, Yuting Sun, Hany Awadalla, Weizhu Chen, Mingyuan Zhou

cs.AI

Resumen

El aprendizaje por refuerzo a partir de la retroalimentación humana (RLHF, por sus siglas en inglés) ha sido ampliamente adoptado para alinear los modelos de lenguaje (LMs) con las preferencias humanas. Los trabajos previos de RLHF suelen adoptar una formulación de bandit, que, aunque intuitiva, ignora la naturaleza secuencial de la generación de LM y puede verse afectada por el problema de la recompensa escasa. Mientras que trabajos recientes proponen RLHF a nivel de token denso, tratando cada token como una acción puede ser demasiado sutil para la asignación adecuada de recompensas. En este documento, buscamos combinar lo mejor de ambos en el entrenamiento y la utilización de un modelo de recompensa a nivel de segmento, que asigna una recompensa a cada segmento de texto semánticamente completo que abarca una secuencia corta de tokens. Para el aprendizaje de recompensas, nuestro método permite la segmentación dinámica de texto y es compatible con conjuntos de datos estándar de preferencias de secuencia. Para el entrenamiento efectivo de LM basado en RL con recompensa de segmento, generalizamos los normalizadores de recompensa de bandit escalares clásicos en funciones normalizadoras con conciencia de ubicación e interpolamos la recompensa de segmento para una mayor densificación. Con estos diseños, nuestro método tiene un rendimiento competitivo en tres benchmarks populares de RLHF para políticas de LM: AlpacaEval 2.0, Arena-Hard y MT-Bench. Se realizaron estudios de ablación para demostrar aún más nuestro método.

English

Reinforcement learning from human feedback (RLHF) has been widely adopted to align language models (LMs) with human preference. Prior RLHF works typically take a bandit formulation, which, though intuitive, ignores the sequential nature of LM generation and can suffer from the sparse reward issue. While recent works propose dense token-level RLHF, treating each token as an action may be oversubtle to proper reward assignment. In this paper, we seek to get the best of both by training and utilizing a segment-level reward model, which assigns a reward to each semantically complete text segment that spans over a short sequence of tokens. For reward learning, our method allows dynamic text segmentation and compatibility with standard sequence-preference datasets. For effective RL-based LM training against segment reward, we generalize the classical scalar bandit reward normalizers into location-aware normalizer functions and interpolate the segment reward for further densification. With these designs, our method performs competitively on three popular RLHF benchmarks for LM policy: AlpacaEval 2.0, Arena-Hard, and MT-Bench. Ablation studies are conducted to further demonstrate our method.