Session 6.5 — Lab notebook

# ~1 min. Colab has torch and transformers; this installs nothing new on most runtimes.
import importlib.util, sys
IN_COLAB = 'google.colab' in sys.modules
if importlib.util.find_spec('transformers') is None:
    %pip install -q transformers

import numpy as np, torch, matplotlib.pyplot as plt
import torch.nn.functional as F
from transformers import (GPT2LMHeadModel, GPT2TokenizerFast,
                          AutoTokenizer, AutoModelForSequenceClassification)

torch.set_num_threads(2)
plt.rcParams.update({'figure.figsize': (7, 4.4), 'axes.grid': True, 'grid.alpha': 0.3})

SEED = 0            # change me in Explore ③
N_PAIRS = 400       # preference pairs the proxy trains on
PROJ_DIM = 48       # proxy capacity: features are projected down to this many dimensions
POOL = 48           # completions sampled per evaluation prompt
N_MAX = 32          # largest best-of-n; POOL must exceed it, see the note in ③
print('Colab:', IN_COLAB)

gen_tok = GPT2TokenizerFast.from_pretrained('gpt2'); gen_tok.pad_token = gen_tok.eos_token
gen = GPT2LMHeadModel.from_pretrained('gpt2').eval()          # the policy we sample from

GOLD = 'distilbert-base-uncased-finetuned-sst-2-english'
gold_tok = AutoTokenizer.from_pretrained(GOLD)
gold_m = AutoModelForSequenceClassification.from_pretrained(GOLD).eval()   # the gold scorer

@torch.no_grad()
def sample(prompt, k, new_tokens=16):
    """k sampled completions of one prompt."""
    ids = gen_tok([prompt] * k, return_tensors='pt', padding=True)
    out = gen.generate(**ids, do_sample=True, top_k=50, temperature=1.0,
                       max_new_tokens=new_tokens, pad_token_id=gen_tok.eos_token_id)
    return [gen_tok.decode(o, skip_special_tokens=True) for o in out]

@torch.no_grad()
def gold_score(texts):
    """Unbounded 'true quality': positive logit minus negative logit."""
    b = gold_tok(texts, return_tensors='pt', padding=True, truncation=True, max_length=64)
    lg = gold_m(**b).logits
    return (lg[:, 1] - lg[:, 0]).numpy()

@torch.no_grad()
def features(texts):
    """Frozen mean-pooled GPT-2 hidden states: the proxy never sees the gold model."""
    b = gen_tok(texts, return_tensors='pt', padding=True, truncation=True, max_length=64)
    h = gen.transformer(**b).last_hidden_state
    m = b['attention_mask'].unsqueeze(-1).float()
    return ((h * m).sum(1) / m.sum(1)).numpy()

print('gold scorer:', GOLD)

TRAIN_PROMPTS = ['The film was', 'I thought the restaurant', 'My experience with the product',
                 'The hotel we stayed in', 'Their customer service',
                 'The book I finished last night', 'The new phone', 'Our trip to the coast']

torch.manual_seed(SEED); rng = np.random.default_rng(SEED)
pool = [t for p in TRAIN_PROMPTS for t in sample(p, 16)]
pool_gold = gold_score(pool)
pool_feats = features(pool)

print(f'{len(pool)} completions, gold from {pool_gold.min():+.1f} to {pool_gold.max():+.1f}')
order = np.argsort(pool_gold)
print(f'\nworst:  {pool_gold[order[0]]:+.1f}  {pool[order[0]][:88]!r}')
print(f'best:   {pool_gold[order[-1]]:+.1f}  {pool[order[-1]][:88]!r}')

proj = rng.normal(size=(pool_feats.shape[1], PROJ_DIM)) / np.sqrt(pool_feats.shape[1])
X = torch.tensor(pool_feats @ proj, dtype=torch.float32)

# preference pairs, labelled by gold; skip near-ties, which carry no signal
pairs = []
while len(pairs) < N_PAIRS:
    i, j = rng.choice(len(pool), 2, replace=False)
    if abs(pool_gold[i] - pool_gold[j]) < 0.3:
        continue
    pairs.append((i, j) if pool_gold[i] > pool_gold[j] else (j, i))
pairs = np.array(pairs)
cut = int(0.8 * len(pairs)); train, test = pairs[:cut], pairs[cut:]

reward = torch.nn.Linear(PROJ_DIM, 1, bias=False)
opt = torch.optim.Adam(reward.parameters(), lr=5e-3)
for step in range(120):
    loss = -F.logsigmoid(reward(X[train[:, 0]]).squeeze(-1)
                         - reward(X[train[:, 1]]).squeeze(-1)).mean()
    opt.zero_grad(); loss.backward(); opt.step()

with torch.no_grad():
    acc = (reward(X[test[:, 0]]).squeeze(-1) > reward(X[test[:, 1]]).squeeze(-1)).float().mean().item()
print(f'trained on {len(train)} pairs, final loss {loss.item():.3f}')
print(f'held-out preference accuracy: {acc:.2f}')
print()
print('Expect roughly 0.6 to 0.75. A published reward model reaches about 0.7 on real human')
print('preferences, so this toy is not as far from the real thing as it deserves to be.')

EVAL_PROMPTS = ['The concert last night', 'This laptop', 'The service at the bank',
                'My flight home', 'The coffee shop downstairs', 'That documentary',
                'The software update', 'Their delivery', 'The meal we ordered',
                'This pair of shoes', 'The museum tour', 'My new headphones',
                'The train service', "This week's episode", 'Our landlord', 'The conference talk']

# The slowest cell: about 600 short completions. Two to four minutes on a Colab CPU.
proxy_pools, gold_pools = [], []
for p in EVAL_PROMPTS:
    outs = []
    for _ in range(POOL // 8):
        outs += sample(p, 8)
    f = torch.tensor(features(outs) @ proj, dtype=torch.float32)
    with torch.no_grad():
        proxy_pools.append(reward(f).squeeze(-1).numpy())
    gold_pools.append(gold_score(outs))
print(f'{len(EVAL_PROMPTS)} prompts x {POOL} completions scored by both models')

NS = [n for n in (1, 2, 4, 8, 16, 32, 64) if n <= N_MAX]
REPEATS = 40

proxy_curve, gold_curve = [], []
for n in NS:
    pv, gv = [], []
    for pr, go in zip(proxy_pools, gold_pools):
        for _ in range(REPEATS):
            subset = rng.choice(len(pr), n, replace=False)
            pick = subset[int(np.argmax(pr[subset]))]      # the proxy chooses
            pv.append(pr[pick]); gv.append(go[pick])       # gold only watches
    proxy_curve.append(np.mean(pv)); gold_curve.append(np.mean(gv))

kl = np.array([np.log(n) - (n - 1) / n for n in NS])
print(f"{'n':>4}  {'KL':>5}  {'proxy':>8}  {'gold':>8}")
for i, n in enumerate(NS):
    print(f'{n:>4}  {kl[i]:>5.2f}  {proxy_curve[i]:>8.2f}  {gold_curve[i]:>8.2f}')
print()
print(f'proxy rose {proxy_curve[-1] - proxy_curve[0]:+.2f} across the range; gold rose {gold_curve[-1] - gold_curve[0]:+.2f}.')
print(f'you optimised the proxy to KL = {kl[-1]:.2f} nats.')

# Gao et al. measure optimisation by d = sqrt(KL), NOT by the KL itself, and fit
# R(d) = d(alpha - beta d) with R(0) := 0. The square root is not decoration: fit a
# downward parabola in the wrong variable and you will invent curvature that is not there.
d = np.sqrt(kl)
delta = np.array(gold_curve) - gold_curve[0]         # R(0) := 0, as in the paper
A = np.stack([d, -d ** 2], axis=1)
(alpha, beta), *_ = np.linalg.lstsq(A, delta, rcond=None)

print(f'fitted alpha = {alpha:.3f}, beta = {beta:+.4f}')
if beta > 0:
    peak_d = alpha / (2 * beta)
    peak_kl = peak_d ** 2
    if peak_kl > 12:                                 # past anything Gao et al. validated
        print(f'predicted peak at d = {peak_d:.2f}, i.e. KL = {peak_kl:.0f} nats. That is far')
        digits = int((peak_kl + 1) / np.log(10)) + 1
        print(f'beyond the 10 nats the published form was checked at, and needs an n of about')
        print(f'{digits} digits. The fit has not located a peak; it has run out of curvature and')
        print('put the vertex somewhere off the map.')
    else:
        n_needed = np.exp(peak_kl + 1)               # KL ~ log n - 1 for large n
        print(f'predicted peak at d = {peak_d:.2f}, i.e. KL = {peak_kl:.1f} nats, n of about {n_needed:,.0f}')
else:
    print('beta came out negative, so the fitted parabola turns the wrong way and has no peak')
    print('at any positive distance. Over the range you covered, gold reward is still rising')
    print('and shows no sign of bending down. That is a result about your data. Report it.')
print()
print(f'You measured out to {kl[-1]:.1f} nats. Gao et al. fitted this form on data to n = 1,000')
print('(about 6 nats) and then validated it at n = 60,000 (about 10 nats). Best-of-n on a CPU')
print('cannot reach that: n = 1,000 here would be 20,000 generations per prompt.')
print()
print('So the turnover in 6.4 is real and you cannot see it from here. What you can do is fit')
print('the published form to your own data and extrapolate, which Session 2.5 told you to')
print('distrust. Explore ③ asks you to find out how much to distrust it.')

# The picture the numbers describe. Proxy and gold are shifted to start at zero so they
# share an axis; the dashed line is Gao's form with YOUR fitted alpha and beta.
fig, ax = plt.subplots()
ax.plot(d, np.array(proxy_curve) - proxy_curve[0], 'o-', label='proxy (what you optimise)')
ax.plot(d, delta, 's-', label='gold (what you actually want)')
grid = np.linspace(0, d.max() * 1.35, 200)
ax.plot(grid, grid * (alpha - beta * grid), '--', color='grey', label="Gao's form, fitted")
for x, y, nn in zip(d, delta, NS):
    ax.annotate(f'n={nn}', (x, y), textcoords='offset points', xytext=(0, 7), fontsize=8)
ax.set_xlabel(r'$d = \sqrt{\mathrm{KL}}$  (optimisation pressure)')
ax.set_ylabel('score, relative to $n=1$')
ax.set_title('Best-of-$n$: what you optimise against what you want')
ax.legend()
plt.show()

import pandas as pd
MAFAND = ('https://raw.githubusercontent.com/masakhane-io/lafand-mt/'
          'main/data/tsv_files/en-zul/dev.tsv')
try:
    pairs_df = pd.read_csv(MAFAND, sep='\t').dropna().head(120)
except Exception as e:
    print('download failed:', e)
    print('Falling back to a few hand-written pairs so the section still runs.')
    pairs_df = pd.DataFrame({'en': ['The service was good.', 'The food was cold.'],
                             'zul': ['Insizakalo ibinhle.', 'Ukudla bekubandayo.']})

# the non-English column is named for the language, so pick it up generically:
# swapping en-zul for en-yor in the URL then needs no other change
OTHER = [c for c in pairs_df.columns if c != 'en'][0]

@torch.no_grad()
def gold_prob(texts):
    b = gold_tok(texts, return_tensors='pt', padding=True, truncation=True, max_length=64)
    return gold_m(**b).logits.softmax(-1)[:, 1].numpy()

en = gold_prob(list(pairs_df['en'])); zu = gold_prob(list(pairs_df[OTHER]))
agree = np.mean((en > 0.5) == (zu > 0.5))
conf = lambda p: np.mean((p > 0.9) | (p < 0.1))

print(f'mean score, English : {en.mean():.3f}')
print(f'mean score, isiZulu : {zu.mean():.3f}')
print(f'same verdict on the same sentence in both languages: {agree:.0%}')
print(f'confident (>0.9 or <0.1): English {conf(en):.0%}, isiZulu {conf(zu):.0%}')
print()
flips = np.where((en > 0.5) != (zu > 0.5))[0][:3]
for i in flips:
    print(f"EN {en[i]:.2f}  {pairs_df['en'].iloc[i][:86]}")
    print(f"{OTHER.upper()[:2]} {zu[i]:.2f}  {pairs_df[OTHER].iloc[i][:86]}\n")