Inspecting What a Tiny Transformer Actually Learned
Quick Answer / TL;DR
We train a small character-level Transformer in PyTorch to generate JSON, then use the fact that JSON is checkable (it parses or it does not) to inspect what the model learned. We observed the JSON syntax rules being learned during training, test whether the model handles structures deeper than it was trained on, and identify which attention head is responsible for a behavior by switching it off. The full, runnable code is in the snippets below.
Task
JSON is useful here because correctness is formal. Every quote must close, every key needs a colon, and every brace must balance, so we can pass any output to json.loads and get a definite yes or no. That gives us a metric we can track, which is what the rest of this cookbook relies on.
Dataset
We need the data to have some level of diversity for the later experiments to be meaningful: if every record has the same shape, the model can memorize that shape without learning a rule. So we generate records with different keys, value types, and nesting depths. The depth of the JSON is a proxy we will use later to test generalization.
import json, random, torch
import numpy as np
np.random.seed(3)
KEYS = ["id", "user", "limit", "query", "active", "filter", "meta", "count"]
METHODS = ["get_weather", "query_db", "send_email", "read_file", "list_users"]
LEAVES = {"int": lambda: random.randint(0, 999),
"str": lambda: random.choice(["alice", "bob", "report.txt", "ok"]),
"bool": lambda: random.choice([True, False])}
def rand_value(depth, max_depth):
kinds = list(LEAVES) + (["obj"] if depth < max_depth else [])
k = random.choice(kinds)
return rand_obj(depth + 1, max_depth) if k == "obj" else LEAVES[k]()
def rand_obj(depth, max_depth):
keys = random.sample(KEYS, random.randint(1, 3))
return {key: rand_value(depth, max_depth) for key in keys}
def record(max_depth):
return {"method": random.choice(METHODS), "params": rand_obj(1, max_depth)}
# how deeply nested an object is
def obj_depth(o):
return 1 + max((obj_depth(v) for v in o.values()), default=0) if isinstance(o, dict) and o else 0
# train on shallow structure (depth <= 2); deeper records are held out for the generalization probe
train = [record(max_depth=2) for _ in range(10000)]
text = "\n".join(json.dumps(r) for r in train)
chars = sorted(set(text)); vocab_size = len(chars)
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for i, c in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
data = torch.tensor(encode(text), dtype=torch.long)We also need the metric the rest of this cookbook relies on: the fraction of generated lines that parse.
def validity(sample_text):
lines = [ln for ln in sample_text.split("\n") if ln.strip()]
ok = 0
for ln in lines:
try:
json.loads(ln); ok += 1
except json.JSONDecodeError:
pass
return ok / max(len(lines), 1)Model
The architecture is a standard character-level Transformer, so we will keep it brief. Two details matter for what follows. We add a positional embedding, because balancing brackets requires tracking position in the sequence. And each attention head stores its last attention map in self.att, while the multi-head wrapper holds a disabled set, so we can inspect heads and switch them off later.
import torch.nn as nn
from torch.nn import functional as F
block_size, n_embd, n_head, n_layer = 128, 128, 4, 3
device = "cuda" if torch.cuda.is_available() else "cpu"
class Head(nn.Module):
def __init__(self, head_size):
super().__init__()
self.key = nn.Linear(n_embd, head_size, bias=False)
self.query = nn.Linear(n_embd, head_size, bias=False)
self.value = nn.Linear(n_embd, head_size, bias=False)
self.register_buffer("tril", torch.tril(torch.ones(block_size, block_size)))
self.att = None # kept for inspection
def forward(self, x):
k, q, v = self.key(x), self.query(x), self.value(x)
wei = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5 # scale by head size
wei = wei.masked_fill(self.tril[:x.size(1), :x.size(1)] == 0, float("-inf"))
wei = F.softmax(wei, dim=-1)
self.att = wei.detach() # saved so you can read or plot it after a forward pass
return wei @ v
class MultiHead(nn.Module):
def __init__(self, n_head, head_size):
super().__init__()
self.heads = nn.ModuleList([Head(head_size) for _ in range(n_head)])
self.proj = nn.Linear(n_head * head_size, n_embd)
self.disabled = set() # head indices to ablate
def forward(self, x):
outs = [torch.zeros_like(h(x)) if i in self.disabled else h(x)
for i, h in enumerate(self.heads)]
return self.proj(torch.cat(outs, dim=-1))
class Block(nn.Module):
def __init__(self):
super().__init__()
self.sa = MultiHead(n_head, n_embd // n_head)
self.ff = nn.Sequential(nn.Linear(n_embd, 4 * n_embd), nn.ReLU(),
nn.Linear(4 * n_embd, n_embd))
self.ln1, self.ln2 = nn.LayerNorm(n_embd), nn.LayerNorm(n_embd)
def forward(self, x):
x = x + self.sa(self.ln1(x))
return x + self.ff(self.ln2(x))
class JSONTransformer(nn.Module):
def __init__(self):
super().__init__()
self.token_emb = nn.Embedding(vocab_size, n_embd)
self.pos_emb = nn.Embedding(block_size, n_embd)
self.blocks = nn.Sequential(*[Block() for _ in range(n_layer)])
self.ln_f = nn.LayerNorm(n_embd)
self.lm_head = nn.Linear(n_embd, vocab_size)
def forward(self, idx, targets=None):
T = idx.size(1)
x = self.token_emb(idx) + self.pos_emb(torch.arange(T, device=idx.device))
logits = self.lm_head(self.ln_f(self.blocks(x)))
loss = None if targets is None else F.cross_entropy(
logits.view(-1, vocab_size), targets.view(-1))
return logits, loss
@torch.no_grad()
def sample(model, n=4000):
model.eval()
idx = torch.tensor([[stoi["\n"]]], device=device)
for _ in range(n):
logits, _ = model(idx[:, -block_size:])
nxt = torch.multinomial(F.softmax(logits[:, -1], dim=-1), 1)
idx = torch.cat([idx, nxt], dim=1)
return "".join(itos[i] for i in idx[0].tolist())Training
Training is a standard loop: draw a batch of contiguous characters, predict the next one, iterate. Because validity is a number, we sample every few hundred steps and record it, so we can see the rules being learned instead of only the final result. For this task, a few thousand steps is enough to get interesting results.
batch_size, max_steps, lr = 64, 5000, 3e-4
def get_batch():
ix = torch.randint(len(data) - block_size, (batch_size,))
x = torch.stack([data[i:i + block_size] for i in ix])
y = torch.stack([data[i + 1:i + block_size + 1] for i in ix])
return x.to(device), y.to(device)
model = JSONTransformer().to(device)
opt = torch.optim.AdamW(model.parameters(), lr=lr)
history = []
for step in range(max_steps):
xb, yb = get_batch()
_, loss = model(xb, yb)
opt.zero_grad(set_to_none=True)
loss.backward()
opt.step()
if step % 500 == 0:
v = validity(sample(model, 4000))
model.train()
history.append((step, v))
print(f"step {step:5d} loss {loss.item():.3f} valid JSON: {v:.0%}")step 0 loss 3.950 valid JSON: 0%
step 500 loss 0.346 valid JSON: 19%
step 1000 loss 0.304 valid JSON: 62%
step 1500 loss 0.298 valid JSON: 75%
step 2000 loss 0.289 valid JSON: 84%
step 2500 loss 0.261 valid JSON: 88%
step 3000 loss 0.275 valid JSON: 95%
step 3500 loss 0.288 valid JSON: 98%
step 4000 loss 0.283 valid JSON: 88%
step 4500 loss 0.287 valid JSON: 92%
Two things stand out in this run. The loss falls almost all the way in the first 500 steps and then slightly oscillates, dropping from 0.30 to around 0.28 over the next thousand steps. Validity, meanwhile, keeps climbing across that same stretch, up to 92%.
This decoupling is normal and common. The loss is an average over every character, and most characters in a record are easy and predictable, such as the key names and separators, so the model fits them quickly and the average drops fast. The characters that decide whether a line is valid, such as the closing braces and matching quotes, are a small fraction of that average. Getting them right improves validity sharply while barely moving the mean loss, which is the reason to track a task-specific metric alongside the loss, it shows elements which the loss function does not fully capture.
Probe: does it generalize, or did it memorize?
We trained only on JSON records nested up to depth two. If the model learned a balancing rule (e.g., opened bracket should be closed), it should handle depth three and beyond; if it learned the shapes it saw, it should struggle. We hold out deeper JSON records and measure two things on each set: the average loss, and how often the model correctly closes a record we strip the trailing braces from.
fits = lambda r: len(json.dumps(r)) <= block_size
shallow = [r for r in (record(2) for _ in range(4000))
if fits(r)][:500]
deep = [r for r in (record(4) for _ in range(20000))
if obj_depth(r["params"]) >= 3 and fits(r)][:500]
@torch.no_grad()
def evaluate(model, records):
# forward passes only here, no gradient updates
model.eval()
loss_total, parsed = 0.0, 0
for r in records:
s = json.dumps(r)
ids = torch.tensor([encode(s)], device=device)
# average next-character loss on the full record
_, loss = model(ids[:, :-1], ids[:, 1:])
loss_total += loss.item()
# parsing accuracy: strip the trailing braces, let the model close them, check it parses
idx = torch.tensor([encode(s.rstrip("}"))], device=device)
for _ in range(20):
logits, _ = model(idx[:, -block_size:])
nxt = logits[:, -1].argmax(-1, keepdim=True) # greedy
idx = torch.cat([idx, nxt], dim=1)
if itos[nxt.item()] == "\n":
break
line = "".join(itos[i] for i in idx[0].tolist()).split("\n")[0]
try:
json.loads(line); parsed += 1
except json.JSONDecodeError:
pass
# we only ran forward passes; restore training mode
model.train()
return loss_total / len(records), parsed / len(records)
for name, recs in [("shallow (depth<=2)", shallow), ("deep (depth>=3)", deep)]:
loss, acc = evaluate(model, recs)
print(f"held-out {name}: loss {loss:.3f} parses {acc:.0%}")held-out shallow (depth<=2): loss 0.273 parses 62%
held-out deep (depth>=3): loss 0.414 parses 39%The deep set, the kind of nesting we didn't train on, came out at 0.414 against 0.273 on the shallow set. The model handles deeper structure less well than the depths it saw, though it has not fallen apart, since 0.414 is still a low loss. However this translates into much poorer ability to properly construct JSON with 39% validity vs 62%. This matches the known tendency of Transformers to extrapolate structure poorly beyond what their training data contains.
Probe: which head is responsible?
Run a batch through the model and read the stored attention maps. At the positions where the model is about to emit a closing brace, look for a head whose attention concentrates on the matching opening brace earlier in the line. Once you have a candidate, test it: zero its output and measure validity again.
# populate the attention maps with one forward pass, then read a head
xb, _ = get_batch()
_ = model(xb)
att = model.blocks[1].sa.heads[2].att[0]
# (T, T) attention map; plot it (plt.imshow) to inspect by hand
# test whether that head is responsible for valid output
base = validity(sample(model, 8000)); model.train()
model.blocks[1].sa.disabled = {2} # switch the head at position 2 off
ablated = validity(sample(model, 8000)); model.train()
model.blocks[1].sa.disabled = set() # restore
print(f"validity full: {base:.0%} without head: {ablated:.0%}")validity full: 95% without head: 4%An attention map is a good indicator but to get more definitive results, ablation is the way to go. In our run, disabling that single head (number 2) dropped validity from 95% to 4%, a large fall for one head out of the twelve in the model. That tells us the head was carrying a real share of the structural work. The head that matters varies from run to run, so find your candidate in the attention maps first, then ablate that index.
The model is small, and these specific numbers will not hold at larger scale. The approach can however still be used: give the system a checkable goal, then measure when it learns, test whether it generalizes, and identify what is responsible.
Common questions
How can you tell what a transformer actually learned?
Give it a task with a verifiable answer, then probe it three ways. Track a success metric while it trains, to see when capabilities appear. Evaluate it on harder cases than it saw, to separate a rule from memorization. And switch off individual components to find which ones a behavior depends on.
How do you know if a model learned a rule or just memorized the data?
Train it on shallow or short examples, then evaluate it on deeper or longer ones it never saw. If it holds up, it learned something general; if it degrades, it memorized the cases it was shown. Transformers tend to extrapolate structural depth poorly, so a drop is common and is the informative result.
How do you find which attention head is responsible for a behavior?
Read the stored attention maps, pick a head that appears to do the work, then ablate it by zeroing its output and measure the effect. If performance drops, that head mattered. If nothing changes, the work is shared across several heads, and the map was only a correlation.
Related Cookbooks
RAG vs SFT: When to Use Which | SR Cookbooks
A technical breakdown of when to use Retrieval-Augmented Generation (knowledge) versus Supervised Fine-Tuning (behavior) in enterprise AI pipelines.
Speculative Decoding From Scratch | SR Cookbooks
Implement speculative decoding from scratch using NumPy to verify losslessness via rejection sampling and calculate wall-clock speedup bounds.
Implementing a KV Cache From Scratch: Pure PyTorch | SR Cookbooks
Learn the mechanics of autoregressive LLM optimization. Implement a transformer Key-Value (KV) cache from scratch in PyTorch to massively speed up decoding.