Aprendizaje de Formas Explicables de Recompensa Densa mediante Optimización Bayesiana

Learning Explainable Dense Reward Shapes via Bayesian Optimization

April 22, 2025

Autores: Ryan Koo, Ian Yang, Vipul Raheja, Mingyi Hong, Kwang-Sung Jun, Dongyeop Kang

cs.AI

Resumen

Los flujos de trabajo actuales de aprendizaje por refuerzo a partir de retroalimentación humana (RLHF) para la alineación de modelos de lenguaje de gran escala (LLM) suelen asignar recompensas escalares a secuencias, utilizando el token final como un indicador sustituto de la calidad de toda la secuencia. Sin embargo, esto conduce a una retroalimentación dispersa y a una asignación de crédito a nivel de token subóptima. En este trabajo, enmarcamos la configuración de recompensas como un problema de optimización centrado en la asignación de crédito a nivel de token. Proponemos una función de configuración de recompensas que aprovecha métodos de explicabilidad como SHAP y LIME para estimar recompensas por token a partir del modelo de recompensas. Para aprender los parámetros de esta función de configuración, empleamos un marco de optimización bi-nivel que integra Optimización Bayesiana y entrenamiento de políticas para manejar el ruido de las estimaciones de recompensas por token. Nuestros experimentos muestran que lograr un mejor equilibrio en la atribución de recompensas a nivel de token conduce a mejoras en el rendimiento respecto a las líneas base en tareas posteriores y encuentra una política óptima más rápido durante el entrenamiento. Además, demostramos teóricamente que los métodos de explicabilidad que son funciones de atribución aditiva de características mantienen la política óptima como la recompensa original.

English

Current reinforcement learning from human feedback (RLHF) pipelines for large language model (LLM) alignment typically assign scalar rewards to sequences, using the final token as a surrogate indicator for the quality of the entire sequence. However, this leads to sparse feedback and suboptimal token-level credit assignment. In this work, we frame reward shaping as an optimization problem focused on token-level credit assignment. We propose a reward-shaping function leveraging explainability methods such as SHAP and LIME to estimate per-token rewards from the reward model. To learn parameters of this shaping function, we employ a bilevel optimization framework that integrates Bayesian Optimization and policy training to handle noise from the token reward estimates. Our experiments show that achieving a better balance of token-level reward attribution leads to performance improvements over baselines on downstream tasks and finds an optimal policy faster during training. Furthermore, we show theoretically that explainability methods that are feature additive attribution functions maintain the optimal policy as the original reward.