Lernen erklärbarer dichter Belohnungsformen durch Bayessche Optimierung

Learning Explainable Dense Reward Shapes via Bayesian Optimization

April 22, 2025

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

cs.AI

Zusammenfassung

Aktuelle Pipelines für Reinforcement Learning aus menschlichem Feedback (RLHF) zur Ausrichtung großer Sprachmodelle (LLM) weisen in der Regel skalare Belohnungen Sequenzen zu, wobei das letzte Token als Ersatzindikator für die Qualität der gesamten Sequenz verwendet wird. Dies führt jedoch zu spärlichem Feedback und suboptimaler Token-spezifischer Belohnungszuweisung. In dieser Arbeit formulieren wir die Belohnungsgestaltung als ein Optimierungsproblem, das sich auf die Token-spezifische Belohnungszuweisung konzentriert. Wir schlagen eine Belohnungsgestaltungsfunktion vor, die Erklärbarkeitsmethoden wie SHAP und LIME nutzt, um Token-spezifische Belohnungen aus dem Belohnungsmodell abzuschätzen. Um die Parameter dieser Gestaltungsfunktion zu lernen, verwenden wir ein bi-level-Optimierungsframework, das Bayesian Optimization und Policy-Training integriert, um das Rauschen in den Token-Belohnungsschätzungen zu handhaben. Unsere Experimente zeigen, dass eine bessere Balance bei der Token-spezifischen Belohnungszuweisung zu Leistungsverbesserungen gegenüber den Baselines bei nachgelagerten Aufgaben führt und während des Trainings eine optimale Policy schneller findet. Darüber hinaus zeigen wir theoretisch, dass Erklärbarkeitsmethoden, die als feature-additive Attributionsfunktionen fungieren, die optimale Policy der ursprünglichen Belohnung beibehalten.

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.