{"id":"8153ae72-552c-43d4-b0ca-559085ae8836","arxiv_id":"2504.16272","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":5.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":2,"one_line_summary":"Reward shaping based on SHAP/LIME token attributions, with weights optimized by Bayesian optimization, improves RLHF training speed and downstream win rates while preserving the optimal policy.","lead":"What if an AI model could get feedback on every word it writes, not just a final score? This paper tests giving language models dense, token-level rewards computed from explanation tools like SHAP, with weights tuned by Bayesian optimization, and reports faster and better RLHF training.","discovery_kind":"extension","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Policy-invariance proof in Appendix B does not connect to the multiplicative/softmax reward shape in Eq. 4; the guarantee is trivially true only if softmax normalization is applied, and is unproven as stated.","rationale":"The central claim has two parts: a theoretical policy-invariance guarantee and an empirical improvement from BO-shaped dense rewards. The theoretical guarantee is the more load-bearing because the empirical comparison is noisy (no seeds, tiny training budgets, and mixed results in Table 1). The reader's weakest_assumption correctly identifies that Appendix B's additive potential-based proof does not match Eq. 4's multiplicative/softmax reward. Examining the appendix confirms the mismatch: Eq. 8 is tautological given the chosen Phi, and Eq. 7 does not follow from additive attribution local accuracy. The proof therefore does not establish the claimed invariance. The concern is not fatal because a simple total-return argument (softmax probabilities sum to one) recovers policy invariance for the undiscounted finite-horizon objective actually used in RLHF. However, that argument changes the paper's stated mechanism: potential-based shaping is not what provides the guarantee, and the BO weights affect optimization dynamics rather than the asymptotic objective. I would keep the reader's conditional verdict: the proof needs correction or replacement, and the empirical claims would benefit from multi-seed runs with error bars, but the core idea remains viable and worth further development.","tokens_in":17272,"tokens_out":10959,"duration_ms":110182,"concrete_test":"Implement the exact reward from Section 3.1 and Algorithm 1 on a small controllable MDP (e.g., a 3-token generation task with a fixed reward model), enumerate all trajectories, and compare the sum of shaped per-token rewards to the original terminal reward. If equality holds for every trajectory, the invariance claim is established by the trivial sum argument and Appendix B should be replaced accordingly. If equality fails for some BO-sampled weights, the policy-invariance claim is false as implemented. Separately, symbolically test whether the softmax-shaped per-token reward r'_t = softmax(w1*phi + w2)_t * R_total can be written as gamma*Phi(s_{t+1}) - Phi(s_t) for any state potential; if not, the potential-based proof in Appendix B cannot be valid.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The central theoretical claim is that additive feature attribution methods preserve the optimal policy via potential-based shaping. Section 3.1 defines the shaped reward multiplicatively as r'(s,a) = W^T E * r(s,a), and in practice softmax-normalizes per-token scores and broadcasts the sparse reward over the sequence. Appendix B instead proves invariance for the additive potential-based form R' = R + F, with F = gamma*Phi(s') - Phi(s), and sets Phi(s) = w2 * sum_{i=1}^s phi_i so that F becomes w2 * phi_{s'}. No step shows that the implemented multiplicative/softmax reward equals R + F for any state potential. The appendix's Eq. 7 also assumes cumulative attribution up to state s' approximates the prefix reward, which is not a consequence of SHAP/LIME local accuracy: local accuracy only fixes the sum over all tokens. Thus the stated proof does not support the invariance claim. A correct but much weaker argument exists: under softmax normalization the per-token rewards sum to the original scalar reward for every trajectory, so the undiscounted finite-horizon total return is identical and the optimal policy is preserved. That argument does not rely on potential-based shaping and does not justify Eq. 4 as written. If softmax normalization is omitted from Eq. 4, the total shaped return becomes w1 * sum(phi_i) * R + w2 * M * R, which is not generally invariant.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes a reward-shaping method for RLHF in which per-token feature attributions (SHAP, LIME, and attention) from the reward model are combined into a dense token-level reward, with combination weights selected by Bayesian optimization in an outer loop while PPO trains in the inner loop. The authors claim that additive feature attribution functions preserve the optimal policy via potential-based reward shaping, and report improved reward scores and win rates on HH-RLHF and Ultrafeedback, including evaluations on AlpacaEval-2 and MTBench.","tokens_in":17602,"tokens_out":4953,"duration_ms":48088,"significance":"If the theoretical claim were established, the paper would provide a principled way to densify RLHF rewards without changing the intended optimal policy, and the combination of explainability with Bayesian optimization is a reasonable research direction. The release of code and the use of independent benchmarks such as AlpacaEval-2 and MTBench are strengths. However, the policy-invariance proof as stated does not apply to the implemented reward transformation, and the empirical evaluation rests on single runs with a very small training budget, so the central claims are not yet supported.","major_comments":[{"comment":"The central policy-invariance proof does not apply to the reward actually used. Eq. (4) defines the shaped reward multiplicatively as r'(s,a) = (w1 E_SHAP(s,a) + w2) r(s,a), and the implementation described in Section 3.1 softmax-normalizes the per-token scores before broadcasting the scalar reward. Appendix B instead proves invariance for an additive potential-based shaping R' = R + F with F = gamma Phi(s') - Phi(s), identifying Phi(s) = w2 sum_{i=1}^s phi_i. No step connects the multiplicative/softmax form to R + F. Eq. (7) also assumes that cumulative attribution up to state s' approximates the prefix reward, which is not implied by SHAP local accuracy; local accuracy only fixes the sum over all tokens. The correct weaker argument--that softmax normalization makes the per-token rewards sum to the original scalar reward for every trajectory, so the finite-horizon total return is unchanged--is not the argument given and does not justify Eq. (4) as written. Without normalization, the total shaped return is (w1 sum_i phi_i + w2 M) R, which is not generally invariant. This is a load-bearing gap because policy invariance is the paper's main theoretical contribution.","section":"Eq. (4); Appendix B"},{"comment":"The empirical evaluation does not support the strength of the empirical claims. Table 1 reports a single run per method with no multiple seeds or error bars, and the BO training budget is only 25 trials with 80 PPO samples per trial. Under this noise level, differences such as BO-SHAP-LIME (6.58) versus BO-SHAP-Attn (6.47) versus LIME* (5.86) on HH-RLHF cannot be distinguished from random variation. The paper itself acknowledges that the d=4 combination degrades performance (Section 4.2), which undercuts the claim that BO 'balances' token-level rewards. Rerunning with several seeds and reporting means with standard deviations or confidence intervals for the headline numbers is necessary before the empirical conclusions can be assessed.","section":"Section 4, Table 1"},{"comment":"The outer BO loop is optimized on average validation reward from the same reward model whose test-split predictions form the 'Score' column of Table 1. This creates selection pressure toward reward-model overfitting, and the independent benchmarks are the only protection against circularity. The paper should report the gap between validation and test reward for the BO-selected weights and should consider using a separate reward model or a preference-based metric for the outer objective. As written, the 'Score' improvements are not a clean measure of alignment quality.","section":"Sections 3.2, 4.1, 4.2"}],"minor_comments":[{"comment":"The numerical values used in the example are inconsistent with the table: f({1,3}) is given as 0.6 but the table lists 1.3, and f({2,3}) is given as 0.7 but the table lists 1.7. This makes the example's arithmetic unreliable.","section":"Appendix A, Example 1"},{"comment":"Please define the dimensions of E and W and clarify whether r(s,a) is a scalar or a per-token vector; the notation W^T E · r is ambiguous as written.","section":"Section 3.1, Eq. (4)"},{"comment":"The paper states that sampled weights satisfy sum_i w_i = 1, but the manual baseline sets w = 0.8 for a single dense reward type; please clarify whether the manual baseline also includes a constant term that sums to one.","section":"Section 3.1, Section 4.1"},{"comment":"The caption says the shading represents the standard error (95% confidence interval), but the paper does not report the number of runs or seeds used to compute this quantity; please clarify.","section":"Figure 4"},{"comment":"There are several typographical issues: 'backpropogate' should be 'backpropagate' in Section 4.2, 'LL AMA' appears in Section 4, and 'Also inspects' in the Related Work paragraph should have consistent capitalization.","section":"Minor wording"}],"recommendation":"major_revision","confidential_remarks":"The theoretical result is the main reason to consider this paper, but the proof does not match the implemented shaping function. If the authors cannot reconcile Eq. (4) with a potential-based form, they should either correct the theorem to the softmax-return invariance argument or substantially weaken the claim. The empirical section needs multiple seeds and error bars before the reported improvements can be taken seriously. I do not think the current version meets the bar for acceptance, but the idea is promising enough to warrant a major revision."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Here's the short version: the combination of SHAP/LIME token attributions with BO-learned weights for RLHF reward shaping is new and sensible, and the authors are honest that the full four-signal configuration does not help. The central theoretical assertion, however, is not supported as written. Eq. 4 defines the shaped reward multiplicatively: r' = (w1 * E_SHAP + w2) * r, with softmax normalization in practice. Appendix B proves invariance for an additive potential-based form R' = R + F, with F = gamma*Phi(s') - Phi(s). There is no step linking the two. The proof's Eq. 7 also assumes a cumulative attribution property that does not follow from SHAP/LIME local accuracy. The weaker invariance argument via softmax normalization exists — per-token rewards sum to the original scalar reward, so the finite-horizon total return is identical — but the paper does not make it. This is load-bearing: the abstract promises a theoretical guarantee that the proof does not deliver.\n\nWhat is genuinely good: the idea is novel relative to Chan et al. (2024), the code is public, and the empirical section at least tries open benchmarks in addition to reward-model score. The BO framing is a reasonable way to handle noisy attribution estimates, and the paper calls out the failure of the d=4 case rather than hiding it.\n\nThe soft spots beyond the proof: no multiple seeds, no error bars on the headline table, and the training budget is tiny (25 BO trials, 80 PPO samples each). Some BO configurations underperform the single-source baselines, which weakens the \"BO helps\" message. The weights are selected on the validation split of the same reward model used for test reporting, so the small wins could be part of selection pressure.\n\nWho is this for? Someone working on token-level credit assignment in RLHF, particularly if they want a quick overview of why explainability-based dense rewards might help. A serious referee should get a chance at this, but the paper is not ready as is. My recommendation: send to peer review, but insist on a corrected proof (or a straight statement of the trivial softmax invariance), multi-seed runs, and error bars.","headline":"A plausible new dense-reward recipe for RLHF, but the policy-invariance proof doesn't match the implemented reward and the experiments are too noisy to carry the claim.","tokens_in":18135,"tokens_out":2811,"would_cite":false,"duration_ms":25516,"reading_group":"maybe","serious_thinker":"no","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"This paper claims that per-token reward shaping from SHAP/LIME explanations, with weights tuned by Bayesian optimization, preserves the original RLHF objective while speeding up training and improving downstream benchmark performance.","keywords":["RLHF","reward shaping","token-level credit assignment","additive feature attribution","SHAP","LIME","Bayesian optimization","policy invariance"],"falsifier":"Run the shaped reward of Eq. 4 on a small MDP whose optimal policy is known, solve both the shaped and unshaped MDPs exhaustively, and check whether the optimal policies coincide for all sampled weight vectors; a single weight vector where they differ falsifies policy invariance.","tokens_in":17057,"feed_emoji":"🎯","tokens_out":9176,"duration_ms":79156,"temperature":0.7,"pith_summary":"This paper tries to establish that the sparse, sequence-level reward used in reinforcement learning from human feedback can be replaced by a dense, token-level reward without changing what the policy optimizes. The proposed shaping takes per-token credit scores from additive feature attribution methods such as SHAP and LIME, combines them with learned weights, and broadcasts the scalar reward across the sequence. The weights are selected in an outer Bayesian-optimization loop while PPO trains the policy in the inner loop. The authors argue that any shaping built from an additive feature attribution function is a potential-based reward shaping, so the optimal policy of the original reward is preserved, and report that the BO-tuned shapes improve training speed, value-function stability, and downstream benchmark performance. If right, this gives RLHF a principled way to give language models fine-grained per-token feedback instead of a single final-token surrogate.","feed_headline":"SHAP/LIME token rewards speed RLHF and preserve the optimal policy","feed_subtitle":"A Bayesian loop tunes per-token reward weights, so LLMs get finer feedback than a single score at the end of a reply.","key_machinery":"The load-bearing object is the additive feature attribution model, $g(z') = \\phi_0 + \\sum_{i=1}^{M}\\phi_i z'_i$, a locally fitted linear model whose coefficients $\\phi_i$ are per-token credit scores for the reward model's prediction. The paper combines these scores with a constant into $E=[E_{\\mathrm{SHAP}},\\mathbf{1}]$ and shapes the reward as $r'(s,a)=W^\\top E \\cdot r(s,a)$, with $W$ drawn from a Bayesian-optimization loop (a Gaussian-process surrogate with a log noisy expected improvement acquisition function) that maximizes the validation reward of the PPO-trained policy. The proof machinery is potential-based reward shaping: defining $\\Phi(s)=w_2\\sum_{i=1}^{s}\\phi_i$, the added term $w_2\\phi_{s'}$ is claimed to equal $\\Phi(s')-\\Phi(s)$, which is the condition under which the optimal policy is unchanged.","core_discovery":"On its own terms, the paper's claim is that reward shaping via additive feature attribution is both safe and useful for RLHF. The shaped reward is written as $r'(s,a) = W^\\top E \\cdot r(s,a)$ with $E = [E_{\\mathrm{SHAP}}, \\mathbf{1}]$, so the scalar reward $r(s,a)$ is multiplied by a convex combination of token-level explanation scores and the constant 1; in practice the explanation vector is softmax-normalized and the scalar reward is broadcast over the sequence. The paper proves (Appendix B) that this family is potential-based by setting $\\Phi(s) = w_2\\sum_{i=1}^{s}\\phi_i(R,x)$, so that the shaping term equals $\\Phi(s')-\\Phi(s)$, and invokes the classic result that potential-based shaping preserves the optimal policy. Empirically, the authors show that unblended SHAP/LIME rewards stabilize PPO value loss, and that BO-tuned combinations such as SHAP+LIME achieve higher validation reward and better AlpacaEval-2 and MT-Bench win rates than sparse RLHF and attention-based dense reward baselines.","pith_inferences":["A natural next test is to compare BO-selected weights with an exhaustive grid over $w_1,w_2$ on a short-sequence task; if a grid matches BO, the gains come from tuning the weights rather than from the acquisition function.","The same additive-attribution recipe could be applied to process-supervised reasoning tasks: per-token SHAP scores of a verifier would give step-level credit for chain-of-thought, and the BO outer loop could tune the balance between step-level and final rewards.","If the invariance claim is repaired to cover the multiplicative form, the framework generalizes to any token-importance signal, such as saliency maps, gradient norms, or learned critics, provided the shaping can be written as a potential difference.","The static weight vector could become prompt-conditional, turning the outer loop into a meta-learning problem, a direction the authors note in the conclusion."],"forward_implications":["Practitioners can add per-token feedback from any additive feature attribution method without re-deriving the task objective, because the shaping is claimed to preserve the original optimal policy.","Combining several token-importance signals (SHAP, LIME, attention) and tuning their convex weights via Bayesian optimization can outperform any single dense-reward signal within a modest trial budget.","Dense explanation rewards make PPO value learning more stable, shown by lower value-head loss, which can reduce the variance of RLHF training runs.","Because the BO outer loop is a hyperparameter search over reward shapes, the final policy is only as good as the validation reward used to select weights; adding more signal sources requires more trials before performance is guaranteed."],"supporting_citations":[{"why":"Supplies the potential-based reward shaping theorem that the paper invokes to claim policy invariance for additive attribution shapings.","marker":"Ng et al. (1999)"},{"why":"Defines LIME and the additive local surrogate model family used to compute per-token explanation scores.","marker":"Ribeiro et al. (2016)"},{"why":"Provides SHAP values as the additive feature attribution estimator behind E_SHAP in the shaped reward.","marker":"Lundberg & Lee (2017)"},{"why":"Sets the attention-based dense reward baseline and the procedure of broadcasting the scalar reward over token-level probabilities.","marker":"Chan et al. (2024)"},{"why":"Supplies the log Noisy Expected Improvement acquisition function that the Bayesian optimizer uses to handle noisy reward evaluations.","marker":"Ament et al. (2025)"},{"why":"Provides the Ax implementation of the Gaussian-process surrogate and acquisition loop used for weight selection.","marker":"Bakshy et al. (2018)"},{"why":"Supplies the HH-RLHF helpfulness dataset used for the main training and evaluation runs.","marker":"Bai et al. (2022)"},{"why":"Supplies the Ultrafeedback dataset used for the second set of RLHF experiments.","marker":"Cui et al. (2024)"},{"why":"Defines the length-controlled AlpacaEval-2 win-rate metric used to compare shaped-reward policies on open benchmarks.","marker":"Dubois et al. (2025)"},{"why":"Defines PPO, the policy-gradient algorithm used as the lower-level optimizer inside the bilevel loop.","marker":"Schulman et al. (2017)"}],"fun_headline_variants":["Bayesian-tuned SHAP/LIME rewards accelerate RLHF alignment","Optimal policy preserved with BO-tuned SHAP/LIME token rewards","SHAP/LIME dense rewards: faster RLHF, policy preservation proven","Token-level credit via SHAP/LIME, BO-tuned, speeds RLHF"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The policy-invariance proof applies to a reward built by adding a potential-based bonus, but the method's actual reward is built by multiplying the sparse reward by a weighted blend of token scores; if those two constructions are not equivalent, the guarantee does not cover the trained method.","fun_headline_variants_meta":{"raw":{"variants":["Bayesian-tuned SHAP/LIME rewards accelerate RLHF alignment","Optimal policy preserved with BO-tuned SHAP/LIME token rewards","SHAP/LIME dense rewards: faster RLHF, policy preservation proven","Token-level credit via SHAP/LIME, BO-tuned, speeds RLHF"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000403,"raw_usage":{"total_tokens":2103,"prompt_tokens":951,"completion_tokens":1152,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":567,"completion_tokens_details":{"reasoning_tokens":1072}},"tokens_in":567,"tokens_out":1152,"duration_ms":10130,"temperature":1.0,"reasoning_tokens":1072,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-16T11:08:11.404944+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Run the shaped reward of Eq. 4 on a small MDP whose optimal policy is known, solve both the shaped and unshaped MDPs exhaustively, and check whether the optimal policies coincide for all sampled weight vectors; a single weight vector where they differ falsifies policy invariance.","supporting_citations":[{"cited_title":"A unified approach to interpreting model predictions","cited_arxiv_id":null,"evidence_quote":"Provides SHAP values as the additive feature attribution estimator behind E_SHAP in the shaped reward."},{"cited_title":"Ae: A domain-agnostic platform for adaptive experimentation","cited_arxiv_id":null,"evidence_quote":"Provides the Ax implementation of the Gaussian-process surrogate and acquisition loop used for weight selection."}],"review_version":1}