Обучение интерпретируемых плотных форм вознаграждений с использованием байесовской оптимизации

Learning Explainable Dense Reward Shapes via Bayesian Optimization

April 22, 2025

Авторы: Ryan Koo, Ian Yang, Vipul Raheja, Mingyi Hong, Kwang-Sung Jun, Dongyeop Kang

cs.AI

Аннотация

Современные подходы к обучению с подкреплением на основе человеческой обратной связи (RLHF) для согласования больших языковых моделей (LLM) обычно присваивают скалярные награды последовательностям, используя финальный токен как суррогатный индикатор качества всей последовательности. Однако это приводит к разреженной обратной связи и субоптимальному распределению наград на уровне токенов. В данной работе мы формулируем формирование наград как задачу оптимизации, сосредоточенную на распределении наград на уровне токенов. Мы предлагаем функцию формирования наград, которая использует методы объяснимости, такие как SHAP и LIME, для оценки наград на уровне токенов на основе модели наград. Для обучения параметров этой функции мы применяем двухуровневую оптимизационную структуру, которая интегрирует байесовскую оптимизацию и обучение политики для обработки шума в оценках наград на уровне токенов. Наши эксперименты показывают, что достижение более сбалансированного распределения наград на уровне токенов приводит к улучшению производительности по сравнению с базовыми подходами на последующих задачах и ускоряет нахождение оптимальной политики в процессе обучения. Кроме того, мы теоретически показываем, что методы объяснимости, являющиеся аддитивными функциями атрибуции признаков, сохраняют оптимальную политику, соответствующую исходной награде.

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.