Source-linked AI summary
Transforming and Combining Rewards for Aligning Large Language Models
Zihao Wang, Chirag Nagpal, Jonathan Berant, Jacob Eisenstein, Alex D'Amour, Sanmi Koyejo, Victor Veitch
TL;DR
The paper asks how to choose among ranking-preserving reward transformations and how to combine rewards for multiple properties. It derives a centered log-sigmoid transformation for Bradley-Terry rewards and reports substantial alignment improvements over raw rewards. The transformation also gives summed rewards a probabilistic interpretation as requiring outputs to be good across properties.
Problem
The paper asks whether a reward transformation can improve alignment and how to combine reward models when outputs should satisfy multiple properties.
Source: §1 · pp. 1–2Method
For Bradley-Terry rewards, the paper derives LSC-transformation: applying log-sigmoid to rewards centered by a prompt-specific reference.
Source: §1 · p. 2Results
The transformed reward yields substantial alignment improvements over the raw-reward baseline across KL levels in helpfulness-and-harmlessness alignment.
Source: §1 · p. 2Takeaways & Limitations
Summing transformed rewards corresponds to the logical AND of outputs being good on each measured property.
Source: §1 · p. 2In general, individual rewards reveal only marginal property distributions, while logical AND may depend on interactions, so aggregation requires an extra assumption.
Source: §4 · p. 6
Abstract
from arXiv · showhide
A common approach for aligning language models to human preferences is to first learn a reward model from preference data, and then use this reward model to update the language model. We study two closely related problems that arise in this approach. First, any monotone transformation of the reward model preserves preference ranking; is there a choice that is ``better'' than others? Second, we often wish to align language models to multiple properties: how should we combine multiple reward models? Using a probabilistic interpretation of the alignment procedure, we identify a natural choice for transformation for (the common case of) rewards learned from Bradley-Terry preference models. The derived transformation is straightforward: we apply a log-sigmoid function to the centered rewards, a method we term ``LSC-transformation'' (log-sigmoid-centered transformation). This transformation has two important properties. First, it emphasizes improving poorly-performing outputs, rather than outputs that already score well. This mitigates both underfitting (where some prompts are not improved) and reward hacking (where the model learns to exploit misspecification of the reward model). Second, it enables principled aggregation of rewards by linking summation to logical conjunction: the sum of transformed rewards corresponds to the probability that the output is ``good'' in all measured properties, in a sense we make precise. Experiments aligning language models to be both helpful and harmless using RLHF show substantial improvements over the baseline (non-transformed) approach.
1 Introduction
The paper asks whether reward transformations can improve alignment and how multiple reward models should be combined. It proposes a probabilistic interpretation that motivates a centered log-sigmoid transformation and reports substantial improvements when aligning for helpfulness and harmlessness.
1 Introduction
- The paper studies whether transforming learned rewards can improve alignment and how reward models for multiple desired properties should be combined.The motivating properties include helpfulness, harmlessness, factuality, and creativity.Source: §1 · pp. 1–2
- A probabilistic interpretation treats alignment to one property as conditioning on outputs being good, and alignment to multiple properties as conditioning on all being good.For Bradley-Terry rewards under KL-constrained expected-reward maximization, the paper identifies a natural transformation.Source: §1 · p. 2
- The LSC-transformation applies log-sigmoid to the reward centered by a prompt-specific reference value.The paper calls alignment using this transformation LSC-alignment.Source: §1 · p. 2
- Summing transformed rewards corresponds to logical AND across properties, while diminishing marginal utility can encourage improvement on weak prompts and discourage reward hacking.The paper reports substantial improvements over untransformed rewards for helpfulness-and-harmlessness alignment across KL-regularization levels.Source: §1 · p. 2
2 Preliminaries
The preliminaries describe learning a reward model from human preference data and using it to update a supervised-finetuned language model. Standard alignment maximizes expected reward while regularizing divergence from that starting model.
2 Preliminaries
- Preference data commonly consists of a prompt and two responses, with a human indicating the preferred response; Bradley-Terry rewards are a common model for such data.The reward function is fit using a maximum log-likelihood objective.Source: §2 · p. 2
- The base language model is commonly supervised-finetuned on winning preference examples, producing the starting policy π0 for further alignment.The alignment step aims to retain information from earlier training while increasing expected reward.Source: §2 · p. 3
- Standard practice maximizes expected reward with a KL-divergence penalty, while Figure 1 compares transformed and untransformed reward aggregation across KL weights.Its y-axis reports improvement over the SFT model in both helpfulness and harmlessness, judged by an external evaluator.Source: §2 · p. 3, Figure 1 · p. 3
3 Reward Transformation
The reward transformation is derived by defining good outputs relative to a reference and matching KL-regularized alignment to a probabilistic target. The resulting centered log-sigmoid utility has diminishing returns above the reference, whose choice materially affects its behavior.
3 Reward Transformation
- The alignment goal is to sample from the base policy conditioned on an output being good, with a hyperparameter controlling how strongly goodness reweights the policy.The paper defines goodness relative to a preferred reference output and connects the target utility to the log-probability of goodness.Source: §3 · pp. 3–5
- Log-sigmoid transformation reflects diminishing returns: Bradley-Terry rewards treat successive improvements equally even when the later improvement adds little utility.The example contrasts the substantial gain from response 1 to 2 with the marginal gain from response 2 to 3.Source: Figure 2 · p. 4
- For pointwise reward models trained on binary good/bad labels, the right utility is log-sigmoid of the reward, which is the logit of goodness probability.The paper distinguishes this from pairwise rewards, whose individual values are not identifiable and whose prompt-specific shifts make nonlinear transformations sensitive to arbitrary offsets.Source: §3 · pp. 4–5
- For Bradley-Terry rewards, the LSC utility is log-sigmoid of the reward minus a prompt-specific reference reward.The theorem states this choice satisfies the alignment goal when the Bradley-Terry model and KL-regularized utility maximization hold; the reference reward is the reward of the chosen reference response.Source: §3 · p. 5
- Raw rewards can sharply upweight marginal improvements between already-good responses, whereas the transformed reward reduces that effect.In the example, raw-reward alignment increases relative probability by a factor of 10, while transformed-reward alignment increases it by 1.01.Source: §3 · p. 5
- The reference acts as a hyperparameter: utility is linear below it and rapidly diminishing above it, so it should represent an achievable good response where the learned reward is meaningful.The paper reports the 85th quantile of base-distribution examples as a good default.Source: §3 · p. 6
4 Reward Aggregation
The section formalizes reward aggregation as aligning outputs to be good on all target properties, then derives a transformed-reward utility under conditional independence of property judgments. It contrasts this approach with raw-reward summation, which can let strong performance on one property offset weak performance on another.
Reward Aggregation
- The individual reward models reveal only marginal property distributions, so constructing the probability of being good on all properties is impossible in full generality without an extra assumption.The section introduces conditional independence of judgments as that assumption.Source: §4 · p. 6
- Summing transformed rewards yields the alignment goal of being good across all properties, provided property judgments are conditionally independent given the prompt and response.The theorem’s result also assumes Bradley-Terry models for all properties and KL-regularized utility maximization.Source: §4 · pp. 6–7
- Raw-reward summation allows strong performance on one property to offset poor performance on another, unlike the goal of being good on all target properties.The section frames the target as, for example, being helpful and harmless.Source: §4 · pp. 6–7
- Approximately 1.8 is the relative probability ratio under transformed-reward alignment for the response that does well on both properties; raw-logit summation does not upweight that response.The example uses two properties and sets γ = 1 for the raw-reward comparison.Source: §4 · p. 7
5 Experiments
Experiments find that LSC-transformation improves single-reward alignment and that summing transformed rewards improves alignment to helpfulness and harmlessness. The experiments also examine the reward-distribution, shortcut, and aggregation effects associated with these improvements.
5.2 LSC-Transformation Improves Alignment
- Across all KL distances and both evaluation strategies, transformed reward dominates raw reward in single-reward alignment.The evaluation strategies are a prompted PALM-2 evaluator and a T5-XXL evaluator.Source: §5.2 · p. 9
- At approximately matched KL, transformed alignment yields fewer very high and very low reward samples than raw-reward alignment.The authors interpret these distributions as indicating less reward hacking and less underfitting, respectively.Source: §5.2 · p. 9, Figure 4 · p. 10
- Transformed-reward alignment substantially mitigates list-formatting and therapy-referral shortcuts that raw-reward alignment exploits.The shortcuts are identified using heuristic methods.Source: §5.2 · pp. 9–10, Figure 5 · p. 10
5.3.1 Further Experimental Setup
- For aggregation, the experiments evaluate both RLHF and best-of-k sampling, ranking SFT samples by combined reward, with k increasing from 1 to 191.The evaluation counts a win only when a sample exceeds the SFT median on both helpfulness and harmlessness.Source: §5.3.1 · pp. 10–12
- The authors caution that their aggregation reference-reward choices are likely not optimal and that further hyperparameter optimization could improve results.The setup uses SFT medians for best-of-k and the SFT 75%-quantile for helpfulness in RLHF.Source: §5.3.1 · p. 11
5.3.2 Aggregation Results
- Transformed reward improves over the SFT model for both best-of-k and RLHF, with the advantage over raw-reward aggregation especially pronounced at high KL.At low KL in RLHF and in best-of-k, the authors attribute the win mainly to the logical-AND effect; at high KL, they report reward hacking with raw rewards.Source: §5.3.2 · pp. 12–13
- At matched K for best-of-k and KL for RLHF, transformed aggregation produces more balanced helpfulness and harmlessness reward distributions than the baseline.For best-of-k, k = 191 has a different distribution from RLHF at KL ≈14, but the transformed-aggregation effect is consistent across both.Source: Figure 7 · p. 12, §5.3.2 · p. 13
6 Ablation Studies
The LSC transformation combines centering and log-sigmoid transformation, but neither component alone improves alignment: centering leaves performance unchanged, while log-sigmoid alone hurts it.
6 Ablation Studies
- Centering alone does not change RLHF performance relative to the baseline, despite reducing variance in policy-gradient updates.It does not change the optimization target of the KL-regularized RLHF objective.Source: §6 · p. 13
- Log-sigmoid transformation alone hurts performance relative to baseline because it uses a constant reference value of zero.A reference that is too low can saturate utility while rewards are still underfit; one that is too high can make the transformation ineffective.Source: §6 · pp. 13–14
- A fixed reference reward is not meaningful because Bradley–Terry rewards are non-identifiable up to prompt-dependent offsets.The paper instead chooses reference rewards based on generated responses to set a meaningful threshold.Source: §6 · p. 14
7 Discussion and Related Work
The discussion positions the transformation as complementary to other reward-hacking mitigations and as a way to aggregate rewards for multiple properties. It also connects bounded utility to related work and identifies direct preference-based alignment as a question for future study.
7 Discussion and Related Work
- The transformation is complementary to reward-hacking mitigations including reward-model averaging, constrained optimization, regularization, iterative preference collection, and data-bias mitigation.The authors say these approaches could be used in combination.Source: §7 · p. 14
- Prior multi-objective approaches include weighted sums, constrained MDPs, social-welfare aggregation, and a single helpfulness-and-harmlessness reward model.The cited single-model approach encountered harmlessness reward hacking and adjusted the helpfulness-data proportion; summing transformed rewards is presented as an alternative.Source: §7 · p. 14
- The transformation applies to strategies that maximize expected utility, and the authors report gains in best-of-k experiments that rank samples using combined reward models.For alignment methods using preference labels directly, aggregating rankings remains necessary when targeting multiple properties.Source: §7 · p. 14
- The paper frames its approach partly as incorporating the need for bounded utility [Aza+23] within the standard utility-maximization pipeline.Whether transformed rewards improve alignment compared with ranked pairs alone remains a future-work question.Source: §7 · p. 14
A.1 Implementation for reference value prediction
The appendix describes prompt-specific reward-reference prediction, PALM-2 evaluation, and heuristics for finding shortcut behaviors in aligned outputs.
A.1 Implementation for reference value prediction
- The reference-value implementation assumes a Gaussian reward distribution per prompt and predicts its mean and standard deviation from sampled rewards.Training data uses 64 SFT-policy responses per prompt; the predicted statistics are used to estimate a chosen quantile.Source: §A.1 · p. 17
A.2 PALM-2 Evaluation Details
- PALM-2 compares aligned-policy and SFT responses using zero-shot prompts, with both response orderings and eight samples per ordering decided by majority vote.The two orderings are used to counter positional bias; Figure 9 is the evaluation prompt template.Source: §A.2 · p. 17
A.3 Heuristics for shortcut discovery
- Shortcut discovery uses list-formatted helpfulness responses and analyzes harmlessness generations at KL ≈14 for recurring profession-related advice.The authors identify a dominant topic around “you should talk to [profession]” and search for responses containing related profession words.Source: §A.3 · p. 17
A.4 RLHF training details
RLHF training uses PPO and sweeps its regularization settings to obtain policies at different KL values. The transformed-reward runs use less KL regularization than raw-reward runs at matched KL.
A.4 RLHF training details
- PPO training sweeps γ to produce aligned policies at different KL values, using a heuristic to predict converged KL for each setting.The parameterization α = γ/(1 + γ) is used, with values listed in Table 2.Source: §A.4 · p. 17, Table 2 · p. 19
- At the same KL value, transformed-reward alignment uses less KL regularization than raw-reward alignment.The authors explain that log-sigmoid transformation prevents using KL budgets once reward saturates.Source: §A.4 · p. 17
- The appendix identifies Table 1 as the RLHF hyperparameter table and Figure 9 as the PALM-2 evaluation prompt.Table 2 lists α values for the task-method combinations.Source: Figure 9 · p. 18, Table 1 · p. 19, Table 2 · p. 19
B More Experiment Results
The experiments report that reward transformation mitigates reward overoptimization, especially for high-reward responses, while raw-reward aggregation worsens reward hacking. Additional figures report more concentrated reward distributions and better KL–win-rate trade-offs for transformed rewards.
B More Experiment Results
- Figure 11 reports more concentrated reward distributions with transformed rewards, a trend observed across KL values.Its caption characterizes the result as more uniform reward improvements.Source: B More Experiment Results · p. 19, Figure 11 · p. 20
- Reward transformation alleviates overoptimization most for larger-reward responses, while raw-reward aggregation exacerbates reward hacking relative to single-reward alignment.The policy can over-optimize helpfulness despite harmlessness having a smaller weight.Source: B More Experiment Results · p. 19, Figure 10 · p. 20
- Figure 12 reports better KL and win-rate trade-offs for transformed rewards, comparing win-rates against SFT random samples and reward quantiles.The quantile comparisons use the 85th percentile for helpfulness and 95th percentile for harmlessness.Source: B More Experiment Results · p. 19, Figure 12 · p. 21