Apprendimento di Forme Dense di Ricompensa Spiegabile tramite Ottimizzazione Bayesiana

Learning Explainable Dense Reward Shapes via Bayesian Optimization

April 22, 2025

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

cs.AI

Abstract

Le attuali pipeline di apprendimento per rinforzo da feedback umano (RLHF) per l'allineamento di modelli linguistici di grandi dimensioni (LLM) assegnano tipicamente ricompense scalari alle sequenze, utilizzando il token finale come indicatore surrogato della qualità dell'intera sequenza. Tuttavia, ciò porta a feedback sparsi e a un'assegnazione di credito a livello di token subottimale. In questo lavoro, inquadriamo la modellazione delle ricompense come un problema di ottimizzazione focalizzato sull'assegnazione di credito a livello di token. Proponiamo una funzione di modellazione delle ricompense che sfrutta metodi di spiegabilità come SHAP e LIME per stimare le ricompense per token dal modello di ricompensa. Per apprendere i parametri di questa funzione di modellazione, utilizziamo un framework di ottimizzazione bilevel che integra l'ottimizzazione bayesiana e l'addestramento della policy per gestire il rumore derivante dalle stime delle ricompense a livello di token. I nostri esperimenti dimostrano che il raggiungimento di un migliore bilanciamento nell'attribuzione delle ricompense a livello di token porta a miglioramenti delle prestazioni rispetto ai benchmark sui task downstream e trova una policy ottimale più rapidamente durante l'addestramento. Inoltre, dimostriamo teoricamente che i metodi di spiegabilità che sono funzioni di attribuzione additive rispetto alle feature mantengono la policy ottimale come la ricompensa 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.