Bradley-Terry turns a ranking into a loss, the KL penalty stops reward hacking, and DPO does it without a reward model.
The notebook carries the source, the install step and every run below, so nothing has to be on your machine. Open it, then choose Runtime › Run all.
Colab opens it read only. Click Copy to Drive to keep
your changes. You can also download
alignment.ipynb, or run it
locally: python3 alignment.py.
python3 alignment.py
All three stages of preference tuning.
python3 alignment.py --bt
Four preference pairs. The pair the model gets backwards costs 1.1711, more than three times the pair it gets right, and an undecided pair costs exactly log 2.
python3 alignment.py --kl
A hack scoring 5.00 beats an honest answer scoring 4.00 until beta reaches 0.10, where its distance from the reference finally costs more than the reward it buys.
python3 alignment.py --dpo
The implicit reward margin starts at 0 with a loss of log 2, the same number the undecided Bradley-Terry pair gives, and falls as the policy separates the two responses.
Download alignment.py
· served verbatim at https://nlp.jcrlabz.com/code/alignment.py
"""Preference tuning, worked as Chapter 20 works it.
Three things, all reproducing the chapter's tables:
the reward model Bradley-Terry turns a ranking into a loss
the KL penalty why reward hacking needs a leash, in numbers
DPO the same preference, with no reward model
python3 alignment.py # all three
python3 alignment.py --bt # Bradley-Terry and the reward-model loss
python3 alignment.py --kl # reward hacking and the penalty that stops it
python3 alignment.py --dpo # direct preference optimisation
Install: nothing, the Python standard library is enough
"""
import argparse
import math
def sigmoid(z):
return 1.0 / (1.0 + math.exp(-z))
# ------------------------------------------------------------ Bradley-Terry
# Four preference pairs. r_w is the reward the model gives the response the
# human preferred, r_l the one it gives the rejected response. The model is
# right on the first two, wrong on the third, and undecided on the fourth.
PAIRS = [
("clear answer", 1.20, 0.40),
("answer with steps", 2.10, 0.30),
("confident bluff", 0.30, 1.10),
("two paraphrases", 0.75, 0.75),
]
def show_bt():
print("\nSTEP 1 the reward model learns from comparisons\n")
print(" Nobody can score helpfulness on an absolute scale, but people can")
print(" say which of two answers they prefer. Bradley-Terry turns that")
print(" comparison into a probability:\n")
print(" P(y_w > y_l) = sigma(r_w - r_l)\n")
print(" pair r_w r_l r_w-r_l P(y_w>y_l) loss")
print(" " + "-" * 64)
total = 0.0
for name, rw, rl in PAIRS:
d = rw - rl
p = sigmoid(d)
loss = -math.log(p)
total += loss
print(f" {name:18} {rw:5.2f} {rl:5.2f} {d:+7.2f} {p:9.4f} {loss:7.4f}")
print(" " + "-" * 64)
print(f" mean loss over the four pairs{' ' * 29}{total / len(PAIRS):7.4f}\n")
print(" Read the third row. The model prefers the bluff, so the pair costs")
print(f" {-math.log(sigmoid(0.30 - 1.10)):.4f}, more than three times the first row. Training")
print(" pushes r_w up and r_l down until that ordering flips.\n")
print(" Read the fourth. Equal rewards give exactly 0.5 and a loss of")
print(f" {-math.log(0.5):.4f} = log 2, the most a single comparison can tell you")
print(" when the model has no opinion.")
# ------------------------------------------------------------- the KL leash
# Two candidate responses the policy could learn to produce. The first scores
# well with the reward model and badly with people: that is reward hacking.
CANDIDATES = [
("honest answer", 4.00, 1.00),
("reward hack", 5.00, 12.00),
]
BETAS = [0.00, 0.05, 0.10, 0.30]
def show_kl():
print("\nSTEP 2 the penalty is not an afterthought\n")
print(" The policy is optimised against the reward model, not against")
print(" people. Left alone it finds text the reward model loves and humans")
print(" do not. The objective keeps it on a leash:\n")
print(" maximise E[r(x,y)] - beta * KL(pi_theta || pi_ref)\n")
print(" candidate r KL" + "".join(f" beta={b:.2f}" for b in BETAS))
print(" " + "-" * (31 + 11 * len(BETAS)))
scores = {}
for name, r, kl in CANDIDATES:
row = [r - b * kl for b in BETAS]
scores[name] = row
cells = "".join(f" {v:9.2f}" for v in row)
print(f" {name:15} {r:5.2f} {kl:6.2f}{cells}")
print(" " + "-" * (31 + 11 * len(BETAS)))
print("\n winner per beta: ", end="")
flip = None
for i, b in enumerate(BETAS):
best = max(scores, key=lambda n: scores[n][i])
print(f"beta={b:.2f} -> {best} ", end="")
if best == "honest answer" and flip is None:
flip = b
print()
print(f"\n With no penalty the hack wins, because 5.00 beats 4.00 and the")
print(f" reward model is the only judge. The ordering flips at beta = {flip:.2f},")
print(" where the hack's distance from the starting model finally costs")
print(" more than the extra reward it buys.\n")
print(" Too large a beta is its own failure. The policy then stays so close")
print(" to the reference that preference training changes nothing.")
# ---------------------------------------------------------------------- DPO
# Log-probabilities of the preferred and rejected response under the policy
# being trained and under the frozen reference it started from.
BETA_DPO = 0.1
LOGP = [
("at the start", -2.30, -2.30, -2.20, -2.20),
("after one step", -2.00, -2.30, -2.50, -2.20),
("over-corrected", -1.40, -2.30, -3.40, -2.20),
]
def show_dpo():
print("\nSTEP 3 the same preference, with no reward model\n")
print(" DPO drops the reward model and the reinforcement learning loop. It")
print(" reads an implicit reward straight off the policy's own log-ratio:\n")
print(" r_hat(y) = beta * ( log pi_theta(y) - log pi_ref(y) )")
print(" loss = -log sigma( r_hat(y_w) - r_hat(y_l) )\n")
print(f" beta = {BETA_DPO}\n")
print(" stage r_hat(y_w) r_hat(y_l) margin sigma loss")
print(" " + "-" * 66)
for name, tw, rw_, tl, rl_ in LOGP:
rhat_w = BETA_DPO * (tw - rw_)
rhat_l = BETA_DPO * (tl - rl_)
margin = rhat_w - rhat_l
s = sigmoid(margin)
loss = -math.log(s)
print(f" {name:16} {rhat_w:+10.4f} {rhat_l:+10.4f} {margin:+7.4f} {s:7.4f} {loss:7.4f}")
print(" " + "-" * 66)
print("\n At the start the policy is the reference, so both implicit rewards")
print(f" are 0, the margin is 0 and the loss is log 2 = {-math.log(0.5):.4f}. That is the")
print(" same number the undecided Bradley-Terry pair gave in Step 1, and for")
print(" the same reason: no opinion yet.\n")
print(" One step raises the preferred response and lowers the rejected one,")
print(" and the loss falls. The third row is the warning. Pushing the margin")
print(" further keeps shrinking the loss, but the policy has now moved a long")
print(" way from the reference, which is what beta is there to price.")
def main():
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--bt", action="store_true")
ap.add_argument("--kl", action="store_true")
ap.add_argument("--dpo", action="store_true")
a = ap.parse_args()
run_all = not (a.bt or a.kl or a.dpo)
if a.bt or run_all:
show_bt()
if a.kl or run_all:
show_kl()
if a.dpo or run_all:
show_dpo()
if __name__ == "__main__":
main()