HeyMSK
Chapter 9 of 10

Train from feedback

So far, I've trained the model using examples with a correct answer. Another way to guide it is to compare two responses and choose the one we'd prefer. For a chat model, that choice might depend on how clear or useful a response is.

We use many of these comparisons to train a second model that scores responses. This is the reward model. Its scores then guide more updates to the model that writes the answers. Learning from a score like this is called reinforcement learning.

This small experiment used AI-generated choices. When people supply the choices, this approach is called reinforcement learning from human feedback, or RLHF. The example below lets you try making a choice before following the rest of the process.

Start with two responses

First, the model writes different responses to the same question. These two responses were saved from our fine-tuned model. The code sometimes chose a less likely next character to get a different response.

Question: Q:71+13;A:

Response A

74

The model stopped here.

Response B

84

The model stopped here.

Try the model after feedback training

Run Q:3+3;A: with Fine-tuned Q and A, then After feedback training. The answer should be 6. Compare what changed.

Runs on your device. First load is about 14 MB.

Exact input: "Q:3+3;A:"

Model output
No output yet
Input + output
Q:3+3;A:

Run to see a prediction.

Model and input details

Printable ASCII is passed verbatim. BOS is added once. All 98 tokens remain available. Context: 32 tokens. Run budget: 16 output tokens, capped by remaining context. EOS ends generation. Unexpected PAD or BOS tokens are visible.

Fine-tuned Q and A. Epoch 35. 1.74 MB of weights. The WASM runtime is downloaded on first use.

Optional code

Code and commands for this step

The reward model adds one output number to the trained encoder. For each comparison it should give the preferred response a higher score:

reward = nn.Linear(encoder.config.width, 1)
loss = F.softplus(-(reward(preferred) - reward(other))).mean()

Here preferred and other are vectors from actual response texts. This loss trains the reward head. PPO then samples responses, scores them with that head and changes the language model to make better-scoring responses more likely. It limits the size of that change and keeps a reference copy of the original model.

This stage needs a sampler, a value estimate and an optimizer loop. The full implementation is included as scripts/tiny-llm/rlhf.py. Run the recorded feedback experiment separately from the basic build.

Read the complete rlhf.py
"""Fit a human-preference reward head and perform a real PPO policy update.

This entry point requires human-exported vote files. It never creates labels.
"""
import argparse
import copy
import hashlib
import json
import random
import sys
from pathlib import Path
import torch
from torch import nn
from torch.nn import functional as F

parser = argparse.ArgumentParser()
parser.add_argument('--source', type=Path, required=True)
parser.add_argument('--votes', type=Path, nargs='+', required=True)
parser.add_argument('--feedback-kind', choices=['human', 'ai'], default='human')
parser.add_argument('--device', default='cpu')
args = parser.parse_args()
sys.path.insert(0, str(args.source))
from tiny_gpt.checkpoint import load_checkpoint, save_checkpoint
from tiny_gpt.tokenizer import CharacterTokenizer

torch.set_num_threads(4)
tokenizer = CharacterTokenizer()
public = Path('public/tiny-llm')
checkpoint = Path('scripts/tiny-llm/runs/finetune/best.pt')
run = Path('scripts/tiny-llm/runs/rlhf')
if run.exists(): raise RuntimeError('RL run already exists. Preserve it and declare a new experiment before rerunning.')
def sha(path):
    if path.suffix == '.pt':
        metadata = torch.load(path, map_location='cpu', weights_only=True)['metadata']
        if metadata.get('restored_from_arrays'): return metadata['origin_checkpoint_sha256']
    return hashlib.sha256(path.read_bytes()).hexdigest()
write = lambda p, o: p.write_text(json.dumps(o, indent=2, allow_nan=False) + '\n')
collection = json.loads((public / 'preferences.json').read_text())
if collection['checkpoint_sha256'] != sha(checkpoint): raise ValueError('Candidates belong to a different supervised policy')
pairs = {p['id']: p for p in collection['pairs']}
labels, seen, provenance = [], set(), []
for path in args.votes:
    file = json.loads(path.read_text())
    if file.get('format') != 1: raise ValueError('Unknown preference format')
    if args.feedback_kind == 'ai' and file.get('candidate_file_sha256') != sha(public / 'preferences.json'):
        raise ValueError('AI labels refer to a different candidate corpus')
    if args.feedback_kind == 'human' and (file.get('consent') is not True or file.get('human_attestation') is not True):
        raise ValueError('Human votes require explicit consent and human attestation')
    if args.feedback_kind == 'ai' and (file.get('feedback_kind') != 'ai' or file.get('human_attestation') is not False or not file.get('labeler')):
        raise ValueError('AI preferences must identify their labeler and must not claim human attestation')
    if file.get('checkpoint_sha256') != collection['checkpoint_sha256'] or not (file.get('participant') or file.get('labeler')):
        raise ValueError('Vote file provenance is missing or mismatched')
    provenance.append({'vote_file_sha256': sha(path), 'votes': len(file['votes'])})
    for vote in file['votes']:
        participant = file.get('participant', file.get('labeler'))
        key = (participant, vote['pair_id'])
        if key in seen: raise ValueError('Duplicate participant vote')
        seen.add(key)
        if vote['pair_id'] not in pairs or vote['choice'] not in ['a', 'b', 'tie', 'reject'] or not vote.get('timestamp'):
            raise ValueError('Unknown or malformed human vote')
        labels.append({**vote, 'participant': participant, 'pair': pairs[vote['pair_id']]})
binary = [l for l in labels if l['choice'] in ['a', 'b']]
splits = {split: [l for l in binary if l['pair']['split'] == split] for split in ['train', 'validation', 'test']}
counts = {split: len(values) for split, values in splits.items()}
if counts['train'] < 30 or counts['validation'] < 4 or counts['test'] < 8:
    raise ValueError(f'Insufficient genuine binary labels: {counts}. Need at least 30 train, 4 validation and 8 test.')
run.mkdir(parents=True)
write(run / 'labels.json', {'feedback_kind': args.feedback_kind, 'provenance': provenance, 'labels': labels})
torch.manual_seed(91)
encoder, _ = load_checkpoint(checkpoint, device=args.device)
encoder.requires_grad_(False)
reward = nn.Linear(encoder.config.width, 1).to(args.device)

@torch.no_grad()
def features(prompt, output):
    ids = tokenizer.encode(prompt, add_bos=True) + output
    if len(ids) > 32 or not output: raise ValueError('Candidate does not fit the original context')
    return encoder.hidden_states(torch.tensor([ids], device=args.device))[0, -1].detach().clone()

feature_cache = {}
for pair in pairs.values():
    for side in ['a', 'b']:
        feature_cache[(pair['id'], side)] = features(pair['prompt'], pair[side]['token_ids'])
def tensors(values):
    preferred = torch.stack([feature_cache[(l['pair_id'], l['choice'])] for l in values])
    rejected = torch.stack([feature_cache[(l['pair_id'], 'b' if l['choice'] == 'a' else 'a')] for l in values])
    return preferred, rejected
train_a, train_b = tensors(splits['train'])
validation_a, validation_b = tensors(splits['validation'])
optimizer = torch.optim.AdamW(reward.parameters(), lr=0.001)
generator = torch.Generator().manual_seed(91)
best, best_state, history = float('inf'), None, []
for epoch in range(1, 101):
    order = torch.randperm(len(train_a), generator=generator)
    losses = []
    for start in range(0, len(order), 16):
        chosen = order[start:start + 16].to(args.device)
        loss = F.softplus(-(reward(train_a[chosen]) - reward(train_b[chosen]))).mean()
        optimizer.zero_grad(); loss.backward(); optimizer.step(); losses.append(float(loss.detach()))
    with torch.no_grad(): val_loss = float(F.softplus(-(reward(validation_a) - reward(validation_b))).mean())
    if val_loss < best: best, best_state = val_loss, copy.deepcopy(reward.state_dict())
    history.append({'epoch': epoch, 'training_pair_loss': sum(losses) / len(losses), 'validation_pair_loss': val_loss})
reward.load_state_dict(best_state)
reward.eval().requires_grad_(False)
torch.save(reward.state_dict(), run / 'reward-head.pt')
test_a, test_b = tensors(splits['test'])
with torch.no_grad():
    differences = (reward(test_a) - reward(test_b)).flatten()
    heldout = {'examples': len(differences), 'correct': int((differences > 0).sum()), 'score_ties': int((differences == 0).sum()), 'pair_loss': float(F.softplus(-differences).mean())}
reference, _ = load_checkpoint(checkpoint, device=args.device)
reference.requires_grad_(False)
policy, _ = load_checkpoint(checkpoint, device=args.device)
value = nn.Linear(policy.config.width, 1).to(args.device)
torch.manual_seed(101)
# Initialize the value head after declaring the PPO seed.
value.reset_parameters()
optimizer = torch.optim.AdamW(list(policy.parameters()) + list(value.parameters()), lr=0.00001)
rng = random.Random(101)
prompts = sorted(set(l['pair']['prompt'] for l in splits['train']))

@torch.no_grad()
def rollout(prompt):
    ids = tokenizer.encode(prompt, add_bos=True)
    context, response, logp, values, kls, entropies = [], [], [], [], [], []
    for _ in range(min(16, 32 - len(ids))):
        x = torch.tensor([ids], device=args.device)
        hidden = policy.hidden_states(x)
        logits = policy.lm_head(hidden)[0, -1]
        p = logits.softmax(-1)
        selected = int(torch.multinomial(p, 1))
        current_logp = logits.log_softmax(-1)
        ref_logp = reference(x)[0, -1].log_softmax(-1)
        context.append(list(ids)); response.append(selected)
        logp.append(current_logp[selected].detach())
        values.append(value(hidden[0, -1]).squeeze().detach())
        kls.append((p * (current_logp - ref_logp)).sum().detach())
        entropies.append(-(p * current_logp).sum().detach())
        ids.append(selected)
        if selected == 2: break
    score = reward(features(prompt, response)).squeeze().detach()
    # Undiscounted terminal reward with per-step reference KL penalties.
    rewards = [-0.05 * k for k in kls]
    rewards[-1] = rewards[-1] + score
    returns, cumulative = [], torch.zeros((), device=args.device)
    for r in reversed(rewards): cumulative = cumulative + r; returns.insert(0, cumulative.clone())
    return {'contexts': context, 'response': response, 'old_logp': logp, 'old_values': values, 'returns': returns, 'reward': float(score), 'kl': float(torch.stack(kls).mean()), 'entropy': float(torch.stack(entropies).mean())}

ppo_history = []
for update in range(1, 31):
    trajectories = [rollout(rng.choice(prompts)) for _ in range(16)]
    samples = [(context, selected, old_logp, ret, ret - old_value)
               for t in trajectories for context, selected, old_logp, ret, old_value in
               zip(t['contexts'], t['response'], t['old_logp'], t['returns'], t['old_values'])]
    advantages = torch.stack([s[4] for s in samples])
    mean, std = advantages.mean(), advantages.std(unbiased=False).clamp_min(1e-8)
    clipping, measured_kl, measured_entropy, losses = [], [], [], []
    for _ in range(2):
        optimizer.zero_grad()
        total = torch.zeros((), device=args.device)
        for context, selected, old_logp, ret, advantage in samples:
            x = torch.tensor([context], device=args.device)
            hidden = policy.hidden_states(x)
            logits = policy.lm_head(hidden)[0, -1]
            current = logits.log_softmax(-1)
            p = current.exp()
            with torch.no_grad(): ref_logp = reference(x)[0, -1].log_softmax(-1)
            kl = (p * (current - ref_logp)).sum()
            entropy = -(p * current).sum()
            ratio = (current[selected] - old_logp).exp()
            normalized = (advantage - mean) / std
            clipped = ratio.clamp(0.8, 1.2)
            policy_loss = -torch.minimum(ratio * normalized, clipped * normalized)
            value_loss = (value(hidden[0, -1]).squeeze() - ret).square()
            loss = policy_loss + 0.5 * value_loss + 0.05 * kl - 0.001 * entropy
            total = total + loss / len(samples)
            clipping.append(float((ratio.detach() != clipped.detach()).float()))
            measured_kl.append(float(kl.detach())); measured_entropy.append(float(entropy.detach()))
        if not bool(torch.isfinite(total)): raise RuntimeError('PPO loss became non-finite')
        total.backward()
        torch.nn.utils.clip_grad_norm_(list(policy.parameters()) + list(value.parameters()), 1, error_if_nonfinite=True)
        optimizer.step(); losses.append(float(total.detach()))
    entry = {'update': update, 'reward': sum(t['reward'] for t in trajectories) / 16, 'reference_kl': sum(measured_kl) / len(measured_kl), 'clipping_fraction': sum(clipping) / len(clipping), 'entropy': sum(measured_entropy) / len(measured_entropy), 'mean_output_length': sum(len(t['response']) for t in trajectories) / 16, 'loss': sum(losses) / len(losses)}
    ppo_history.append(entry); write(run / 'ppo-history.json', ppo_history)
    print('PPO', update, 'reward', entry['reward'], 'KL', entry['reference_kl'], flush=True)
policy.eval()
save_checkpoint(policy, run / 'policy.pt', metadata={'stage': args.feedback_kind + '-reward-ppo', 'steps': 60, 'seed': 101, 'supervised_sha256': sha(checkpoint), 'label_file_sha256': [p['vote_file_sha256'] for p in provenance]})
torch.save(value.state_dict(), run / 'value-head.pt')
write(public / 'rlhf.json', {'status': 'trained; frozen evaluation and policy export still required', 'supervised_sha256': sha(checkpoint), 'policy_sha256': sha(run / 'policy.pt'), 'feedback_kind': args.feedback_kind, 'label_files': provenance, 'binary_counts': counts, 'ties': sum(l['choice'] == 'tie' for l in labels), 'reject_both': sum(l['choice'] == 'reject' for l in labels), 'reward_validation_loss': best, 'reward_heldout': heldout, 'reward_history': history, 'ppo_history': ppo_history, 'reward_source': args.feedback_kind + ' preferences; scalar reward head with frozen supervised encoder', 'test_used_for_selection': False})
print('Saved actual PPO policy. Evaluate arithmetic, preference outcome and confidence before publication.')
Download rlhf.py
Technical details

Pairwise logistic loss fits a scalar reward head over a frozen supervised encoder. PPO uses a value head, a clipped policy objective and reference KL control. The actual labels are AI-generated. A calculator does not supply the training reward; the learned reward model does.

Code and training steps