RLHF — 从偏好到奖励的实验#
验证 rlhf.tex 中的核心结论：
奖励建模：Bradley-Terry 模型从偏好对学出奖励，验证集精度随数据量上升——但奖励只在偏好数据覆盖区内可信，覆盖区外单调外推；
Reward hacking：无 KL 约束（β=0）的 PPO 先带来真实提升，随后击穿奖励模型——代理奖励一路上涨，真实分数崩塌；
KL 锚定：适中的 β 把策略稳定在真提升区；过大的 β 把策略钉死在 SFT 附近。
环境是一个完全合成的玩具序列任务（词表 24、序列长 12），使真实质量可精确计算：真实分数 = token 价值均值 + 强调 token 的凹形加分（每个加 0.3 分、超过 3 个后每个倒扣 0.6 分——「适量加分、过量崩坏」，类比真实 RLHF 中的长度偏置）。SFT 策略几乎不产生超过 3 个强调 token 的序列，因此偏好数据只覆盖加分的上升段。
流程骨架仿照 RLHF，任务仍是玩具序列：SFT 预训练 → SFT 采样 + 真实质量按 BT 模型标注偏好 → 训练奖励模型（均值池化 MLP，按 SFT 分布白化）→ PPO 微调（批均值基线替代 Critic，完整 PPO 见第 7 章）。
输出图：
fig1_reward_model.pdffig2_reward_hacking.pdffig3_kl_anchor.pdf
预计运行时间：GPU 约 5 分钟；CPU 约 20–30 分钟（3 seeds）。
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib
import matplotlib.pyplot as plt
matplotlib.rcParams.update({
'font.family': 'serif',
'font.serif': ['Times New Roman', 'DejaVu Serif'],
'pdf.fonttype': 42,
'ps.fonttype': 42,
'axes.labelsize': 11,
'xtick.labelsize': 9,
'ytick.labelsize': 9,
'legend.fontsize': 9,
})
BLUE = '#2166AC'
RED = '#D6604D'
GRAY = '#808080'
FIGSIZE = (7.2, 4.8)
OUTDIR = '.'
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
VOCAB = 24 # ordinary tokens 0..23
BOS = VOCAB
SEQ_LEN = 12
E_TOKEN = 12 # the "emphasis" token (mid value ~0)
T_BASE = 1.5 # temperature of the base corpus distribution
D_MODEL = 64
SEEDS = (42, 43, 44)
TOKEN_VALUES = np.linspace(-1.0, 1.0, VOCAB)
TOKEN_VALUES_T = torch.tensor(TOKEN_VALUES, dtype=torch.float32, device=DEVICE)
PEAK_C = 3 # emphasis bonus peaks at 3 occurrences
BONUS_UP = 0.3 # per occurrence up to the peak
BONUS_DOWN = 0.6 # per occurrence beyond the peak
def bonus(count):
up = BONUS_UP * torch.clamp(count.float(), max=PEAK_C)
down = BONUS_DOWN * torch.clamp(count.float() - PEAK_C, min=0)
return up - down
def true_score(seqs):
"""Ground-truth quality: mean token value + concave emphasis bonus."""
vals = TOKEN_VALUES_T[seqs].mean(dim=1)
return vals + bonus((seqs == E_TOKEN).sum(dim=1))
def set_seed(seed):
np.random.seed(seed)
torch.manual_seed(seed)
def style_axes(ax):
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.grid(True, linestyle='--', linewidth=0.5, alpha=0.7)
print(f'Device: {DEVICE}')
Device: cuda
class TinyLM(nn.Module):
"""Char-level causal Transformer LM (~120k params)."""
def __init__(self):
super().__init__()
self.emb = nn.Embedding(VOCAB + 1, D_MODEL)
self.pos = nn.Embedding(SEQ_LEN + 1, D_MODEL)
layer = nn.TransformerEncoderLayer(D_MODEL, 4, 128, batch_first=True,
dropout=0.0, norm_first=True)
self.encoder = nn.TransformerEncoder(layer, 2, enable_nested_tensor=False)
self.head = nn.Linear(D_MODEL, VOCAB)
mask = torch.triu(torch.full((SEQ_LEN + 1, SEQ_LEN + 1), float('-inf')), 1)
self.register_buffer('mask', mask)
def logits(self, inp):
L = inp.size(1)
h = self.emb(inp) + self.pos(torch.arange(L, device=inp.device))
h = self.encoder(h, mask=self.mask[:L, :L])
return self.head(h)
@torch.no_grad()
def generate(self, batch):
inp = torch.full((batch, 1), BOS, dtype=torch.long, device=DEVICE)
for _ in range(SEQ_LEN):
logits = self.logits(inp)[:, -1]
nxt = torch.multinomial(F.softmax(logits, dim=-1), 1)
inp = torch.cat([inp, nxt], dim=1)
return inp[:, 1:]
def log_probs(self, seqs):
inp = torch.cat([torch.full((seqs.size(0), 1), BOS, dtype=torch.long,
device=DEVICE), seqs[:, :-1]], dim=1)
logits = self.logits(inp)
return torch.log_softmax(logits, dim=-1).gather(2, seqs.unsqueeze(2)).squeeze(2)
class RewardModel(nn.Module):
"""Mean-pooled MLP scored by Bradley-Terry; whitened on the SFT distribution."""
def __init__(self):
super().__init__()
self.emb = nn.Embedding(VOCAB, D_MODEL)
self.net = nn.Sequential(nn.Linear(D_MODEL, D_MODEL), nn.ReLU(),
nn.Linear(D_MODEL, 1))
self.mu = 0.0
self.sigma = 1.0
def raw(self, seqs):
return self.net(self.emb(seqs).mean(dim=1)).squeeze(1)
def forward(self, seqs):
return (self.raw(seqs) - self.mu) / self.sigma
@torch.no_grad()
def calibrate(self, base_seqs):
r = self.raw(base_seqs)
self.mu = r.mean().item()
self.sigma = r.std().item() + 1e-8
def sample_base_corpus(n):
"""The 'human text' distribution: i.i.d. tokens, softmax of values / T_BASE."""
probs = F.softmax(torch.tensor(TOKEN_VALUES / T_BASE), dim=0).numpy()
return torch.tensor(np.random.choice(VOCAB, size=(n, SEQ_LEN), p=probs),
dtype=torch.long, device=DEVICE)
def train_sft(steps=800, batch=256, lr=3e-3):
lm = TinyLM().to(DEVICE)
opt = torch.optim.Adam(lm.parameters(), lr=lr)
for _ in range(steps):
loss = -lm.log_probs(sample_base_corpus(batch)).mean()
opt.zero_grad(); loss.backward(); opt.step()
return lm
def make_pref_pairs(lm, n_pairs, noise=0.1):
"""Sample two completions from the policy, label by noisy true score (BT)."""
a = lm.generate(n_pairs)
b = lm.generate(n_pairs)
sa, sb = true_score(a), true_score(b)
a_wins = torch.bernoulli(torch.sigmoid((sa - sb) / noise)).bool()
return (torch.where(a_wins.unsqueeze(1), a, b),
torch.where(a_wins.unsqueeze(1), b, a))
def train_rm(winners, losers, base, steps=600, batch=128, lr=1e-3):
rm = RewardModel().to(DEVICE)
opt = torch.optim.Adam(rm.parameters(), lr=lr)
n = winners.size(0)
for _ in range(steps):
idx = torch.randint(0, n, (batch,), device=DEVICE)
loss = -F.logsigmoid(rm.raw(winners[idx]) - rm.raw(losers[idx])).mean()
opt.zero_grad(); loss.backward(); opt.step()
rm.calibrate(base)
return rm
@torch.no_grad()
def rm_accuracy(rm, winners, losers):
return (rm(winners) > rm(losers)).float().mean().item()
@torch.no_grad()
def rm_vs_count_curve(rm, n=512):
"""Mean RM score and true score of sequences with exactly c emphasis tokens."""
rows = []
for c in range(SEQ_LEN + 1):
seqs = sample_base_corpus(n)
seqs = seqs.masked_fill(seqs == E_TOKEN, E_TOKEN + 1)
for i in range(n):
pos = torch.randperm(SEQ_LEN, device=DEVICE)[:c]
seqs[i, pos] = E_TOKEN
rows.append((rm(seqs).mean().item(), true_score(seqs).mean().item()))
return rows
def ppo_finetune(sft, rm, beta, iters=150, batch=256, k_epochs=4, mb=64,
clip=0.2, lr=3e-5):
"""PPO-clip against the learned reward with a per-sequence KL penalty.
Sequence-level bandit formulation with a batch-mean baseline instead of a
critic, to keep the focus on the reward problem (full PPO: chapter 7).
"""
policy = TinyLM().to(DEVICE)
policy.load_state_dict(sft.state_dict())
ref = TinyLM().to(DEVICE)
ref.load_state_dict(sft.state_dict())
ref.eval()
opt = torch.optim.Adam(policy.parameters(), lr=lr)
history = {'proxy': [], 'truth': [], 'kl': [], 'count_e': []}
for it in range(iters):
with torch.no_grad():
seqs = policy.generate(batch)
lp_old = policy.log_probs(seqs)
lp_ref = ref.log_probs(seqs)
proxy = rm(seqs)
kl_seq = (lp_old - lp_ref).sum(dim=1)
shaped = proxy - beta * kl_seq
adv = (shaped - shaped.mean()) / (shaped.std() + 1e-8)
for _ in range(k_epochs):
perm = torch.randperm(batch, device=DEVICE)
for s in range(0, batch, mb):
idx = perm[s:s + mb]
lp_new = policy.log_probs(seqs[idx])
ratio = torch.exp((lp_new - lp_old[idx]).sum(dim=1))
s1 = ratio * adv[idx]
s2 = torch.clamp(ratio, 1 - clip, 1 + clip) * adv[idx]
loss = -torch.min(s1, s2).mean()
opt.zero_grad(); loss.backward()
nn.utils.clip_grad_norm_(policy.parameters(), 1.0)
opt.step()
history['proxy'].append(proxy.mean().item())
history['truth'].append(true_score(seqs).mean().item())
history['kl'].append(kl_seq.mean().item())
history['count_e'].append((seqs == E_TOKEN).float().sum(dim=1).mean().item())
return history
print('=== SFT + preference data + reward model, per seed ===')
runs = {}
for seed in SEEDS:
set_seed(seed)
sft = train_sft()
base = sft.generate(2000)
winners, losers = make_pref_pairs(sft, 4000)
rm = train_rm(winners, losers, base)
vw, vl = make_pref_pairs(sft, 2000)
acc = rm_accuracy(rm, vw, vl)
ce = (base == E_TOKEN).float().sum(dim=1)
runs[seed] = {'sft': sft, 'base': base, 'rm': rm,
'winners': winners, 'losers': losers}
print(f' seed {seed}: base truth {true_score(base).mean():+.3f} | '
f'P(c_E > {PEAK_C}) = {(ce > PEAK_C).float().mean():.4f} | '
f'RM val acc {acc:.3f}')
=== SFT + preference data + reward model, per seed ===
seed 42: base truth +0.389 | P(c_E > 3) = 0.0010 | RM val acc 0.866
seed 43: base truth +0.384 | P(c_E > 3) = 0.0015 | RM val acc 0.853
seed 44: base truth +0.386 | P(c_E > 3) = 0.0005 | RM val acc 0.853
Figure 1 — 奖励模型：数据内可信，数据外单调外推#
左图：奖励模型的验证集偏好精度随偏好对数量上升——数据越多，代理越接近真实偏好。右图是本章的机制核心：在偏好数据覆盖区（SFT 策略几乎只产生 \(c_E \le 3\) 的序列，灰色区域）内，奖励模型与真实分数一致；覆盖区外，真实分数在 \(c_E=3\) 见顶后崩坏，而 ReLU 网络在支撑集外线性外推，奖励模型给出的分数一路上扬——它从数据里学到的是「强调 token 越多越好」。这条分叉就是 reward hacking 的入口。
print('=== Experiment 1: reward model quality and extrapolation ===')
PAIR_SIZES = (250, 500, 1000, 2000, 4000)
acc_matrix = np.zeros((len(SEEDS), len(PAIR_SIZES)))
curves = []
for si, seed in enumerate(SEEDS):
set_seed(seed + 1000)
r = runs[seed]
vw, vl = make_pref_pairs(r['sft'], 2000)
for ni, n in enumerate(PAIR_SIZES):
rm_n = train_rm(r['winners'][:n], r['losers'][:n], r['base'])
acc_matrix[si, ni] = rm_accuracy(rm_n, vw, vl)
curves.append(rm_vs_count_curve(r['rm']))
print(f' seed {seed}: acc {np.round(acc_matrix[si], 3)}')
curves = np.array(curves) # (seeds, c, [rm, true])
=== Experiment 1: reward model quality and extrapolation ===
seed 42: acc [0.822 0.827 0.846 0.847 0.855]
seed 43: acc [0.823 0.835 0.854 0.862 0.861]
seed 44: acc [0.824 0.828 0.841 0.852 0.848]
fig, axes = plt.subplots(1, 2, figsize=(10, 4.2))
ax = axes[0]
mu, sd = acc_matrix.mean(0), acc_matrix.std(0)
ax.plot(PAIR_SIZES, mu, color=BLUE, marker='o', linewidth=1.5)
ax.fill_between(PAIR_SIZES, mu - sd, mu + sd, color=BLUE, alpha=0.15)
ax.axhline(0.5, color=GRAY, linestyle='--', linewidth=1.0, label='random guess')
ax.set_xscale('log')
ax.set_xticks(PAIR_SIZES)
ax.set_xticklabels(PAIR_SIZES)
style_axes(ax)
ax.set_xlabel('Preference pairs')
ax.set_ylabel('Validation accuracy')
ax.set_title('Reward model accuracy vs data', pad=8)
ax.legend(loc='lower right')
ax = axes[1]
cs = np.arange(SEQ_LEN + 1)
rm_mu, rm_sd = curves[:, :, 0].mean(0), curves[:, :, 0].std(0)
ax.plot(cs, rm_mu, color=BLUE, marker='o', markersize=3, linewidth=1.5,
label='RM score (whitened)')
ax.fill_between(cs, rm_mu - rm_sd, rm_mu + rm_sd, color=BLUE, alpha=0.15)
ax2 = ax.twinx()
ax2.plot(cs, curves[0, :, 1], color=RED, marker='s', markersize=3,
linewidth=1.5, label='true score')
ax.axvspan(0, PEAK_C, color=GRAY, alpha=0.15)
ax.text(PEAK_C / 2, ax.get_ylim()[1] * 0.9, 'data\ncoverage',
ha='center', fontsize=8, color=GRAY)
style_axes(ax)
ax2.spines['top'].set_visible(False)
ax.set_xlabel('Emphasis-token count $c_E$')
ax.set_ylabel('RM score (whitened)', color=BLUE)
ax2.set_ylabel('True score', color=RED)
ax.set_title('In-support agreement, out-of-support divergence', pad=8)
h1, l1 = ax.get_legend_handles_labels()
h2, l2 = ax2.get_legend_handles_labels()
ax.legend(h1 + h2, l1 + l2, loc='center right')
fig.tight_layout()
fig.savefig(f'{OUTDIR}/fig1_reward_model.pdf', bbox_inches='tight')
plt.show()
print('Saved fig1_reward_model.pdf')
Saved fig1_reward_model.pdf
Figure 2 — Reward hacking：代理奖励上涨，真实分数崩塌#
β=0（无 KL 锚定）的 PPO 全程只看奖励模型。前约 30 轮是真提升：策略把强调 token 用到甜点区附近，真实分数从 +0.39 升到约 +1.0；随后优化冲出数据覆盖区，把强调 token 打满到 12 个——代理奖励继续上涨到 +7 以上，真实分数崩到 −4.5。代理指标越来越好、真实目标越来越差，与第 9 章离线 RL 的「Q 涨、回报跌」同构（3 seeds，阴影为 \(\pm 1\sigma\)）。
print('=== Experiment 2: PPO with beta = 0 (no KL anchor) ===')
hists = {}
for seed in SEEDS:
set_seed(seed + 2000)
r = runs[seed]
hists[(seed, 0.0)] = ppo_finetune(r['sft'], r['rm'], beta=0.0)
h = hists[(seed, 0.0)]
print(f" seed {seed}: peak truth {max(h['truth']):+.3f} | "
f"final proxy {h['proxy'][-1]:+.3f} | final truth {h['truth'][-1]:+.3f} | "
f"final c_E {h['count_e'][-1]:.1f}")
=== Experiment 2: PPO with beta = 0 (no KL anchor) ===
seed 42: peak truth +1.068 | final proxy +7.050 | final truth -4.457 | final c_E 12.0
seed 43: peak truth +0.999 | final proxy +7.051 | final truth -4.457 | final c_E 12.0
seed 44: peak truth +1.060 | final proxy +8.038 | final truth -4.457 | final c_E 12.0
def stack(key, beta):
return np.array([hists[(seed, beta)][key] for seed in SEEDS])
fig, ax = plt.subplots(figsize=FIGSIZE)
its = np.arange(1, 151)
proxy = stack('proxy', 0.0)
ax.plot(its, proxy.mean(0), color=BLUE, linewidth=1.5, label='proxy reward (RM)')
ax.fill_between(its, proxy.mean(0) - proxy.std(0), proxy.mean(0) + proxy.std(0),
color=BLUE, alpha=0.15)
ax.set_xlabel('PPO iteration')
ax.set_ylabel('Proxy reward (RM, whitened)', color=BLUE)
ax2 = ax.twinx()
truth = stack('truth', 0.0)
ax2.plot(its, truth.mean(0), color=RED, linewidth=1.5, label='true score')
ax2.fill_between(its, truth.mean(0) - truth.std(0), truth.mean(0) + truth.std(0),
color=RED, alpha=0.15)
ax2.set_ylabel('True score', color=RED)
style_axes(ax)
ax2.spines['top'].set_visible(False)
ax.set_title('Fig 2. Reward hacking under $\\beta=0$ (3 seeds, mean $\\pm$ std)', pad=8)
h1, l1 = ax.get_legend_handles_labels()
h2, l2 = ax2.get_legend_handles_labels()
ax.legend(h1 + h2, l1 + l2, loc='center right')
fig.tight_layout()
fig.savefig(f'{OUTDIR}/fig2_reward_hacking.pdf', bbox_inches='tight')
plt.show()
print('Saved fig2_reward_hacking.pdf')
Saved fig2_reward_hacking.pdf
Figure 3 — KL 锚定：β 的三种结局#
在奖励里加上 \(-\beta\,\mathrm{KL}(\pi \| \pi_{\mathrm{ref}})\) 把策略拴在 SFT 参考策略附近。β=0 崩塌（同 Figure 2）；β=0.5 把最终 KL 压在约 2 nats，策略停在数据覆盖区边缘的甜点附近，真实分数稳定在约 +0.84；β=2.0 锚得太紧，策略几乎回不出 SFT 的水平（约 +0.5）。KL 锚定不消除奖励模型的缺陷，只是限制策略走进代理与真实分道扬镳的区域——β 是「利用代理奖励」与「防止击穿」之间的权衡（3 seeds，阴影为 \(\pm 1\sigma\)）。
print('=== Experiment 3: KL coefficient sweep ===')
for beta in (0.5, 2.0):
for seed in SEEDS:
set_seed(seed + 3000 + int(beta * 10))
r = runs[seed]
hists[(seed, beta)] = ppo_finetune(r['sft'], r['rm'], beta=beta)
truths = np.array([hists[(s, beta)]['truth'][-1] for s in SEEDS])
kls = np.array([hists[(s, beta)]['kl'][-1] for s in SEEDS])
print(f' beta {beta}: final truth {truths.mean():+.3f} ± {truths.std():.3f} | '
f'final KL {kls.mean():.1f}')
=== Experiment 3: KL coefficient sweep ===
beta 0.5: final truth +0.840 ± 0.019 | final KL 1.7
beta 2.0: final truth +0.525 ± 0.011 | final KL 0.1
fig, ax = plt.subplots(figsize=FIGSIZE)
its = np.arange(1, 151)
base_truth = np.mean([true_score(runs[s]['base']).mean().item() for s in SEEDS])
for beta, color, label in [(0.0, RED, r'$\beta=0$ (no anchor)'),
(0.5, BLUE, r'$\beta=0.5$'),
(2.0, GRAY, r'$\beta=2.0$')]:
truth = stack('truth', beta)
ax.plot(its, truth.mean(0), color=color, linewidth=1.5, label=label)
ax.fill_between(its, truth.mean(0) - truth.std(0),
truth.mean(0) + truth.std(0), color=color, alpha=0.15)
ax.axhline(base_truth, color='black', linestyle='--', linewidth=1.0,
label='SFT baseline')
style_axes(ax)
ax.set_xlabel('PPO iteration')
ax.set_ylabel('True score')
ax.set_title('Fig 3. KL anchoring (3 seeds, mean $\\pm$ std)', pad=8)
ax.legend(loc='lower left')
fig.tight_layout()
fig.savefig(f'{OUTDIR}/fig3_kl_anchor.pdf', bbox_inches='tight')
plt.show()
print('Saved fig3_kl_anchor.pdf')
Saved fig3_kl_anchor.pdf
小结#
奖励模型是近似的：BT 偏好学习在数据覆盖区内可信；覆盖区外，网络按自己的归纳偏置外推——本例中 ReLU 网络把「强调 token 适量加分」外推成「越多越好」；
优化器专找分叉：PPO 不知道什么是「真实质量」，它只会沿代理奖励最陡的方向走——恰恰是代理与真实分道扬镳的方向，真实分数先升后崩；
KL 锚定是缓解不是治愈：β 限制策略离开参考策略的距离，把优化关在数据覆盖区附近；β 过小挡不住击穿，β 过大浪费了偏好信息——它是超参，需要像 Figure 3 一样扫出来；
与第 9 章同构：离线 RL 的外推误差被
max放大，RLHF 的外推误差被 PPO 放大——同一个病：对数据支撑外的乐观估计，交给优化器就会被优化器放大。下一章 DPO 用「不学奖励模型」的方式绕开其中一半问题。