Overoptimizing a reward model
Each dot is an answer the base model might write. Pushing toward the reward model's favorites moves the orange average; the chart tracks the score it reports and the true quality.
Reward model score+1.52
True quality+1.12
Distance from base model (KL)0.64 nats
Average length+0.84
Challenges
0 of 3 complete
- With the reward model as it starts (it learned none of the padding penalty), find the pressure where true quality peaks. (not yet)
- Push hard enough that the chosen answers are worse than the base model's average, while the reward model reports a gain of more than 2. (not yet)
- Teach the reward model enough about padding that more pressure never hurts true quality, all the way to the strongest push. (not yet)
Try this
- Slowly raise the optimization pressure. Watch the orange average answer drift, and watch both lines on the chart.
- Switch between best of n and the KL penalty. Is the pattern the same?
- Raise how much the reward model knows about padding, then push hard again.
The KL-penalized optimum, computed exactlyPython with NumPy
import numpy as np
# Every possible answer, described by two numbers in standard units:
# its length and its substance (useful content). The base model writes
# answers near (0, 0) more often, following a bell curve in both.
grid = np.linspace(-6, 6, 241)
length, substance = np.meshgrid(grid, grid)
base = np.exp(-(length**2 + substance**2) / 2)
base /= base.sum()
# The true goal: substance, plus some length, minus a penalty for padding.
true_quality = substance + length - (length**2 - 1)
# The reward model learned "longer is better" and none of the padding penalty.
reward_model = substance + length
def optimize(beta):
"""Exact optimum of E[reward model] - beta * KL(policy || base)."""
w = base * np.exp((reward_model - reward_model.max()) / beta)
policy = w / w.sum()
kl = np.sum(policy * np.log(policy / base))
gain = lambda score: np.sum(policy * score) - np.sum(base * score)
return kl, gain(reward_model), gain(true_quality), np.sum(policy * length)
print(" beta KL reward model true quality avg length")
for beta in [4, 2, 1, 0.6, 0.4, 0.3]:
kl, rm, tq, ln = optimize(beta)
print(f"{beta:5.2f} {kl:5.1f} {rm:+13.2f} {tq:+14.2f} {ln:+12.2f}") beta KL reward model true quality avg length
4.00 0.1 +0.50 +0.44 +0.25
2.00 0.2 +1.00 +0.75 +0.50
1.00 1.0 +2.00 +1.00 +1.00
0.60 2.8 +3.33 +0.56 +1.67
0.40 6.2 +5.00 -1.24 +2.50
0.30 11.0 +6.65 -4.37 +3.32