Leerbare verklaarbare dichte beloningsvormen via Bayesiaanse optimalisatie

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

Samenvatting

Huidige pijplijnen voor reinforcement learning met menselijke feedback (RLHF) voor de afstemming van grote taalmodelen (LLM) kennen doorgaans scalaire beloningen toe aan sequenties, waarbij het laatste token wordt gebruikt als een surrogaatindicator voor de kwaliteit van de hele sequentie. Dit leidt echter tot schaarse feedback en suboptimale toewijzing van credits op tokenniveau. In dit werk formuleren we beloningsvorming als een optimalisatieprobleem gericht op de toewijzing van credits op tokenniveau. We stellen een beloningsvormingsfunctie voor die gebruikmaakt van verklaarbaarheidsmethoden zoals SHAP en LIME om per-token beloningen te schatten vanuit het beloningsmodel. Om de parameters van deze vormingsfunctie te leren, gebruiken we een bi-level optimalisatieraamwerk dat Bayesiaanse optimalisatie en beleidstraining integreert om ruis van de tokenbeloningsschattingen te hanteren. Onze experimenten tonen aan dat het bereiken van een betere balans in de toewijzing van beloningen op tokenniveau leidt tot prestatieverbeteringen ten opzichte van basislijnen bij downstream taken en een optimaal beleid sneller vindt tijdens de training. Bovendien tonen we theoretisch aan dat verklaarbaarheidsmethoden die functies zijn voor additieve attributie van kenmerken, het optimale beleid behouden als de oorspronkelijke beloning.

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.