Paper deep dive

Learning Explainable Dense Reward Shapes via Bayesian Optimization

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

Year: 2025Venue: arXiv preprintArea: Alignment TrainingType: EmpiricalEmbeddings: 89

Models: LLaMA-3.2, OpenLLaMA

Intelligence

Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 95%

Last extracted: 3/12/2026, 6:04:23 PM

Summary

The paper introduces a bilevel optimization framework to improve reinforcement learning from human feedback (RLHF) by densifying sparse rewards. It utilizes explainability methods (SHAP, LIME) to estimate token-level contributions and employs Bayesian Optimization to learn optimal weights for these reward-shaping functions, leading to faster policy convergence and improved performance on downstream tasks.

Entities (5)

Bayesian Optimization · optimization-algorithm · 99%LIME · explainability-method · 99%SHAP · explainability-method · 99%PPO · reinforcement-learning-algorithm · 95%RLHF · methodology · 95%

Relation Signals (3)

SHAP → estimates → Token-level Rewards

confidence 95% · We propose a reward-shaping function leveraging explainability methods such as SHAP and LIME to estimate per-token rewards

Bayesian Optimization → optimizes → Reward Shaping Function

confidence 95% · we employ a bilevel optimization framework that integrates Bayesian Optimization and policy training to handle noise from the token reward estimates.

Reward Shaping → improves → RLHF

confidence 92% · achieving a better balance of token-level reward attribution leads to performance improvements over baselines on downstream tasks

Cypher Suggestions (0)

No Cypher suggestions yet.

Abstract

Abstract: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.

PDF

Open source PDF →Open local PDF →

Trouble viewing inline? Open PDF directly →

Full Text

88,888 characters extracted from source content.

Expand or collapse full text

Learning Explainable Dense Reward Shapes via Bayesian Optimization Ryan Koo1, Ian Yang2, Vipul Raheja3, Mingyi Hong1, Kwang-Sung Jun4, Dongyeop Kang1 University of Minnesota1, Georgia Institute of Technology2, Grammarly3, University of Arizona4 koo00017, mhong, dongyeop@umn.edu, iyang30@gatech.edu, vipul.raheja@grammarly.com, kjun@cs.arizona.edu Abstract 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. The code is publicly available.111https://github.com/minnesotanlp/explainable-dense-rewards 1 Introduction One of the fundamental challenges in reinforcement learning (RL) arises from the sparsity of rewards, where feedback signals are typically only provided at the end of a trajectory, with little to no evaluative information about the intermediate states. As a result, this limitation leads to gaps about the favorability of certain intermediate states, as the agent lacks the required granular feedback regarding which actions were beneficial to the outcome. A similar challenge has been observed in many recent applications of reinforcement learning from human feedback (RLHF), where sparse rewards are common (Zheng et al., 2023b; Chaudhari et al., 2024). Given the sequential nature of language models and the token-level value function typically optimized in RLHF (Zhong et al., 2025; Rafailov et al., 2024), it is crucial to determine the contribution of individual tokens to the overall reward. The challenge lies in the fact that the reward is often awarded at the end of sequence generation and represents the quality of the entire sequence as a scalar, which is known to be unstable (Razin et al., 2024; Engstrom et al., 2020) and encodes a low-bandwidth signal that does not help determine the relative quality of intermediate tokens. Hence, to address this limitation, it can be beneficial to assign token-level rewards for fine-grained feedback to the policy (Wu et al., 2023; Xie et al., 2024). However, assigning scores to tokens as they get autoregressively generated is computationally intensive, and collecting a large set of fine-grained human annotations for a supervised learning setup is expensive and subject to high disagreement. Thus, some works have explored ”reward shaping” techniques to transform sparse rewards into dense rewards (Sutton & Barto, 2018; Ng et al., 1999), thus facilitating more efficient and interpretable optimization, as well as providing finer-grained control over intermediate decisions. For instance, recent attempts examine process rewards (Lightman et al., 2023; Uesato et al., 2022), which provide intermediate feedback for chain-of-thought generations or directly utilize the attention map (Chan et al., 2024) of the reward model to redistribute the reward. However, these steps are complex and often require high-quality human feedback, and the attention on each token may not be directly correlated to its output as an explanation (Jain & Wallace, 2019). In this work, we propose a two-part approach to densify sparse rewards: (1) using explainability methods to construct a dense reward signal, and (2) using Bayesian optimization to learn the weights for a new reward shaping function composed of the explainability scores. We frame reward shaping as a bi-level optimization problem: at the higher level, we optimize the coefficients of the shaped reward function, while at the lower level, we learn the corresponding optimal policy, as illustrated in Figure 1. First, we estimate token-level contributions using explainability techniques such as LIME (Ribeiro et al., 2016) or SHAPley values (Lundberg & Lee, 2017). Next, since these explainability methods are known to be sensitive to noise (Li et al., 2020), we treat them as uncertain estimates and seek to learn an optimal weighting scheme over them. However, due to the complexity of the reward landscape, exhaustively evaluating all possible weight configurations is computationally infeasible. To address this, we treat the problem as black-box optimization and employ Bayesian optimization to learn the best reward coefficients as a natural method robust to noisy objective functions (Fröhlich et al., 2020; Daulton et al., 2022). We demonstrate that explainability offers a natural way to extract more information from the reward model to densify rewards and satisfies as potential-based reward shaping (Ng et al., 1999) so that we satisfy policy invariance or that the optimal policy with the original reward function remains unchanged. Furthermore, we introduce a new optimization setup incorporating different sources of token-level information and show that Bayesian Optimization can aid in learning the best reward-shaping function. Empirically, we show that explainability methods positively impact the RL training compared to sparse rewards through accelerating learning and more stable updates in the value function. We also show that adding Bayesian Optimization to properly shape rewards improves generation quality on downstream tasks compared to its naive setups. Figure 1: Overview of the bilevel optimization setup to find the best reward shape and the optimal policy. The pipeline involves an outer and inner training loop, where the outer step optimizes the Bayesian optimization model and samples the weights for our reward shape. The inner step optimizes the classic RLHF objective. 2 Preliminaries In this section, we first introduce the classic RLHF approach and build the relevant background for explainability and Bayesian optimization for our method. 2.1 Token-level MDP for RLHF We cast RL for language modeling in the sequential decision-making setting as a Markov Decision Process (MDP) by tuple ℳ=(,,,γ,r)ℳM=(S,A,P,γ,r)M = ( S , A , P , γ , r ). Here the state space SS contains states st=x,y0:tsubscriptsubscript:0s_t=\x,\,y_0:t\sitalic_t = x , y0 : t , where x denotes the input prompt and y0:tsubscript:0y_0:ty0 : t the sequence produced up to token t. The action set AA consists of next‑token choices atsubscripta_taitalic_t sampled from the model’s token distribution. Transition dynamics are governed by the kernel PP, which assigns a probability P⁢(s′∣s,a)conditionalsuperscript′P(s s,a)P ( s′ ∣ s , a ) to moving from state s to s′ after taking action a. A scalar discount factor γ∈[0,1)01γ∈[0,1)γ ∈ [ 0 , 1 ) down‑weights future returns, and the r⁢(st,at)subscriptsubscriptr(s_t,a_t)r ( sitalic_t , aitalic_t ) is the reward. In the RLHF pipeline, a model is first trained over a preference dataset =(x,yc,yr)subscriptsubscriptD=\(x,y_c,y_r)\D = ( x , yitalic_c , yitalic_r ) consisting of prompts x and a pair of chosen and rejected responses yc,yrsubscriptsubscripty_c,y_ryitalic_c , yitalic_r via supervised fine-tuning (SFT) followed by a two-stage process (1) reward modeling and (2) RL training. From the SFT model, the reward function is modeled under a preference model encoded over the sentence level and trained as a binary classification problem via maximum likelihood: P⁢(yc≻yr)=exp⁡r⁢(x,yc)exp⁡r⁢(x,yc)+exp⁡r⁢(x,yr)succeedssubscriptsubscriptsubscriptsubscriptsubscriptP(y_c y_r)= \r(x,y_c)\ \r(x,y_c)\+ \r(x,% y_r)\P ( yitalic_c ≻ yitalic_r ) = divide start_ARG exp r ( x , yitalic_c ) end_ARG start_ARG exp r ( x , yitalic_c ) + exp r ( x , yitalic_r ) end_ARG After learning the reward model, a language model is usually trained to optimize the reward via a policy gradient algorithm such as PPO (Schulman et al., 2017) with KL-regularization to maximize the expected accumulated rewards. Here, since we optimize over a finite horizon, we set γ=11γ=1γ = 1 and consider the undiscounted return: J∗⁢(πθ)=maxθ⁡at∼πθ⁢[∑t=0Tr⁢(st,at)−β⁢log⁡πθ⁢(at|st)πref⁢(at|st)|s0]superscriptsubscriptsubscriptsubscriptsimilar-tosubscriptsubscriptdelimited-[]superscriptsubscript0subscriptsubscriptconditionalsubscriptconditionalsubscriptsubscriptsubscriptrefconditionalsubscriptsubscriptsubscript0J^*( _θ)= _θE_a_t _θ [ _% t=0^Tr(s_t,a_t)-β _θ(a_t|s_t) _% ref(a_t|s_t) |s_0 ]J∗ ( πitalic_θ ) = maxitalic_θ blackboard_Ea start_POSTSUBSCRIPT t ∼ πitalic_θ end_POSTSUBSCRIPT [ ∑t = 0T r ( sitalic_t , aitalic_t ) - β log divide start_ARG πitalic_θ ( aitalic_t | sitalic_t ) end_ARG start_ARG πref ( aitalic_t | sitalic_t ) end_ARG | s0 ] where πθsubscript _θπitalic_θ is the language model agent and πrefsubscriptref _refπref is a reference policy, typically the SFT model. In practice222We use the TRL (von Werra et al., 2020) implementation of PPO for our method, the policy gradient optimizes a token-level value function given by rewards: r⁢(st,at)=r⁢(x,y)−β⁢log⁡πθ⁢(at|st)πref⁢(at|st)if ⁢y⁢ is terminal−β⁢log⁡πθ⁢(at|st)πref⁢(at|st) otherwise subscriptsubscriptcasessubscriptconditionalsubscriptsubscriptsubscriptrefconditionalsubscriptsubscriptif is terminalsubscriptconditionalsubscriptsubscriptsubscriptrefconditionalsubscriptsubscript otherwise r(s_t,a_t)= casesr(x,y)-β _θ(a_t|s_t)% _ref(a_t|s_t)&if y is terminal\\ -β _θ(a_t|s_t) _ref(a_t|s_t)&% otherwise casesr ( sitalic_t , aitalic_t ) = start_ROW start_CELL r ( x , y ) - β log divide start_ARG πitalic_θ ( aitalic_t | sitalic_t ) end_ARG start_ARG πref ( aitalic_t | sitalic_t ) end_ARG end_CELL start_CELL if y is terminal end_CELL end_ROW start_ROW start_CELL - β log divide start_ARG πitalic_θ ( aitalic_t | sitalic_t ) end_ARG start_ARG πref ( aitalic_t | sitalic_t ) end_ARG end_CELL start_CELL otherwise end_CELL end_ROW As a result, the signal received is very sparse (Zhong et al., 2025; Chan et al., 2024) in which the model receives no relevant feedback about intermediate token generations beyond the KL-regularized term. Thus, it is of interest to understand how we can design rewards on the token level to provide quality, dense signals to make learning easier and potentially improve performance. 2.2 Explainability as Token-level Rewards We first introduce a general definition for estimating token-level information with explainability in the textual environment. To explain the predictions of a complex black-box function f (i.e., our reward model), often a small linear model g is used to locally approximate it concerning an input x of the prompt and completion. Here, we consider only the family of additive feature attribution functions (Ribeiro et al., 2016; Lundberg & Lee, 2017) such that the optimal policy with the original reward remains unchanged (Ng et al., 1999): g⁢(z′)=ϕ0+∑i=1Mϕi⁢zi′subscriptitalic-ϕ0superscriptsubscript1subscriptitalic-ϕsubscriptsuperscript′g(z )= _0+ _i=1^M _iz _ig ( z′ ) = ϕ0 + ∑i = 1M ϕitalic_i z′italic_i (1) In the text domain, we define zi′∈0,1Msubscriptsuperscript′01z _i∈\0,1\^Mz′italic_i ∈ 0 , 1 M as a binary mask that masks the i-th token from the original input x where M is the number of tokens in the sequence such that g⁢(z′)≈f⁢(hx⁢(z′))superscript′subscriptℎsuperscript′g(z )≈ f(h_x(z ))g ( z′ ) ≈ f ( hitalic_x ( z′ ) ) and hxsubscriptℎh_xhitalic_x is mapping a function that reconstructs z′ into a valid text input. The value of interest for each token is ϕisubscriptitalic-ϕ _iϕitalic_i, representing the feature effect or their marginal contribution to the prediction f⁢(x)f(x)f ( x ). To solve for the explainability score ϕisubscriptitalic-ϕ _iϕitalic_i for each token, we optimize the local accuracy to sample x by minimizing the penalized linear regression problem: ℰ=arg⁢ming∈G⁡ℒ⁢(f,g,πx)+Ω⁢(g)ℰsubscriptargminℒsubscriptΩE= *arg\,min_g∈ GL(f,g, _x)+ (g)E = start_OPERATOR arg min end_OPERATORg ∈ G L ( f , g , πitalic_x ) + Ω ( g ) (2) where G defines the set of local functions (i.e., small linear model), πxsubscript _xπitalic_x is a kernel transformation to the input (i.e., the SHAPley kernel in Eq. 6), and Ω Ω is a complexity penalty of g. Solving the regression problem gives us ℰEE or the set of explanation scores per token ϕisubscriptitalic-ϕ _iϕitalic_i for a given sequence x. The form of Eq. 2 is general for any method of additive feature attribution functions, and we see later that formulating this as a reward-shaping function optimizes for the same policy as the original sparse reward objective in Section 3.1. We detail a complete example of solving the explainability objective in Appendix A to produce token-level scores. 2.3 Bayesian Optimization for Reward Shaping To optimize the token-level weights in our reward shaping function, we employ Bayesian Optimization (BO) as a zeroth-order method to efficiently search for high-reward configurations. BO aims to find the global maximizer x∗=arg⁡maxx⁡f⁢(x)superscriptsubscriptx^*= _xf(x)x∗ = arg maxitalic_x f ( x ) of a black-box function f, which in our case represents the scalar reward produced by a shaped reward model under a given weight configuration. Since f is expensive to evaluate, BO uses a surrogate model—typically a Gaussian Process (GP)—to approximate f from past observations and suggest new candidate points. In this work, it is sufficient to understand the GP as a distribution over a collection of function outputs mapping the observed function values y=f⁢(x)y=f(x)y = f ( x ) (here, the reward prediction) to a multivariate Normal: f⁢(x)∼⁢(μ⁢(x),Σ⁢(x))similar-toΣf(x) (μ(x), (x))f ( x ) ∼ N ( μ ( x ) , Σ ( x ) ), where μ⁢(x)=⁢[y|x]delimited-[]conditionalμ(x)=E[y|x]μ ( x ) = blackboard_E [ y | x ] is the mean function, and Σ⁢(x)=Cov⁢[y|x]ΣCovdelimited-[]conditional (x)=Cov[y|x]Σ ( x ) = Cov [ y | x ] is the covariance function. To update the GP model, we can combine the observed values with the prior distribution to update the belief about f and generate a posterior distribution, which we continue to do iteratively. After building the surrogate model, BO samples points by maximizing the expected utility333Formally, it is usually also referred to as the ”reward” but we separate definitions to avoid confusion from RL rewards. based on the current posterior under some acquisition function to sample areas of improvement and balance between exploration and exploitation to choose the next point. At each BO step, we use the surrogate to select a new candidate weight vector by maximizing an acquisition function u, which quantifies the expected utility (i.e., predicted improvement in validation reward) of sampling a new point: w=arg⁡maxw⁡u⁢(w∣1:t−1)subscriptconditionalsubscript:11w= _wu(w _1:t-1)w = arg maxitalic_w u ( w ∣ D1 : t - 1 ) (3) where 1:t−1subscript:11D_1:t-1D1 : t - 1 is accumulated observations from the objective function. For u, we adopt log Noisy Expected Improvement (Ament et al., 2025), which handles noisy evaluations well and is easy to optimize numerically. We implement BO using the Ax API (Bakshy et al., 2018), which manages the optimization loop for the surrogate GP and acquisition functions. 3 Explainable Reward Shaping as a Bilevel Optimization In this section, we discuss the optimization setup and the preliminary algorithm to find the best reward shape from each token-level score. First, we discuss the formulation of reward shape to compute a dense reward from locally approximating token-level rewards via explainability methods and obtain more fine-grained feedback. Next, we discuss the bilevel optimization problem, where we model BO and the classic RLHF problem as a nested objective to optimize both the reward shape and the policy. 3.1 Computing the Reward Shape At each step of the optimization procedure, after sampling the weights w from the upper-level step, we shape the reward by computing a weighted linear combination of all token-level scores and broadcast the reward over the sequence: r′⁢(s,a)superscript′ r (s,a)r′ ( s , a ) =W⊤⁢⋅r⁢(s,a)=w1⋅ℰSHAP⋅r⁢(s,a)+w2⋅r⁢(s,a)absent⋅superscripttop⋅subscript1subscriptℰSHAP⋅subscript2 =W E· r(s,a)=w_1·% E_SHAP· r(s,a)+w_2· r(s,a)= W⊤ caligraphic_E ⋅ r ( s , a ) = w1 ⋅ ESHAP ⋅ r ( s , a ) + w2 ⋅ r ( s , a ) (4) where =[ℰSHAP,1]subscriptℰSHAP1 E=[E_SHAP,1]caligraphic_E = [ ESHAP , 1 ] is the matrix of token-level explanations (not limited to just SHAP); W is the vector of weights sampled from the Bayesian optimization step. Importantly, we sample each wisubscriptw_iwitalic_i from the BO model such that ∑iwi=1subscriptsubscript1 _iw_i=1∑i witalic_i = 1 to consider the convex combination of all token-level scores over the reward. In practice, we compute the dense reward by first taking the softmax over the token-level scores to get a probability vector over the sequence. Then, following Chan et al. (2024), we distribute the sparse reward or assign credit to each token by broadcasting it over the sequence. As a result, the scores at intermediate states are assigned to provide fine-grained feedback for the next token generated transitioning from state s→s′→superscript′s→ s s → s′. Additionally, we show that any explainability method in the family of additive feature attribution functions follows policy invariance and can be formulated as a potential-based shaping function (Ng et al., 1999). This ensures that adding a reward for transitions between states does not yield a suboptimal policy but better guides the learning agent towards the same optimal policy. We detail a complete proof in Appendix B. Figure 2: Redistribution sequence of the scalar reward prediction over the explanation feature attributions after softmax normalization. A darker red highlights a much stronger positive contribution, while a deeper blue indicates a more negative contribution. 3.2 BO and RLHF as Nested Problems Conventionally, the lower level is an auxiliary problem that helps solve the upper-level problem (Zhang et al., 2023).In contrast, here, we are equally interested in the solution to the lower-level problem as a result of the solution retrieved from the upper-level one. In other words, we have the following bilevel optimization setup: maxw⁡f⁢(w;πθ∗⁢(w))s.t.πθ∗⁢(w)=arg⁢maxπθ∈Π⁡a∼πθ⁢[r′⁢(s,a)−β⁢log⁡πθ⁢(a|s)πref⁢(a|s)]subscriptsuperscriptsubscripts.t.superscriptsubscriptsubscriptargmaxsubscriptΠsubscriptsimilar-tosubscriptdelimited-[]superscript′subscriptconditionalsubscriptrefconditional gathered _wf(w; _θ^*(w)) .t. _% θ^*(w)= *arg\,max_ _θ∈ E_a % _θ [r (s,a)-β _θ(a|s) _% ref(a|s) ] gatheredstart_ROW start_CELL maxitalic_w f ( w ; πitalic_θ∗ ( w ) ) s.t. πitalic_θ∗ ( w ) = start_OPERATOR arg max end_OPERATORπ start_POSTSUBSCRIPT θ ∈ Π end_POSTSUBSCRIPT blackboard_Ea ∼ π start_POSTSUBSCRIPT θ end_POSTSUBSCRIPT [ r′ ( s , a ) - β log divide start_ARG πitalic_θ ( a | s ) end_ARG start_ARG πref ( a | s ) end_ARG ] end_CELL end_ROW (5) where we have the upper-level optimization problem maximizing the acquisition function f in each BO iteration to sample weights w based on the validation set performance (i.e., average validation reward) of the best-performing policy πθ∗subscriptsuperscriptπ^*_θπ∗italic_θ. In the lower-level optimization step, we have the standard RLHF problem to find the best policy πθ∗superscriptsubscript _θ^*πitalic_θ∗ given the weights for reward shaping from the upper-level step. In principle, our optimization procedure can be more concretely understood as a hyperparameter tuning problem, as we do not explicitly parameterize the reward transformation, but where the outer step looks to find the best set of weights w (or hyperparameters) to compute the reward shape for the optimal policy. We outline a practical implementation of our method in Algorithm 1. Models # H-RLHF Ultrafeedback Score Win (%) Score Win (%) MTBench Baselines →SFT - 5.40 – 4.34 – 7.2 →RLHF with Sparse (original) - 5.94 48.7 4.35 57.64 7.17 Dense RLHF →Attention (Chan et al., 2024) 1 6.24 51.40 4.42 61.40 7.28 →SHAP* 1 5.91 53.95 4.46 53.81 7.81 →LIME* 1 5.86 54.36 4.46 59.20 7.28 Dense RLHF with BO →BO-SHAP 1 5.99 55.75 4.37 60.82 7.37 →BO-LIME 1 6.26 57.62 4.37 56.73 7.59 1-7 →BO-SHAP-Attn 2 6.47 56.90 4.48 51.08 7.71 →BO-LIME-Attn 2 6.30 54.09 4.45 55.03 7.03 →BO-SHAP-LIME 2 6.58 56.28 4.51 64.82 7.61 1-7 →BO-SHAP-LIME-Attn 3 5.74 51.00 4.44 59.64 7.38 Table 1: Comparison of H-RLHF (Only Helpful) and Ultrafeedback. Winrate is calculated via AlpacaEval, with the reference model being the SFT model’s generation. * indicates strong baselines that we have implemented. # means the number of dense reward types used. Score refers to the average test set reward and Win (%) refers to the length-controlled win-rate against the baseline SFT generations. We use gpt4-turbo as the judge for both AlpacaEval2 and MTBench. Bolded and underlined numbers highlight the best performance. 4 Experiments Tasks. In this section, we empirically verify the effectiveness of our method on downstream tasks. Namely, we train each model on two different single-turn dialogue datasets: 1) H-RLHF (helpfulness) (Bai et al., 2022) and 2) Ultrafeedback (Cui et al., 2024) datasets. For H-RLHF, we utilize the OpenLLaMA family of models with an instruction-tuned 7B model as the SFT model and the 3B parameter reward model from Dong et al. (2023) following the setup in (Chan et al., 2024). For Ultrafeedback experiments, we utilize LLaMa-3.2-Instruct 3B (Team, 2024) as the SFT model and fine-tune LLaMa-3.2-Instruct 1B on Ultrafeedback preferences as the reward model. To verify the effectiveness of each method, we consider first the average holistic reward predicted by the reward model over the test splits for each dataset. We also evaluate on open benchmarks such as AlpacaEval-2 (Dubois et al., 2025) and MTBench (Zheng et al., 2023a) to examine whether our methods lead to a better local optimum versus optimizing with only the sparse reward while also avoiding reward overfitting. Baselines To measure the improvement of our method, we consider several different baselines: 1) the SFT model to calibrate improvements, 2) RLHF on the sparse reward, 3) attention-based credit (Chan et al., 2024), and 4) each individual explanation method before optimizing the reward shape with Bayesian Optimization. Training setup For every experiment, we adopt PPO as our method’s principle policy gradient algorithm. Due to the computational complexity of PPO and to run a sufficient number of Bayesian trials, we must compromise training samples per trial to avoid inducing too much computation overhead. Hence, instead of running over the whole dataset for each BO trial, we run m=2525m=25m = 25 trials for the Bayesian Optimization step and batch each trial into a small number of training epochs (i.e., 10) of batch size 8 (80 samples) for a total of 2000 samples. To initially sample weights for reward shaping, we employ Sobol sampling for the first five trials to build a prior for the Gaussian Process model and then sample from the GP model for the remaining trials. After each BO iteration, we evaluate the trained model on the validation split 444We construct the validation set by splitting the training dataset 90%/10%. and take the average validation reward as the utility to update the BO surrogate function. To continue training, we employ model checkpointing to resume training from the best checkpoint instead of running a complete training cycle from scratch. Thus, we run at most only twice over the dataset 1:nsubscript:1D_1:nD1 : n (once for BO, once for full PPO). To ensure randomness, we randomly sample a subset from the training dataset at each Bayesian trial to train PPO. We also randomly sample a subset from the validation set at each validation step. We detail our model setups and hyperparameters and discuss the computational complexity in Appendix C and D. (a) SHAP vs. BO variants (b) Validation reward after each BO trial. (c) Mean dense reward shape Figure 3: (Left) The mean training reward per timestep with increasing BO dimensionality. (Middle) The mean validation reward over each BO trial. The highlighted dots indicate the best validation reward received at trial n. (Right) The average dense reward attribution over each trial for SHAP + ATTN. The highlighted row indicates the shape in trial 17 that received the highest validation reward. (a) Explanation-based rewards training curves vs. baselines (b) PPO Value loss Figure 4: Helpfulness. (Left) The average training reward per timestep. (Right) The average value head loss per timestep. The shading represents the standard error (95%percent9595\%95 % confidence interval) as training progresses. 4.1 How well do explanation-based rewards optimize the RLHF objective? We first analyze the impact of explanation-based rewards in optimizing the vanilla RLHF objective. First, we do not expect to see any significant improvements from explanation methods alone. As we formulate explanation scores as a potential-based reward shaping transformation, we are not guaranteed to find better local optima Ng et al. (1999), but a carefully chosen reward shape could improve the learning complexity of reinforcement learning (Fu et al., 2025; Gupta et al., 2022). Foremost, we are interested in whether explainability can provide any meaningful improvements or optimizations over simply considering the sparse, scalar reward. Credit assignment with explainability helps exploration Before applying any Bayesian optimization, we manually select a weighting combination, placing higher emphasis on the token-level scores with w=0.80.8w=0.8w = 0.8 following Chan et al. (2024). Overall, explanation-based methods achieve a high average reward relatively early during training that is competitive with the attention-based baseline observed in Figure 4(a). Furthermore, explanation-based methods achieve more stable updates than sparse rewards from Figure 4(b). In particular, the dense rewards from explainability drastically reduce the PPO value head loss, signaling that we have a good approximation of states to long-term returns for more stable policy updates during training. Explanation-based rewards avoid reward overfitting One of the main risks associated in reinforcement learning from human feedback is reward overfitting. This is particularly undesirable since the reward model is a proxy for human preferences, and overoptimizing its value can hinder ground truth performance (Gao et al., 2022). To study this relationship, we observe the performance of baselines and our methods on open benchmarks for instruction-following and multi-turn dialogue capabilities, which contain data samples not strictly in distribution to our training datasets. From Table 3.2, we see that explanation-based methods do not substantially outperform baselines on the H-RLHF test split; however, they achieve a markedly higher win rate on open benchmarks, indicating that we are not particularly overfitting to the reward function and also maintaining general performance. Notably, we see the sparse reward approach underperforms relative to the SFT model on AlpacaEval2 for the H-RLHF models. We suspect this stems from overfitting on the reward function, which can degrade performance on data splits not directly related to the helpfulness split. To verify this, we examine the win rate against the SFT model only the helpfulness split of AlpacaEval2 and find the win rate improves, being preferred 56.34%±4.03plus-or-minuspercent56.344.0356.34\%± 4.0356.34 % ± 4.03 of the time. Figure 5: (Top) The weight transition between trials for SHAPley scores. (Bottom) The weight transition between trials for LIME scores. The black boxes indicate the ”best weights” sampled by the BO model. 4.2 Does Bayesian optimization help balance token-level rewards better? Next, we evaluate the impact of Bayesian Optimization (BO) in sampling optimal weights to construct a new, dense reward function based on token-level interpretability scores. We consider several combinations of token-level scores derived from SHAP, LIME, and the attention map of the reward model. Specifically, our goal is to assess whether BO can effectively weigh these different sources of token-level importance to produce a more informative reward signal during policy optimization. However, due to the limited number of trials available, we acknowledge that convergence of BO is not guaranteed, particularly given the high-dimensional and potentially non-convex search space of reward weight combinations (Loeppky et al., 2012; Snoek et al., 2012). The challenge of balancing exploration and exploitation in such a setting further complicates the optimization process, especially when evaluating noisy, non-stationary reward surfaces. Despite these limitations, Table 1 and Figure 3(a) show that the weights sampled by BO generally lead to improved downstream task performance compared to baselines and that models guided by BO-optimized dense rewards tend to achieve even higher training rewards. Explore-Exploit Tradeoff To better understand BO’s behavior, we analyze the evolution of weight distributions and dense reward shape across trials in Figure 5 and 3(c). In particular, we highlight that the final combination incorporating all scores with an input dimension of d=44d=4d = 4 does not outperform simpler combinations. In particular, this is expected as by increasing the search space complexity, we should also increase the number of trials; however, due to the computation constraints, we keep the number of trials constant. Ideally, if BO had successfully identified optimal weights, we would expect to see a lower performance bound on any sub-combination of scores; that is, more scores should not hurt performance if adequately weighted. However, the observed degradation in performance for d=44d=4d = 4 implies that BO may not have sufficiently explored or exploited the reward space, potentially due to early convergence to suboptimal regions. Across all trials in Figure 5, we see that after Sobol sampling from the first five trials, BO tends to consistently switch between exploration and exploitation, suggesting that just more exploration is needed. Figure 6: (Ultrafeedback) The top left represents the baseline per-token reward without shaping. The color of each token represents the reward received, with a darker color representing a higher proportion of the reward assigned. A more uniform coloring indicates a more uniform assignment of the scalar reward to each token, while a contrasting light/dark coloring indicates a more skewed assignment. BO and Credit Assignment Additionally, we inspect the before and after of the reward shape on a representative example from the Ultrafeedback (Cui et al., 2024) dataset, where the model is asked to choose two sentences out of three that seamlessly connect with an unfinished story. Here, in Figure 6, we show several variants of the shaping method on a completion from the SFT model where the tokens are colored proportionally to the reward they received after computing the shaping weights. In particular, almost all shaping methods still place importance on the terminating token. We suspect that since the reward model is trained to backpropogate on the terminal token, we cannot ignore it in order to maximize the reward function best. Specifically, we see the combination of SHAP + LIME, which highlights the explanation of the model’s rationale for choosing the first and second sentences. We similarly see this for the d=44d=4d = 4 combination with all token-level scores. However, we see a more uniform assignment of scores as the scalar reward was given zero weight in the final shaped score. Limitations of Current Explainability Methods While Bayesian Optimization (BO) can leverage complementary effects from multiple explainability signals, individual mechanistic interpretability methods, such as a model’s intrinsic attention, remain fundamentally suboptimal. As shown in Figure 6, these methods often fail to align precisely with human perception of relevant features Bereska & Gavves (2024). A promising direction for future research is to incorporate more fine-grained, human-aligned signals, potentially sourced from explicit token-level annotations Hayati et al. (2021) or cognitive indicators like eye movements when reading text De Langis & Kang (2022). 5 Related Work Reward Shaping is a technique in reinforcement learning (RL) that supplements traditionally sparse rewards with more informative reward signals (Hu et al., 2020). Early work in potential-based reward shaping (PBRS) demonstrates that reward shaping can preserve the optimal policy while reducing training time or introducing domain knowledge (Hu et al., 2020; Cao et al., 2024a). However, challenges still arise because it is often difficult to guarantee that the additional reward is helpful for any given task (Hu et al., 2020). Additionally, we note that we are not the first to inspect reward shaping or RLHF under a bi-level optimization lens. For example, bi-level optimization of parameterized reward shaping (BiPaRS) (Hu et al., 2020) adaptively utilizes a given reward function, optimizing both the policy itself, as well as the shaping weight for the reward function. Shen et al. (2024) Also inspects RLHF as a bilevel formulation under a penalty function. Other reward shaping methods focus on adaptability to overcome limitations, such as AlphaPO (Gupta et al., 2025) builds on simple policy optimization (SimPO) by introducing a new scalar parameter α, demonstrating that the shape of the reward function can alter the likelihood displacement. Token-Level RLHF Other methods aim to improve reward signals by distributing the reward at the token-level, hoping to capture more nuance in particular words or phrases (Yoon et al., 2024) or discourses (Kim et al., 2025) in the generation ; examples include Token-Level Continuous Reward (TLCR), which uses a token-level preference discriminator (Yoon et al., 2024), RLMEC which leverages minimum editing constraints to produce token-level supervision (Chen et al., 2024b), and DRLC, which employs LLMs to identify dense positive and negative labelings within a response (Cao et al., 2024b). Other work has also explored using LLMs themselves to provide the dense reward signal (Huang et al., 2024). A common theme with these works is training contrastively: an initial output is compared with a revised output, and the fine-grained reward signal is achieved by comparing the differences (Guo et al., 2023). Methods such as token-level PPO (TPPO) (Ouyang et al., 2024) and token-level DPO (TDPO) (Zeng et al., 2024) further support the approach of token-level fine-tuning. Bayesian Optimization (BO) has also emerged to augment LLM training and fine-tuning (Austin et al., 2024; Liu et al., 2024; Yang et al., ; Gao et al., 2024; Agarwal et al., 2025; Kristiadi et al., 2024; Opsahl-Ong et al., 2024; Chen et al., 2024a). BO functions by iteratively estimating the optimal posterior predictive distribution by adding new information/candidates (Agarwal et al., 2025). Some recent work uses BO with LLMs to perform selection across a discrete search space (Kristiadi et al., 2024; Opsahl-Ong et al., 2024). Other work investigates using BO to optimize prompts for LLMs, such as BOPRO (Agarwal et al., 2025), InstructZero (Chen et al., 2024a), leveraging the Bayesian Optimization to explore the prompt search space. BO can also be leveraged in pre-training, by finding the optimal weighting for checkpoint merging (Liu et al., 2024). BO can also be leveraged for uncertainty estimation (Yang et al., ), bias mitigation (via win rate calibration) (Gao et al., 2024), and guided query generation (Austin et al., 2024). These tasks ultimately benefit from BO’s robustness to noise (Opsahl-Ong et al., 2024) by calibrating for uncertainty in the optimization process. 6 Conclusion In this work, we show that explainability offers an effective way to reshape rewards for token-level credit assignments. Furthermore, our findings suggest that BO is a promising tool for learning and balancing different sources of token-level information for reward shaping. As a combination, we show that explainable dense reward shapes optimized under a Bayesian framework positively impact RL training compared to sparse rewards by accelerating learning and providing more stable updates to the value function. Furthermore, properly attributing tokens via interpretable methods also improves performance over the sparse reward on downstream tasks while also theoretically being policy invariant to stay faithful to the original reward model. In the future, we could consider incorporating contextual information into the BO weights to dynamically shape token-level rewards. For instance, we could learn a set of high-dimensional weights corresponding to the token-embeddings of a policy model’s generation to dynamically compute the reward shape instead of the current statically weighted reward shape. Acknowledgements This work was mainly supported by the research gift from Grammarly and UMN Data Science Institute (DSI) Seed Grant. We also thank Minnesota NLP group members for providing us with valuable feedback and comments on the initial draft. References Agarwal et al. (2025) Dhruv Agarwal, Manoj Ghuhan Arivazhagan, Rajarshi Das, Sandesh Swamy, Sopan Khosla, and Rashmi Gangadharaiah. Searching for optimal solutions with LLMs via bayesian optimization. In The Thirteenth International Conference on Learning Representations, 2025. URL https://openreview.net/forum?id=aVfDrl7xDV. Ament et al. (2025) Sebastian Ament, Samuel Daulton, David Eriksson, Maximilian Balandat, and Eytan Bakshy. Unexpected improvements to expected improvement for bayesian optimization, 2025. URL https://arxiv.org/abs/2310.20708. Austin et al. (2024) David Austin, Anton Korikov, Armin Toroghi, and Scott Sanner. Bayesian optimization with llm-based acquisition functions for natural language preference elicitation. In Proceedings of the 18th ACM Conference on Recommender Systems, RecSys ’24, p. 74–83, New York, NY, USA, 2024. Association for Computing Machinery. ISBN 9798400705052. doi: 10.1145/3640457.3688142. URL https://doi.org/10.1145/3640457.3688142. Bai et al. (2022) Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, Nicholas Joseph, Saurav Kadavath, Jackson Kernion, Tom Conerly, Sheer El-Showk, Nelson Elhage, Zac Hatfield-Dodds, Danny Hernandez, Tristan Hume, Scott Johnston, Shauna Kravec, Liane Lovitt, Neel Nanda, Catherine Olsson, Dario Amodei, Tom Brown, Jack Clark, Sam McCandlish, Chris Olah, Ben Mann, and Jared Kaplan. Training a helpful and harmless assistant with reinforcement learning from human feedback, 2022. URL https://arxiv.org/abs/2204.05862. Bakshy et al. (2018) Eytan Bakshy, Lili Dworkin, Brian Karrer, Konstantin Kashin, Ben Letham, Ashwin Murthy, and Shaun Singh. Ae: A domain-agnostic platform for adaptive experimentation. In NeurIPS Systems for ML Workshop, 2018. URL http://learningsys.org/nips18/assets/papers/87CameraReadySubmissionAE%20-%20NeurIPS%202018.pdf. Bereska & Gavves (2024) Leonard Bereska and Efstratios Gavves. Mechanistic interpretability for ai safety–a review. arXiv preprint arXiv:2404.14082, 2024. Cao et al. (2024a) Meng Cao, Lei Shu, Lei Yu, Yun Zhu, Nevan Wichers, Yinxiao Liu, and Lei Meng. Enhancing reinforcement learning with dense rewards from language model critic. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, p. 9119–9138, Miami, Florida, USA, November 2024a. Association for Computational Linguistics. doi: 10.18653/v1/2024.emnlp-main.515. URL https://aclanthology.org/2024.emnlp-main.515/. Cao et al. (2024b) Meng Cao, Lei Shu, Lei Yu, Yun Zhu, Nevan Wichers, Yinxiao Liu, and Lei Meng. Enhancing reinforcement learning with dense rewards from language model critic. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, p. 9119–9138, Miami, Florida, USA, November 2024b. Association for Computational Linguistics. doi: 10.18653/v1/2024.emnlp-main.515. URL https://aclanthology.org/2024.emnlp-main.515/. Chan et al. (2024) Alex J. Chan, Hao Sun, Samuel Holt, and Mihaela van der Schaar. Dense reward for free in reinforcement learning from human feedback, 2024. URL https://arxiv.org/abs/2402.00782. Chaudhari et al. (2024) Shreyas Chaudhari, Pranjal Aggarwal, Vishvak Murahari, Tanmay Rajpurohit, Ashwin Kalyan, Karthik Narasimhan, Ameet Deshpande, and Bruno Castro da Silva. Rlhf deciphered: A critical analysis of reinforcement learning from human feedback for llms, 2024. URL https://arxiv.org/abs/2404.08555. Chen et al. (2024a) Lichang Chen, Jiuhai Chen, Tom Goldstein, Heng Huang, and Tianyi Zhou. InstructZero: Efficient instruction optimization for black-box large language models. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, p. 6503–6518. PMLR, 21–27 Jul 2024a. URL https://proceedings.mlr.press/v235/chen24e.html. Chen et al. (2024b) Zhipeng Chen, Kun Zhou, Xin Zhao, Junchen Wan, Fuzheng Zhang, Di Zhang, and Ji-Rong Wen. Improving large language models via fine-grained reinforcement learning with minimum editing constraint. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Findings of the Association for Computational Linguistics: ACL 2024, p. 5694–5711, Bangkok, Thailand, August 2024b. Association for Computational Linguistics. doi: 10.18653/v1/2024.findings-acl.338. URL https://aclanthology.org/2024.findings-acl.338/. Cui et al. (2024) Ganqu Cui, Lifan Yuan, Ning Ding, Guanming Yao, Bingxiang He, Wei Zhu, Yuan Ni, Guotong Xie, Ruobing Xie, Yankai Lin, Zhiyuan Liu, and Maosong Sun. Ultrafeedback: Boosting language models with scaled ai feedback, 2024. URL https://arxiv.org/abs/2310.01377. Daulton et al. (2022) Samuel Daulton, Sait Cakmak, Maximilian Balandat, Michael A. Osborne, Enlu Zhou, and Eytan Bakshy. Robust multi-objective bayesian optimization under input noise, 2022. URL https://arxiv.org/abs/2202.07549. De Langis & Kang (2022) Karin De Langis and Dongyeop Kang. A comparative study on textual saliency of styles from eye tracking, annotations, and language models. arXiv preprint arXiv:2212.09873, 2022. Dong et al. (2023) Hanze Dong, Wei Xiong, Deepanshu Goyal, Yihan Zhang, Winnie Chow, Rui Pan, Shizhe Diao, Jipeng Zhang, Kashun Shum, and Tong Zhang. Raft: Reward ranked finetuning for generative foundation model alignment, 2023. URL https://arxiv.org/abs/2304.06767. Dubois et al. (2025) Yann Dubois, Balázs Galambosi, Percy Liang, and Tatsunori B. Hashimoto. Length-controlled alpacaeval: A simple way to debias automatic evaluators, 2025. URL https://arxiv.org/abs/2404.04475. Engstrom et al. (2020) Logan Engstrom, Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Firdaus Janoos, Larry Rudolph, and Aleksander Madry. Implementation matters in deep policy gradients: A case study on ppo and trpo, 2020. URL https://arxiv.org/abs/2005.12729. Fröhlich et al. (2020) Lukas P. Fröhlich, Edgar D. Klenske, Julia Vinogradska, Christian Daniel, and Melanie N. Zeilinger. Noisy-input entropy search for efficient robust bayesian optimization, 2020. URL https://arxiv.org/abs/2002.02820. Fu et al. (2025) Jiayi Fu, Xuandong Zhao, Chengyuan Yao, Heng Wang, Qi Han, and Yanghua Xiao. Reward shaping to mitigate reward hacking in rlhf, 2025. URL https://arxiv.org/abs/2502.18770. Gao et al. (2022) Leo Gao, John Schulman, and Jacob Hilton. Scaling laws for reward model overoptimization, 2022. URL https://arxiv.org/abs/2210.10760. Gao et al. (2024) Yicheng Gao, Gonghan Xu, Zhe Wang, and Arman Cohan. Bayesian calibration of win rate estimation with LLM evaluators. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, p. 4757–4769, Miami, Florida, USA, November 2024. Association for Computational Linguistics. doi: 10.18653/v1/2024.emnlp-main.273. URL https://aclanthology.org/2024.emnlp-main.273/. Guo et al. (2023) Geyang Guo, Ranchi Zhao, Tianyi Tang, Wayne Xin Zhao, and Ji-Rong Wen. Beyond imitation: Leveraging fine-grained quality signals for alignment. arXiv preprint arXiv:2311.04072, 2023. Gupta et al. (2022) Abhishek Gupta, Aldo Pacchiano, Yuexiang Zhai, Sham M. Kakade, and Sergey Levine. Unpacking reward shaping: Understanding the benefits of reward engineering on sample complexity, 2022. URL https://arxiv.org/abs/2210.09579. Gupta et al. (2025) Aman Gupta, Shao Tang, Qingquan Song, Sirou Zhu, Jiwoo Hong, Ankan Saha, Viral Gupta, Noah Lee, Eunki Kim, Siyu Zhu, Parag Agrawal, Natesh Pillai, and S. Sathiya Keerthi. Alphapo – reward shape matters for llm alignment, 2025. URL https://arxiv.org/abs/2501.03884. Hayati et al. (2021) Shirley Anugrah Hayati, Dongyeop Kang, and Lyle Ungar. Does bert learn as humans perceive? understanding linguistic styles through lexica. arXiv preprint arXiv:2109.02738, 2021. Hu et al. (2020) Yujing Hu, Weixun Wang, Hangtian Jia, Yixiang Wang, Yingfeng Chen, Jianye Hao, Feng Wu, and Changjie Fan. Learning to utilize shaping rewards: a new approach of reward shaping. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS ’20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829546. Huang et al. (2024) Chengyu Huang, Zeqiu Wu, Yushi Hu, and Wenya Wang. Training language models to generate text with citations via fine-grained rewards. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), p. 2926–2949, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: 10.18653/v1/2024.acl-long.161. URL https://aclanthology.org/2024.acl-long.161/. Jain & Wallace (2019) Sarthak Jain and Byron C. Wallace. Attention is not explanation, 2019. URL https://arxiv.org/abs/1902.10186. Kim et al. (2025) Zae Myung Kim, Anand Ramachandran, Farideh Tavazoee, Joo-Kyung Kim, Oleg Rokhlenko, and Dongyeop Kang. Align to structure: Aligning large language models with structural information, 2025. Kristiadi et al. (2024) Agustinus Kristiadi, Felix Strieth-Kalthoff, Marta Skreta, Pascal Poupart, Alán Aspuru-Guzik, and Geoff Pleiss. A sober look at llms for material discovery: are they actually good for bayesian optimization over molecules? In Proceedings of the 41st International Conference on Machine Learning, ICML’24. JMLR.org, 2024. Li et al. (2020) Xiaoxiao Li, Yuan Zhou, Nicha C. Dvornek, Yufeng Gu, Pamela Ventola, and James S. Duncan. Efficient shapley explanation for features importance estimation under uncertainty. In Medical Image Computing and Computer Assisted Intervention – MICCAI 2020: 23rd International Conference, Lima, Peru, October 4–8, 2020, Proceedings, Part I, p. 792–801, Berlin, Heidelberg, 2020. Springer-Verlag. ISBN 978-3-030-59709-2. doi: 10.1007/978-3-030-59710-8˙77. URL https://doi.org/10.1007/978-3-030-59710-8_77. Lightman et al. (2023) Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let’s verify step by step, 2023. URL https://arxiv.org/abs/2305.20050. Liu et al. (2024) Deyuan Liu, Zecheng Wang, Bingning Wang, Weipeng Chen, Chunshan Li, Zhiying Tu, Dianhui Chu, Bo Li, and Dianbo Sui. Checkpoint merging via bayesian optimization in llm pretraining. CoRR, 2024. Loeppky et al. (2012) Jason L Loeppky, Jerome Sacks, and William J Welch. Choosing the sample size of a computer experiment: A practical guide. Technometrics, 54(4):435–446, 2012. Lundberg & Lee (2017) Scott M Lundberg and Su-In Lee. A unified approach to interpreting model predictions. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https://proceedings.neurips.c/paper_files/paper/2017/file/8a20a8621978632d76c43dfd28b67767-Paper.pdf. López & Saboyá (2009) S. López and Martha Saboyá. On the relationship between shapley and owen values. Central European Journal of Operations Research, 17:415–423, 12 2009. doi: 10.1007/s10100-009-0100-8. Ng et al. (1999) Andrew Y. Ng, Daishi Harada, and Stuart J. Russell. Policy invariance under reward transformations: Theory and application to reward shaping. In Proceedings of the Sixteenth International Conference on Machine Learning, ICML ’99, p. 278–287, San Francisco, CA, USA, 1999. Morgan Kaufmann Publishers Inc. ISBN 1558606122. Opsahl-Ong et al. (2024) Krista Opsahl-Ong, Michael J Ryan, Josh Purtell, David Broman, Christopher Potts, Matei Zaharia, and Omar Khattab. Optimizing instructions and demonstrations for multi-stage language model programs. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, p. 9340–9366, Miami, Florida, USA, November 2024. Association for Computational Linguistics. doi: 10.18653/v1/2024.emnlp-main.525. URL https://aclanthology.org/2024.emnlp-main.525/. Ouyang et al. (2024) Yichen Ouyang, Lu Wang, Fangkai Yang, Pu Zhao, Chenghua Huang, Jianfeng Liu, Bochen Pang, Yaming Yang, Yuefeng Zhan, Hao Sun, et al. Token-level proximal policy optimization for query generation. arXiv preprint arXiv:2411.00722, 2024. Rafailov et al. (2024) Rafael Rafailov, Joey Hejna, Ryan Park, and Chelsea Finn. From r to q∗superscriptq^*q∗: Your language model is secretly a q-function, 2024. URL https://arxiv.org/abs/2404.12358. Razin et al. (2024) Noam Razin, Hattie Zhou, Omid Saremi, Vimal Thilak, Arwen Bradley, Preetum Nakkiran, Joshua Susskind, and Etai Littwin. Vanishing gradients in reinforcement finetuning of language models, 2024. URL https://arxiv.org/abs/2310.20703. Ribeiro et al. (2016) Marco Tulio Ribeiro, Sameer Singh, and Carlos Guestrin. ”why should i trust you?”: Explaining the predictions of any classifier, 2016. URL https://arxiv.org/abs/1602.04938. Schulman et al. (2017) John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017. URL https://arxiv.org/abs/1707.06347. Shen et al. (2024) Han Shen, Zhuoran Yang, and Tianyi Chen. Principled penalty-based methods for bilevel reinforcement learning and rlhf, 2024. URL https://arxiv.org/abs/2402.06886. Snoek et al. (2012) Jasper Snoek, Hugo Larochelle, and Ryan P. Adams. Practical bayesian optimization of machine learning algorithms, 2012. URL https://arxiv.org/abs/1206.2944. Sutton & Barto (2018) Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. The MIT Press, second edition, 2018. URL http://incompleteideas.net/book/the-book-2nd.html. Team (2024) Meta Llama Team. The llama 3 herd of models, 2024. URL https://arxiv.org/abs/2407.21783. Uesato et al. (2022) Jonathan Uesato, Nate Kushman, Ramana Kumar, Francis Song, Noah Siegel, Lisa Wang, Antonia Creswell, Geoffrey Irving, and Irina Higgins. Solving math word problems with process- and outcome-based feedback, 2022. URL https://arxiv.org/abs/2211.14275. von Werra et al. (2020) Leandro von Werra, Younes Belkada, Lewis Tunstall, Edward Beeching, Tristan Thrush, Nathan Lambert, Shengyi Huang, Kashif Rasul, and Quentin Gallouédec. Trl: Transformer reinforcement learning. https://github.com/huggingface/trl, 2020. Wu et al. (2023) Zeqiu Wu, Yushi Hu, Weijia Shi, Nouha Dziri, Alane Suhr, Prithviraj Ammanabrolu, Noah A. Smith, Mari Ostendorf, and Hannaneh Hajishirzi. Fine-grained human feedback gives better rewards for language model training, 2023. URL https://arxiv.org/abs/2306.01693. Xie et al. (2024) Tianbao Xie, Siheng Zhao, Chen Henry Wu, Yitao Liu, Qian Luo, Victor Zhong, Yanchao Yang, and Tao Yu. Text2reward: Reward shaping with language models for reinforcement learning, 2024. URL https://arxiv.org/abs/2309.11489. (53) Adam X Yang, Maxime Robeyns, Thomas Coste, Zhengyan Shi, Jun Wang, Haitham Bou Ammar, and Laurence Aitchison. Bayesian reward models for llm alignment. In ICML 2024 Workshop on Structured Probabilistic Inference \\\ \&\ Generative Modeling. Yoon et al. (2024) Eunseop Yoon, Hee Suk Yoon, SooHwan Eom, Gunsoo Han, Daniel Nam, Daejin Jo, Kyoung-Woon On, Mark Hasegawa-Johnson, Sungwoong Kim, and Chang Yoo. Tlcr: Token-level continuous reward for fine-grained reinforcement learning from human feedback. In Findings of the Association for Computational Linguistics ACL 2024, p. 14969–14981, 2024. Zeng et al. (2024) Yongcheng Zeng, Guoqing Liu, Weiyu Ma, Ning Yang, Haifeng Zhang, and Jun Wang. Token-level direct preference optimization. In Proceedings of the 41st International Conference on Machine Learning, p. 58348–58365, 2024. Zhang et al. (2023) Yihua Zhang, Prashant Khanduri, Ioannis Tsaknakis, Yuguang Yao, Mingyi Hong, and Sijia Liu. An introduction to bi-level optimization: Foundations and applications in signal processing and machine learning, 2023. URL https://arxiv.org/abs/2308.00788. Zheng et al. (2023a) Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric P. Xing, Hao Zhang, Joseph E. Gonzalez, and Ion Stoica. Judging llm-as-a-judge with mt-bench and chatbot arena, 2023a. URL https://arxiv.org/abs/2306.05685. Zheng et al. (2023b) Rui Zheng, Shihan Dou, Songyang Gao, Yuan Hua, Wei Shen, Binghai Wang, Yan Liu, Senjie Jin, Qin Liu, Yuhao Zhou, Limao Xiong, Lu Chen, Zhiheng Xi, Nuo Xu, Wenbin Lai, Minghao Zhu, Cheng Chang, Zhangyue Yin, Rongxiang Weng, Wensen Cheng, Haoran Huang, Tianxiang Sun, Hang Yan, Tao Gui, Qi Zhang, Xipeng Qiu, and Xuanjing Huang. Secrets of rlhf in large language models part i: Ppo, 2023b. URL https://arxiv.org/abs/2307.04964. Zhong et al. (2025) Han Zhong, Zikang Shan, Guhao Feng, Wei Xiong, Xinle Cheng, Li Zhao, Di He, Jiang Bian, and Liwei Wang. Dpo meets ppo: Reinforced token optimization for rlhf, 2025. URL https://arxiv.org/abs/2404.18922. Appendix A Computing the Explainability Scores from a LLM Reward Model We discuss a simple setting of computing the SHAPley scores and also well-define its terms such as the simplified mapping z′, the mapping function hxsubscriptℎh_xhitalic_x, and show by example the satisfaction of the faithfulness property: g⁢(z′)≈f⁢(hx⁢(z′))superscript′subscriptℎsuperscript′g(z )≈ f(h_x(z ))g ( z′ ) ≈ f ( hitalic_x ( z′ ) ). Consider an input text x: ”I like apples” that receives a reward score r⁢(x)=2.12.1r(x)=2.1r ( x ) = 2.1. Then we consider 3 tokens • Token 1: I • Token 2: like • Token 3: apples The binary vector z′=[z1′,z2′,z3′]superscript′subscript1′subscript2′subscript3′z =[z_1 ,z_2 ,z_3 ]z′ = [ z1′ , z2′ , z3′ ] has each zi′∈0,1superscriptsubscript′01z_i ∈0,1zitalic_i′ ∈ 0 , 1 which masks each token as (1) present or (0) absent. The mapping function hxsubscriptℎh_xhitalic_x reconstructs valid text from the simplified input. For example, if z′=[1,0,1]superscript′101z =[1,0,1]z′ = [ 1 , 0 , 1 ] then hx⁢(z′)subscriptℎsuperscript′h_x(z )hitalic_x ( z′ ) might return ”I [MASK] apples” and r⁢(hx⁢(z′))subscriptℎsuperscript′r(h_x(z ))r ( hitalic_x ( z′ ) ) is the true reward for the masked text. We draw a table for each permutation and a toy example of the reward for each masked text z′ Masked Text hx⁢(z′)subscriptℎsuperscript′h_x(z )hitalic_x ( z′ ) r⁢(hx⁢(z′))subscriptℎsuperscript′r(h_x(z ))r ( hitalic_x ( z′ ) ) [0,0,0] [MASK] [MASK] [MASK] 0 [1,0,0] I [MASK] [MASK] 0.3 [0,1,0] [MASK] like [MASK] 0.5 [0,0,1] [MASK] [MASK] apples 1.2 [1,1,0] I like [MASK] 0.9 [1,0,1] I [MASK] apples 1.3 [0,1,1] [MASK] like apples 1.7 [1,1,1] I like apples 2.1 From this mapping, we can estimate token-level scores at intermediate states by computing the explainability metric or the feature attribution for each token to interpret each prediction of the reward model. Here, we consider the SHAPley kernel, which provides a natural way to identify the marginal contribution of each token i concerning the overall reward: ϕi⁢(f,x)=∑s⊆x\i|s|!⁢(|x|−|s|−1)!|x|!⁢[f⁢(s∪i)−fx⁢(s)]subscriptitalic-ϕsubscript\1delimited-[]subscript _i(f,x)= _s x \i\ |s|!(|x|-|s|-1)!|x|![% f(s∪\i\)-f_x(s)]ϕitalic_i ( f , x ) = ∑s ⊆ x \ i divide start_ARG | s | ! ( | x | - | s | - 1 ) ! end_ARG start_ARG | x | ! end_ARG [ f ( s ∪ i ) - fitalic_x ( s ) ] (6) where s is all subsets of tokens excluding token i that also respect the order of the original sentence, and x is the bag of all tokens in the sequence. Here, we have x=1,2,3123x=\1,2,3\x = 1 , 2 , 3 as the bag of all tokens from the sequence, and s constructs each subset of tokens but also maintains the order of the original sentence. Here, we can define the terms to properly solve Eq. 2 as a penalized linear regression: Ω⁢(g)Ω (g)Ω ( g ) =0absent0 =0= 0 πx⁢(z′)subscriptsuperscript′ _x(z )πitalic_x ( z′ ) =|z′|!⁢(|x|−|z′|−1)!|x|!=|s|!⁢(|x|−|s|−1)!|x|!absentsuperscript′11 = |z |!(|x|-|z |-1)!|x|!= |s|!(|x|-|% s|-1)!|x|!= divide start_ARG | z′ | ! ( | x | - | z′ | - 1 ) ! end_ARG start_ARG | x | ! end_ARG = divide start_ARG | s | ! ( | x | - | s | - 1 ) ! end_ARG start_ARG | x | ! end_ARG ℒ⁢(f,g,πx)ℒsubscript (f,g, _x)L ( f , g , πitalic_x ) =∑z′∈Z[f⁢(hx⁢(z′))−g⁢(z′)]2⁢πx⁢(z′)absentsubscriptsuperscript′delimited-[]subscriptℎsuperscript′2subscriptsuperscript′ = _z ∈ Z [f(h_x(z ))-g(z )% ]^2 _x(z )= ∑z′ ∈ Z [ f ( hitalic_x ( z′ ) ) - g ( z′ ) ]2 πitalic_x ( z′ ) Now, in relation to the simplified input z′, we consider a subset s with the token 22\2\ 2 which corresponds to z′=[0,1,0]superscript′010z =[0,1,0]z′ = [ 0 , 1 , 0 ]. Additionally s∪2=1,2212s∪\2\=\1,2\s ∪ 2 = 1 , 2 corresponds to z′=[1,1,0].superscript′110z =[1,1,0].z′ = [ 1 , 1 , 0 ] . Hence, using Eq. 3 we can compute the feature attribution for each token ϕisubscriptitalic-ϕ _iϕitalic_i. We continue with an example computation of the feature attribution for the token ”i”. Example 1. We compute the marginal contribution of element i=11i=1i = 1 across all subsets s⊆x∖11s x \1\s ⊆ x ∖ 1 , using the Shapley value formula. Case 1: s=∅s= = ∅ Weight: |s|!⁢(|x|−|s|−1)!|x|!=0!⋅(3−0−1)!3!=1⋅2!6=26=131⋅03013⋅1262613 |s|!(|x|-|s|-1)!|x|!= 0!·(3-0-1)!3!= 1% · 2!6= 26= 13divide start_ARG | s | ! ( | x | - | s | - 1 ) ! end_ARG start_ARG | x | ! end_ARG = divide start_ARG 0 ! ⋅ ( 3 - 0 - 1 ) ! end_ARG start_ARG 3 ! end_ARG = divide start_ARG 1 ⋅ 2 ! end_ARG start_ARG 6 end_ARG = divide start_ARG 2 end_ARG start_ARG 6 end_ARG = divide start_ARG 1 end_ARG start_ARG 3 end_ARG Marginal: f⁢(s∪1)−f⁢(s)=f⁢(1)−f⁢(∅)=0.3−0=0.3110.300.3 f(s∪\1\)-f(s)=f(\1\)-f( )=0.3-0=0.3f ( s ∪ 1 ) - f ( s ) = f ( 1 ) - f ( ∅ ) = 0.3 - 0 = 0.3 Contribution: 13⋅0.3=0.1⋅130.30.1 13· 0.3=0.1divide start_ARG 1 end_ARG start_ARG 3 end_ARG ⋅ 0.3 = 0.1 Case 2: s=2s2s=\2\s = 2 Weight: 1!⋅(3−1−1)!3!=1⋅1!6=16⋅13113⋅11616 1!·(3-1-1)!3!= 1· 1!6= 16divide start_ARG 1 ! ⋅ ( 3 - 1 - 1 ) ! end_ARG start_ARG 3 ! end_ARG = divide start_ARG 1 ⋅ 1 ! end_ARG start_ARG 6 end_ARG = divide start_ARG 1 end_ARG start_ARG 6 end_ARG Marginal: f⁢(1,2)−f⁢(2)=0.9−0.5=0.41220.90.50.4 f(\1,2\)-f(\2\)=0.9-0.5=0.4f ( 1 , 2 ) - f ( 2 ) = 0.9 - 0.5 = 0.4 Contribution: 16⋅0.4=0.066⁢6¯≈0.067⋅160.40.066¯60.067 16· 0.4=0.066 6≈ 0.067divide start_ARG 1 end_ARG start_ARG 6 end_ARG ⋅ 0.4 = 0.066 over¯ start_ARG 6 end_ARG ≈ 0.067 Case 3: s=3s3s=\3\s = 3 Weight: 1!⋅(3−1−1)!3!=1⋅1!6=16⋅13113⋅11616 1!·(3-1-1)!3!= 1· 1!6= 16divide start_ARG 1 ! ⋅ ( 3 - 1 - 1 ) ! end_ARG start_ARG 3 ! end_ARG = divide start_ARG 1 ⋅ 1 ! end_ARG start_ARG 6 end_ARG = divide start_ARG 1 end_ARG start_ARG 6 end_ARG Marginal: f⁢(1,3)−f⁢(3)=0.6−0.2=0.41330.60.20.4 f(\1,3\)-f(\3\)=0.6-0.2=0.4f ( 1 , 3 ) - f ( 3 ) = 0.6 - 0.2 = 0.4 Contribution: 16⋅0.4=0.067⋅160.40.067 16· 0.4=0.067divide start_ARG 1 end_ARG start_ARG 6 end_ARG ⋅ 0.4 = 0.067 Case 4: s=2,3s23s=\2,3\s = 2 , 3 Weight: 2!⋅(3−2−1)!3!=2⋅0!6=26=13⋅23213⋅2062613 2!·(3-2-1)!3!= 2· 0!6= 26= % 13divide start_ARG 2 ! ⋅ ( 3 - 2 - 1 ) ! end_ARG start_ARG 3 ! end_ARG = divide start_ARG 2 ⋅ 0 ! end_ARG start_ARG 6 end_ARG = divide start_ARG 2 end_ARG start_ARG 6 end_ARG = divide start_ARG 1 end_ARG start_ARG 3 end_ARG Marginal: f⁢(1,2,3)−f⁢(2,3)=1.0−0.7=0.3123231.00.70.3 f(\1,2,3\)-f(\2,3\)=1.0-0.7=0.3f ( 1 , 2 , 3 ) - f ( 2 , 3 ) = 1.0 - 0.7 = 0.3 Contribution: 13⋅0.3=0.1⋅130.30.1 13· 0.3=0.1divide start_ARG 1 end_ARG start_ARG 3 end_ARG ⋅ 0.3 = 0.1 Continuing this sequence gives eventually the rest of the terms in which summing them computes ϕ1subscriptitalic-ϕ1 _1ϕ1 as: ϕ1=0.1+0.067+0.017+0.133≈0.32subscriptitalic-ϕ10.10.0670.0170.1330.32 _1=0.1+0.067+0.017+0.133≈ 0.32ϕ1 = 0.1 + 0.067 + 0.017 + 0.133 ≈ 0.32 Continuing this with the rest of the SHAPley values gives eventually ϕ1=0.32,ϕ2=0.62,ϕ3=1.17formulae-sequencesubscriptitalic-ϕ10.32formulae-sequencesubscriptitalic-ϕ20.62subscriptitalic-ϕ31.17 _1=0.32, _2=0.62, _3=1.17ϕ1 = 0.32 , ϕ2 = 0.62 , ϕ3 = 1.17 Hence, we see that the summation of all feature attributions give approximately the full reward as well as satisfying the additive surrogate g⁢(z′)=r⁢(∅)+∑i=13ϕi⁢z′≈r⁢(hx⁢(z′))superscript′subscript13subscriptitalic-ϕsuperscript′subscriptℎsuperscript′g(z )=r( )+ _i=1^3 _iz ≈ r(h_x(z% ))g ( z′ ) = r ( ∅ ) + ∑i = 13 ϕitalic_i z′ ≈ r ( hitalic_x ( z′ ) ) In this example, it is important to note that we overview the standard SHAPley formula, which requires computing every partition of tokens, resulting in 2ksuperscript22^k2k operations. However, in practice, we employ the Owen values, which provably produce the same solution as SHAPley values (López & Saboyá, 2009) and also reduce the necessary operations drastically, requiring to compute only x2superscript2x^2x2 (where x is the total sequence length) operations per prediction. Appendix B Proof of Policy Invariance of Explainability Methods Proposition 1. Given a policy πθsubscript _θπitalic_θ and a Markov Decision Process ℳ=(,,,γ,R)ℳM=(S,A,P,γ,R)M = ( S , A , P , γ , R ), any reward shaping function R′ in the family of additive feature attribution methods follows a potential-based shaping function and then the optimal πθsubscript _θπitalic_θ for R′ is also optimal for the original R. Proof. We first have R′⁢(s,a,s′)=R⁢(s,a,s′)+F⁢(s,a,s′)superscript′superscript′R (s,a,s )=R(s,a,s )+F(s,a,s )R′ ( s , a , s′ ) = R ( s , a , s′ ) + F ( s , a , s′ ) where F is a potential-based shaping function defined as F⁢(s,a,s′)=γ⁢Φ⁢(s′)−Φ⁢(s)superscript′Φsuperscript′ΦF(s,a,s )=γ (s )- (s)F ( s , a , s′ ) = γ Φ ( s′ ) - Φ ( s ) for Φ:S↦ℝ:Φmaps-toℝ :S Φ : S ↦ blackboard_R. We also have by locality from Eq. 1 that g⁢(x′)≈Rϕ⁢(s,a,s′)superscript′subscriptitalic-ϕsuperscript′g(x )≈ R_φ(s,a,s )g ( x′ ) ≈ Ritalic_ϕ ( s , a , s′ ) or: ∑i=1s′ϕi⁢(f,x)≈R⁢(s,a,s′)superscriptsubscript1superscript′subscriptitalic-ϕsuperscript′ _i=1^s _i(f,x)≈ R(s,a,s )∑i = 1s start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT ϕitalic_i ( f , x ) ≈ R ( s , a , s′ ) (7) where s′ refers to the index up to state s′ in the sequence. Then, without loss of generality, consider the case with heuristics =ϕs′,1subscriptitalic-ϕsuperscript′1 E=\ _s ,1\caligraphic_E = ϕitalic_s′ , 1 where the constant is just the identity for the original sparse reward. We omit the case for the constant, as it is known that any linear transformation also preserves the optimal policy. Then, by substituting Eq. 4 for F with γ=11γ=1γ = 1 since we care about all future states equally, we have: Φ⁢(s′)−Φ⁢(s)=w2⁢ϕs′⁢(R,x)Φsuperscript′Φsubscript2subscriptitalic-ϕsuperscript′ (s )- (s)=w_2 _s (R,x)Φ ( s′ ) - Φ ( s ) = w2 ϕitalic_s′ ( R , x ) (8) Then defining Φ⁢(s)=w2⁢∑i=1sϕi⁢(R,x)Φsubscript2superscriptsubscript1subscriptitalic-ϕ (s)=w_2Σ _i=1^s _i(R,x)Φ ( s ) = w2 ∑i = 1s ϕitalic_i ( R , x ) we have: Φ⁢(s′)−Φ⁢(s)Φsuperscript′Φ (s )- (s)Φ ( s′ ) - Φ ( s ) =w2⁢∑i=1s′ϕi⁢(R,x)−w2⁢∑i=1sϕi⁢(R,x)=w2⁢ϕs′⁢(R,x)absentsubscript2superscriptsubscript1superscript′subscriptitalic-ϕsubscript2superscriptsubscript1subscriptitalic-ϕsubscript2subscriptitalic-ϕsuperscript′ =w_2 _i=1^s _i(R,x)-w_2 _i=1^s% _i(R,x)=w_2 _s (R,x)= w2 ∑i = 1s start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT ϕitalic_i ( R , x ) - w2 ∑i = 1s ϕitalic_i ( R , x ) = w2 ϕitalic_s′ ( R , x ) (9) Since by Eq. 8 we have that ∑i=1s′ϕi⁢(R,x)−∑i=1sϕi⁢(R,x)=ϕs′⁢(R,x)superscriptsubscript1superscript′subscriptitalic-ϕsuperscriptsubscript1subscriptitalic-ϕsubscriptitalic-ϕsuperscript′ _i=1^s _i(R,x)- _i=1^s _i(R,x)= _s^% (R,x)∑i = 1s start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT ϕitalic_i ( R , x ) - ∑i = 1s ϕitalic_i ( R , x ) = ϕitalic_s′ ( R , x ) ∎ Appendix C Computation Overhead of Explanation Methods To measure the computational overhead of explanation methods, we can represent the time complexity required for inference of the reward model as ⁢(N×F⁢(k))O(N× F(k))O ( N × F ( k ) ) where F⁢(k)F(k)F ( k ) represents the number of floating point operations required in one forward pass on sequence of length k, then our method is upper-bounded by ⁢(N×k2×F⁢(k))superscript2O(N× k^2× F(k))O ( N × k2 × F ( k ) ) computations when also calculating SHAPley values. We run a simple test measuring the total wall clock time, varying batch size and generation length to run inference with and without explainability methods: Appendix D Practical Implementation Details Algorithm 1 Explainable Reward Shaping with Bayesian Optimization 1:Initialize a Gaussian Process (GP) surrogate model: B⁢O←GP←GPBO O ← GP 2:for k=1,…,n1…k=1,…,nk = 1 , … , n do ▷ ▷ Bayesian Optimization iterations 3: wk←B⁢O.next_action()formulae-sequence←subscriptnext_action()w_k← BO.next\_action()witalic_k ← B O . next_action() 4: for xi∈1:nsubscriptsubscript:1x_i _1:nxitalic_i ∈ D1 : n do ▷ ▷ Subsampled set of points DD and train PPO 5: yi←πθ⁢(xi)←subscriptsubscriptsubscripty_i← _θ(x_i)yitalic_i ← πitalic_θ ( xitalic_i ) 6: r⁢(xi,yi);←ℰ⁢(r)←subscriptsubscriptℰr(x_i,y_i);\; E (r)r ( xitalic_i , yitalic_i ) ; caligraphic_E ← E ( r ) ▷ ▷ Compute sparse and token-level scores (Eq. 2) 7: r′←wk⊤⁢←superscript′subscripttopr ← w_k Er′ ← witalic_k⊤ caligraphic_E 8: Update θ via ∇θJ⁢(πθ;r′)subscript∇subscriptsuperscript′ _θJ( _θ;r )∇θ J ( πitalic_θ ; r′ ) 9: end for 10: rv⁢a⁢l←1m⁢∑i=1mr⁢(xi,πθ⁢(xi))←superscript1superscriptsubscript1subscriptsubscriptsubscriptr^val← 1m _i=1^mr (x_i, _θ(x_i)% )ritalic_v a l ← divide start_ARG 1 end_ARG start_ARG m end_ARG ∑i = 1m r ( xitalic_i , πitalic_θ ( xitalic_i ) ) ▷ ▷ Compute avg. validation reward 11: B⁢O.update⁢(wk,rv⁢a⁢l)formulae-sequenceupdatesubscriptsuperscriptBO.update(w_k,r^val)B O . update ( witalic_k , ritalic_v a l ) 12: if rv⁢a⁢l>rb⁢e⁢s⁢tsuperscriptsuperscriptr^val>r^bestritalic_v a l > ritalic_b e s t then 13: Update the PPO checkpoint for next BO iteration 14: end if 15:end for 16:wbest←B⁢O.best_params()formulae-sequence←subscriptbestbest_params()w_best← BO.best\_params()wbest ← B O . best_params() ▷ ▷ Retrieve the best parameters 17:Train new PPO using r′←wbest⊤⁢←superscript′subscriptbesttopr ← w_best Er′ ← wbest⊤ caligraphic_E Component Link Dataset Anthropic/h-rlhf Reward Model Base weqweasdas/h_rlhf_rm_open_llama_3b SFT Model Base VMware/open-llama-7b-open-instruct Component Link Dataset princeton-nlp/llama3-ultrafeedback Reward Model Base meta-llama/Llama-3.2-1B-Instruct SFT Model Base meta-llama/Llama-3.2-3B-Instruct Appendix E Qualitative Examples Prompt: Can you describe the ways in which insomnia is treated? Model Generated Text from H-RLHF Test split Score SHAP (no BO) The main treatments for insomnia are cognitive behavioral therapy (CBT), pharmaceutical drugs, and sleep hygiene practices. Cognitive behavioral therapy involves techniques such as sleep hygiene education, relaxation training, and cognitive restructuring, and is shown to be significantly more effective than medication or simple advice for short-term insomnia. Medication is generally used as a second-line treatment for chronic insomnia, and can include hypnotics such as benzodiazepines or non-benzodiazepine sedatives and inhibitors of the wake-promoting hormones such as histamine and melatonin. Sleep hygiene practices such as avoiding 4.63 SHAP (d=2)2(d=2)( d = 2 ) ”Insomnia is typically treated with one of the three following methods: 1. Medication 2. Cognitive behavioral therapy (CBT) 3. Sleep hygiene practices, such as regular sleep schedules and avoiding heavy meals before bedtime. Medications like benzodiazepines and selective serotonin reuptake inhibitors (SSRIs) can help with sleep problems. Cognitive behavioral therapy can be helpful in changing maladaptive thoughts and behaviors around sleep, such as staying up too late, which prolongs the time spent awake. It can also help teach sleep-promoting” 6.5 SHAP + ATTN (d=3)3(d=3)( d = 3 ) Insomnia is treated based on the underlying cause. The most common treatments include medications such as Benzodiazepines (e.g. triazolam, temazepam, alprazolam), Z-drugs (e.g. doxepin, zolpidem), or Minor tranquillizers (e.g. buspirone). Cognitive behavioral therapy (CBT) and relaxation techniques are also common treatments. </s> 8.63 SHAP + LIME (d=3)3(d=3)( d = 3 ) Treatment for insomnia includes changes in lifestyle and environmental factors, as well as medications and other therapies. Common treatments include cognitive-behavioural therapy, which focuses on identifying and changing thought patterns and other behavioural factors that may contribute to insomnia, and medication that may help regulate the sleep cycle. Other treatments may include relaxation techniques, such as meditation or breathing exercises, which can help calm the mind and promote sleep. </s> 9.31 Table 2: Generation topography across different models as input dimension increases for BO.

Instant research discovery