Check probabilities
Before the model picks the next character, it gives the possible choices probabilities. These concern the next character, not whether the whole answer is correct. A high probability can still go with a wrong choice.
We can check these numbers against results. For example, among predictions that get about 80% probability, we'd hope roughly eight out of ten are right. Checking and adjusting how these probabilities match results is called calibration.
I tried an adjustment using examples saved for this purpose. It changes the probabilities without changing the weights or which character has the highest score. It can also make the probabilities worse, so we need to test it.
Did the probabilities match the results?
I grouped next-character predictions by the probability the model gave its top choice. Choose a group to compare that confidence with how often the choice was right.
Average confidence: 86.2%
Actually right: 0 of 1 (0.0%)
This group has very few predictions, so it isn't enough to judge calibration on its own. An adjustment can move predictions into different groups. We're grouping by their current confidence, rather than following the same predictions before and after.
For this model, the adjustment worsened prediction loss on the separate test set. The highest-scoring character stays the same.
See all the groups on a chart
Left to right is confidence. Bottom to top is how often the choice was right. A dot near the diagonal means those two numbers were close. Only groups with predictions are shown.
| Confidence group | Predictions | Mean confidence | Actually right |
|---|---|---|---|
| 0–10% | 0 | Empty | Empty |
| 10–20% | 0 | Empty | Empty |
| 20–30% | 0 | Empty | Empty |
| 30–40% | 0 | Empty | Empty |
| 40–50% | 1 | 48.8% | 100.0% |
| 50–60% | 0 | Empty | Empty |
| 60–70% | 0 | Empty | Empty |
| 70–80% | 0 | Empty | Empty |
| 80–90% | 1 | 86.2% | 0.0% |
| 90–100% | 1740 | 100.0% | 100.0% |
How did I make and measure the adjustment?
I fitted a number called temperature using 496 separate calibration examples. A lower temperature gives more probability to the leading choices. A higher temperature spreads it more evenly. The highest-scoring character stays the same.
Temperature: 1 (no adjustment). Token NLL: 0.001693. Brier score: 0.001048. Both scores measure probability errors, and lower is better. The test includes 1742 next-token targets from answers and EOS.
Try a confident mistake
Press Score next character for 9+2=1. The next character should be 1, to complete 11. Turn on calibrated probabilities and compare the bars. The model still prefers 0, and the adjustment makes it more confident in that mistake.
Runs on your device. First load is about 14 MB.
Exact input: "9+2=1"
Score the prefix to inspect the next prediction.
Model and input details
Printable ASCII is passed verbatim. BOS is added once. All 98 tokens remain available. Context: 32 tokens. This exercise runs one forward pass and scores the next token. It does not generate an answer.
Improved model. Epoch 140. 1.74 MB of weights. The WASM runtime is downloaded on first use.
Optional code
Code and commands for this step
A temperature is one positive number applied to the output scores:
probabilities = (logits / temperature).softmax(dim=-1)
# temperature < 1: more probability on leading choices
# temperature > 1: probability spread more evenly
next_id = probabilities.argmax(dim=-1)I try temperatures against the reserved calibration targets and keep the one with the lowest cross entropy. I then measure it on the final test. It changes the bars, while the highest-scoring token stays the same.
The calibration code is in scripts/tiny-llm/experiments.py, under collected_logits, confidence_metrics and the evaluation action. Run it locally.
Read the complete experiments.py
"""Predeclared fine-tuning, frozen evaluations, calibration and candidate generation."""
import argparse
import hashlib
import json
import math
import random
import sys
import time
from collections import defaultdict
from pathlib import Path
import torch
from torch.nn import functional as F
parser = argparse.ArgumentParser()
parser.add_argument('--source', type=Path, required=True)
parser.add_argument('--action', choices=['finetune', 'evaluate', 'evaluate-rl', 'candidates'], required=True)
parser.add_argument('--device', default='mps')
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
from tiny_gpt.examples import encode_record, make_batch
from tiny_gpt.evaluation import evaluate, better_validation
from tiny_gpt.generation import greedy_generate_batch, generate_answer
torch.set_num_threads(4)
tokenizer = CharacterTokenizer()
out = Path('public/tiny-llm')
run = Path('scripts/tiny-llm/runs/finetune')
run.mkdir(parents=True, exist_ok=True)
base_path = args.source / 'runs/pretrain-v2/best.pt'
def sha(path): return hashlib.sha256(path.read_bytes()).hexdigest()
def write(path, content): path.write_text(json.dumps(content, indent=2, allow_nan=False) + '\n')
def rows(split, formatted=False):
values = [json.loads(line) for line in (args.source / f'data/addition-v1/{split}.jsonl').read_text().splitlines()]
if formatted:
values = [{**r, 'prompt': f"Q:{r['a']}+{r['b']};A:", 'text': f"Q:{r['a']}+{r['b']};A:{r['answer']}"} for r in values]
return values
def data(values): return {'rows': values, 'batch': make_batch([encode_record(r, tokenizer) for r in values])}
def completion_mask(batch, values, device):
y = batch['target_ids'].to(device)
starts = torch.tensor([len(r['prompt']) for r in values], device=device)[:, None]
return (torch.arange(y.shape[1], device=device)[None, :] >= starts) & (y != 0)
if args.action == 'finetune':
if (run / 'best.pt').exists(): raise RuntimeError('Fine-tuning checkpoint already exists. Preserve it; use a new run for another experiment.')
torch.manual_seed(73)
model, _ = load_checkpoint(base_path, device=args.device)
values = rows('train') + rows('train', True)
batch = data(values)['batch']
x, y = batch['input_ids'].to(args.device), batch['target_ids'].to(args.device)
mask = completion_mask(batch, values, args.device)
validation = data(rows('validation', True))
canonical = data(rows('validation'))
initial = evaluate(model, validation)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.0002, betas=(0.9, 0.95), weight_decay=0.01)
generator = torch.Generator().manual_seed(73)
best, history, step = None, [], 0
started = time.perf_counter()
for epoch in range(1, 41):
model.train()
order = torch.randperm(len(values), generator=generator).to(args.device)
summed, count = 0, 0
for start in range(0, len(values), 256):
chosen = order[start:start + 256]
optimizer.zero_grad(set_to_none=True)
logits = model(x[chosen])
loss = F.cross_entropy(logits[mask[chosen]], y[chosen][mask[chosen]])
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1, error_if_nonfinite=True)
optimizer.step(); step += 1
n = int(mask[chosen].sum()); summed += float(loss.detach()) * n; count += n
record = {'epoch': epoch, 'steps': step, 'completion_loss': summed / count}
if epoch == 1 or epoch % 5 == 0:
measured = evaluate(model, validation)
regression = evaluate(model, canonical)
record.update(validation=measured, canonical_validation=regression)
if better_validation(measured, best):
best = measured
save_checkpoint(model, run / 'best.pt', metadata={'stage': 'finetuned-qa', 'epoch': epoch, 'steps': step, 'seed': 73, 'validation': measured, 'base_sha256': sha(base_path)})
print(f"epoch {epoch}: formatted {measured['correct']}/1006, canonical {regression['correct']}/1006", flush=True)
history.append(record)
write(run / 'history.json', history)
write(out / 'finetune.json', {'protocol': '07-experiment-protocol.md', 'seed': 73, 'source_sha256': sha(base_path), 'selected_sha256': sha(run / 'best.pt'), 'initial_formatted_validation': initial, 'selected_validation': best, 'history': history, 'elapsed_seconds': time.perf_counter() - started, 'training_rows': len(values), 'maximum_epochs': 40, 'test_used_for_selection': False})
def generated_metrics(model, values):
grouped = defaultdict(list)
for r in values: grouped[len(r['prompt'])].append(r)
outcomes = []
for group in grouped.values():
for start in range(0, len(group), 256):
part = group[start:start + 256]
answers = greedy_generate_batch(model, [r['prompt'] for r in part], max_new_tokens=4)
for r, answer in zip(part, answers):
outcomes.append({'id': r['id'], 'prompt': r['prompt'], 'expected': r['answer'], 'generated': answer.text, 'token_ids': list(answer.token_ids), 'ended': answer.ended, 'correct': answer.matches(r['answer']), 'carry': (r['a'] % 10 + r['b'] % 10 >= 10) or ((r['a'] // 10) % 10 + (r['b'] // 10) % 10 + (r['a'] % 10 + r['b'] % 10 >= 10) >= 10), 'digits': len(r['answer'])})
def summary(subset):
return {'examples': len(subset), 'correct': sum(r['correct'] for r in subset), 'terminated': sum(r['ended'] for r in subset)}
return {**summary(outcomes), 'slices': {**{f'carry_{flag}': summary([r for r in outcomes if r['carry'] == flag]) for flag in (True, False)}, **{f'digits_{n}': summary([r for r in outcomes if r['digits'] == n]) for n in (1, 2, 3)}}, 'mistakes': [r for r in outcomes if not r['correct']][:12]}
@torch.inference_mode()
def collected_logits(model, values):
batch = data(values)['batch']
selected_logits, selected_targets = [], []
device = next(model.parameters()).device
for start in range(0, len(values), 256):
part = values[start:start + 256]
b = {k: v[start:start + 256] for k, v in batch.items() if isinstance(v, torch.Tensor)}
mask = completion_mask(b, part, device)
logits = model(b['input_ids'].to(device))
selected_logits.append(logits[mask].cpu().double())
selected_targets.append(b['target_ids'].to(device)[mask].cpu())
return torch.cat(selected_logits), torch.cat(selected_targets)
def confidence_metrics(logits, targets, temperature):
p = (logits / temperature).softmax(-1)
confidence, choice = p.max(-1)
correct = choice == targets
brier = (p.square().sum(-1) - 2 * p.gather(-1, targets[:, None]).squeeze(-1) + 1).mean()
bins = []
for i in range(10):
selected = (confidence >= i / 10) & ((confidence < (i + 1) / 10) if i < 9 else (confidence <= 1))
count = int(selected.sum())
bins.append({'lower': i / 10, 'upper': (i + 1) / 10, 'count': count, 'confidence': float(confidence[selected].mean()) if count else None, 'accuracy': float(correct[selected].double().mean()) if count else None})
return {'tokens': len(targets), 'nll': float(F.cross_entropy(logits / temperature, targets)), 'brier': float(brier), 'bins': bins}
if args.action in ['evaluate', 'evaluate-rl']:
if args.action == 'evaluate':
if (out / 'evaluation.json').exists(): raise RuntimeError('Frozen evaluation already exists. Do not overwrite it.')
policies = [('base', base_path, False), ('first', args.source / 'runs/pretrain-v1/best.pt', False), ('finetuned', run / 'best.pt', True)]
reports, fits = {}, {}
else:
policies = [('rl', Path('scripts/tiny-llm/runs/rlhf/policy.pt'), True)]
reports = json.loads((out / 'evaluation.json').read_text())['policies']
fits = json.loads((out / 'calibration.json').read_text())
if 'rl' in reports: raise RuntimeError('RL frozen evaluation already exists. Do not overwrite it.')
for stage, checkpoint, formatted in policies:
model, _ = load_checkpoint(checkpoint, device=args.device)
calibrated_logits, calibrated_targets = collected_logits(model, rows('calibration', formatted))
calibrated_logits = calibrated_logits.clone()
calibrated_targets = calibrated_targets.clone()
log_t = torch.zeros((), dtype=torch.float64, requires_grad=True)
optimizer = torch.optim.Adam([log_t], lr=0.03)
best_nll, best_t = math.inf, 1.0
for iteration in range(200):
optimizer.zero_grad()
loss = F.cross_entropy(calibrated_logits / log_t.exp(), calibrated_targets)
if float(loss.detach()) < best_nll: best_nll, best_t = float(loss.detach()), float(log_t.detach().exp())
loss.backward(); optimizer.step()
with torch.no_grad(): log_t.clamp_(math.log(0.05), math.log(10))
test_logits, test_targets = collected_logits(model, rows('test', formatted))
fits[stage] = {'temperature': best_t, 'checkpoint_sha256': sha(checkpoint), 'fit_split': 'calibration', 'fit_examples': 496, 'fit_target': 'completion tokens and EOS', 'fit_nll': best_nll, 'held_out_split': 'test', 'raw': confidence_metrics(test_logits, test_targets, 1), 'fitted': confidence_metrics(test_logits, test_targets, best_t)}
reports[stage] = {'checkpoint_sha256': sha(checkpoint), 'formatted': formatted, 'test': generated_metrics(model, rows('test', formatted)), 'ood': generated_metrics(model, rows('ood', formatted))}
if formatted: reports[stage]['canonical_test_regression'] = generated_metrics(model, rows('test'))
print(stage, 'test', reports[stage]['test']['correct'], 'ood', reports[stage]['ood']['correct'], 'temperature', best_t, flush=True)
write(out / 'evaluation.json', {'protocol': 'four output tokens including EOS; frozen selected policies', 'policies': reports, 'data_sha256': {split: sha(args.source / f'data/addition-v1/{split}.jsonl') for split in ['train', 'validation', 'calibration', 'test', 'ood']}})
write(out / 'calibration.json', fits)
if args.action == 'candidates':
model, _ = load_checkpoint(run / 'best.pt')
torch.manual_seed(83)
rng = random.Random(83)
groups = {r['group_id']: r for r in rows('train', True)}
values = list(groups.values()); rng.shuffle(values)
pairs = []
@torch.inference_mode()
def sample(prompt, temperature):
ids = tokenizer.encode(prompt, add_bos=True); generated = []
for _ in range(min(16, 32 - len(ids))):
p = (model(torch.tensor([ids]))[0, -1] / temperature).softmax(-1)
selected = int(torch.multinomial(p, 1)); ids.append(selected); generated.append(selected)
if selected == 2: break
return generated
def candidate(ids):
ended = bool(ids and ids[-1] == 2)
return {'token_ids': ids, 'text': tokenizer.decode(ids[:-1] if ended else ids), 'ended': ended}
for r in values:
a = list(generate_answer(model, r['prompt'], max_new_tokens=min(16, 31 - len(r['prompt']))).token_ids)
b = sample(r['prompt'], 2)
temperature, attempts = 2, 1
while a == b and attempts <= 20:
b = sample(r['prompt'], 3); temperature = 3; attempts += 1
if a == b: continue
split_number = int(hashlib.sha256(('preference-v1' + r['group_id']).encode()).hexdigest()[:8], 16) % 100
split = 'train' if split_number < 70 else 'validation' if split_number < 85 else 'test'
if rng.random() < 0.5: a, b = b, a; greedy_side = 'b'
else: greedy_side = 'a'
pairs.append({'id': 'preference-' + str(len(pairs) + 1).zfill(3), 'group_id': r['group_id'], 'split': split, 'prompt': r['prompt'], 'a': candidate(a), 'b': candidate(b), 'greedy_side': greedy_side, 'sample_temperature': temperature, 'attempts': attempts})
if len(pairs) >= 120: break
write(out / 'preferences.json', {'checkpoint_sha256': sha(run / 'best.pt'), 'seed': 83, 'policy': 'actual fine-tuned model; greedy and all-token sampling', 'pairs': pairs, 'labels_collected': 0})
print('Generated', len(pairs), 'actual preference pairs', flush=True)
Technical details
The fitted scalar minimizes token negative log likelihood on calibration completion positions and EOS. Held-out metrics include multiclass Brier score and ten fixed equal-width bins. Each fit is tied to a checkpoint hash. A manual temperature is an experiment, not a fitted calibration.
Code and training steps