Overoptimizing a reward model

Push a model toward a flawed reward model's favorite answers with best-of-n or a KL penalty, and watch true quality rise, peak, and fall.

IntermediateExplained in Rewarding LLMs, and reward hacking

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.

How the model is pushed
n = 4
0%
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

  1. With the reward model as it starts (it learned none of the padding penalty), find the pressure where true quality peaks. (not yet)
  2. 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)
  3. 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

Try "embedding", "softmax", "overfitting", or "backpropagation".