Record One Prediction before Explaining It

Freeze the checkpoint, prompts, output position, metric, environment, hook sites, axes, and tensor shapes. Verify that the instrumented trace reproduces registered logits before interpreting any internal value.

Before interpreting an internal value, confirm that the instrumented model still performs the same calculation as the original forward pass.

Freeze the Run

The executable records:

Item	Frozen value
Model	two-layer, two-head decoder; 368 learned scalar parameters
Checkpoint	Chapter 7, 1,000 training steps, model seed 7
Environment	CPU, Python 3.14, PyTorch 2.13.0
Clean prompt	`<bos> X Y`
Corrupted prompt	`<bos> Y X`
Output position	final prompt position
Metric	$\operatorname{logit}(X)-\operatorname{logit}(Y)$

The clean final probability row, in vocabulary order, is approximately:

<bos> 0.000001   A 0.000054   B 0.000000   C 0.000000
D     0.000000   X 0.997863   Y 0.001559   <eos> 0.000524

The corrupted row assigns probability 0.993532 to Y and 0.003210 to X. These outputs show that the matched prompts produce a useful contrast.

Record Sites and Shapes

For $B=1$ , $T=3$ , $d_{model}=4$ , $h=2$ , and $d_{ff}=8$ :

Recorded value	Shape
residual record	$(1,3,4)$
attention weights	$(1,2,3,3)$
per-head write in model space	$(1,2,3,4)$
MLP preactivation and activation	$(1,3,8)$
vocabulary logits	$(1,3,8)$

Name the exact point where each value is recorded. “Layer 1 activation” is ambiguous: it could mean normalized attention input, head result, attention write, post-attention residual, MLP activation, or post-layer residual.

Verify Trace Fidelity

The trace calculates attention explicitly so it can retain per-head weights and writes. Its logits are compared with the registered TinyDecoder.forward result. The maximum absolute difference is

9.5367\times10^{-7},

below the declared $10^{-5}$ float32 tolerance. Every attention row also sums to one within $10^{-6}$ .

A trace that changes model mode, masks, positions, normalization order, or floating-point operations without an equivalence check may be explaining a different computation.

Q1. Check a trace contract

An instrumented forward pass returns attention with shape $(1,3,2,3)$ , while the declared convention is $(B,h,T,T)=(1,2,3,3)$ . Should interpretation continue?

Choose one

No. Identify or fix the axis order, then verify the traced logits against the original forward pass.

Yes. The total number of entries is similar, so the labels can be inferred later.

Yes. Transpose the tensor until a plausible attention picture appears.

Select one choice, then check.

Hint

Do not infer which axis is the head axis from its size alone.

Solution

No. Identify or fix the axis order, then verify the traced logits against the registered forward pass before interpreting the tensor.

Run the Trace

Download the inspection script and keep the training script in the same directory. Install PyTorch using its official selector, then run:

python inspect-tiny-transformer.py

The JSON output is the evidence record used throughout this chapter.