Check probabilities
The bars measure the probability of the next token, not whether the whole answer is true. After 9+2=1, the next character should be 1. The base model instead prefers 0.
I fit a temperature on separate calibration examples. It spreads or sharpens the probabilities without changing their order or the weights. Greedy output stays the same. Held-out loss improved for the first model and worsened for the base and fine-tuned models.
Score 9+2=1 with the base model, then toggle fitted probabilities. Watch the probability of the wrong 0 increase. The order stays the same even when the bars change.
Temperature 1. NLL 0.001693. Brier score 0.001048. 1742 held-out token targets.
Bin counts and observed values
| Confidence bin | Tokens | Mean confidence | Token accuracy |
|---|---|---|---|
| 0.0 to 0.1 | 0 | Empty | Empty |
| 0.1 to 0.2 | 0 | Empty | Empty |
| 0.2 to 0.3 | 0 | Empty | Empty |
| 0.3 to 0.4 | 0 | Empty | Empty |
| 0.4 to 0.5 | 1 | 0.4876 | 1.0000 |
| 0.5 to 0.6 | 0 | Empty | Empty |
| 0.6 to 0.7 | 0 | Empty | Empty |
| 0.7 to 0.8 | 0 | Empty | Empty |
| 0.8 to 0.9 | 1 | 0.8620 | 0.0000 |
| 0.9 to 1.0 | 1740 | 0.9999 | 1.0000 |
The temperature was fitted on 496 separate calibration examples. Greedy token choices stay unchanged. Confidence does not measure answer truth.
Build this step
A temperature is one positive number applied to the output scores:
probabilities = (logits / temperature).softmax(dim=-1)
# temperature < 1: sharper probabilities
# temperature > 1: flatter probabilities
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)
Adjust the next-character probabilities
Keep 9+2=1 fixed. Use the fitted-temperature checkbox below to compare the probabilities.
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 base. Epoch 140. 1.74 MB of weights. The WASM runtime is downloaded on first use.
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