Apprentissage de formes de récompenses denses explicables par optimisation bayésienne

Learning Explainable Dense Reward Shapes via Bayesian Optimization

April 22, 2025

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

cs.AI

Résumé

Les pipelines actuels d'apprentissage par renforcement à partir de retours humains (RLHF) pour l'alignement des grands modèles de langage (LLM) attribuent généralement des récompenses scalaires à des séquences, en utilisant le dernier jeton comme indicateur substitut de la qualité de l'ensemble de la séquence. Cependant, cela conduit à des retours épars et à une attribution de crédit au niveau des jetons sous-optimale. Dans ce travail, nous formulons le façonnage des récompenses comme un problème d'optimisation centré sur l'attribution de crédit au niveau des jetons. Nous proposons une fonction de façonnage des récompenses exploitant des méthodes d'explicabilité telles que SHAP et LIME pour estimer les récompenses par jeton à partir du modèle de récompense. Pour apprendre les paramètres de cette fonction de façonnage, nous utilisons un cadre d'optimisation bi-niveaux qui intègre l'optimisation bayésienne et l'entraînement des politiques pour gérer le bruit des estimations de récompense par jeton. Nos expériences montrent qu'un meilleur équilibre dans l'attribution des récompenses au niveau des jetons améliore les performances par rapport aux bases de référence sur les tâches en aval et permet de trouver une politique optimale plus rapidement pendant l'entraînement. De plus, nous démontrons théoriquement que les méthodes d'explicabilité qui sont des fonctions d'attribution additives par caractéristique maintiennent la politique optimale comme la récompense originale.

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.