Aprendendo Formas Explicáveis de Recompensas Densas via Otimização Bayesiana

Learning Explainable Dense Reward Shapes via Bayesian Optimization

April 22, 2025

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

cs.AI

Resumo

Os pipelines atuais de aprendizado por reforço com feedback humano (RLHF) para alinhamento de modelos de linguagem de grande escala (LLM) normalmente atribuem recompensas escalares a sequências, utilizando o token final como um indicador substituto para a qualidade de toda a sequência. No entanto, isso resulta em feedback esparso e atribuição de crédito em nível de token subótima. Neste trabalho, enquadramos a modelagem de recompensas como um problema de otimização focado na atribuição de crédito em nível de token. Propomos uma função de modelagem de recompensas que aproveita métodos de explicabilidade, como SHAP e LIME, para estimar recompensas por token a partir do modelo de recompensa. Para aprender os parâmetros dessa função de modelagem, empregamos um framework de otimização bilevel que integra Otimização Bayesiana e treinamento de políticas para lidar com o ruído das estimativas de recompensa por token. Nossos experimentos mostram que alcançar um melhor equilíbrio na atribuição de recompensas em nível de token leva a melhorias de desempenho em relação às baselines em tarefas subsequentes e encontra uma política ótima mais rapidamente durante o treinamento. Além disso, mostramos teoricamente que métodos de explicabilidade que são funções de atribuição aditivas de características mantêm a política ótima como a recompensa original.

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.