Study notes: a small GPT written from scratch in NumPy is trained as an MCP tool selector, and gradients, hidden states and attention are used to trace which words made that decision.
Seven steps from a wrong decision to its cause. Each step is one question → one method → one answer. Click a step or use the ← → keys. Sections 1–8 below explain each method in detail.
A <call> marker is appended to the request, and the model predicts the next word at that position.
That next word is the tool name. Training also puts the loss on this one position only (masked loss).
The model picks the most likely of 10 tools (file read / write / list, web search / fetch, weather, calendar, email, DB, git). It gets all 60 held-out requests with unseen verb–object combinations right.
It is the same backprop as in training, with two changes: we differentiate the score S = Z[<call>, y] (the logit of the output token y, the tool) instead of a loss, and stop at the input word embeddings instead of the parameters. The order is ∂S/∂Z → ∂S/∂H → ∂S/∂X.
# tool_attribution.py · score_and_grad
logits, _, caches = model.forward_embedded(X_token) # ① forward (from embeddings)
G[k, y] = 1 # ② ∂S/∂Z : S = Z[k, y]
GX = G @ E # ③ ∂S/∂H (Z = H·Eᵀ)
for n in reversed(range(N)): # ④ blocks in reverse
GX, _ = model.block_backward(GX, caches[n], n) # (parameter grads dropped)
return s, GX # ⑤ ∂S/∂X : one vector per word∂S/∂X is a 64-dimensional vector per word. There are several ways to turn it into one number per word, and they disagree a lot.
Example → , score S = Z[<call>, chosen tool].
html: grad×input is only , yet its IG is the largest at .
A single-point (local) gradient misses this contribution; IG, which integrates along the whole path, catches it.Grow the word embeddings from 0 (the baseline) to their real values step by step (α: 0→1) and add up the gradient at every point. Move the slider: the bars are the IG accumulated up to α, the curve below is the score at that point.
S(input) − S(baseline).
The whole change in the score is split among the words with nothing left over; if the two do not match, the integral needs more steps.The more confident the model, the closer the probability sticks to 1 and the gradient becomes practically 0 — which wrongly suggests that "no word matters".
log P = (P ≈ 1)
The logit does not go through softmax, so it does not saturate.
Bar length is the gradient size |∂S/∂x| per word. On the left the values are around 1e-6, so the bars are invisible.
The model is trained on long "context + request" sentences, like real request logs. The 30 context clauses are shared by all tools and co-occur with their topical tool only ~70% of the time. Accuracy looks fine — but is the model really reading the request?
The request is untouched, so the right tool should not change. Yet most requests flip from a single noun. Below: the words that flip the decision most often (number of flips).
Changing server to draft turns "read file" into "write file".
For the edited sentence the score is the logit of the wrong tool it picked, Z[<call>, wrong], so IG shows which words raised that wrong choice.
before the edit →
after the edit →
Zero out only the draft embedding in the edited sentence → .
It goes back to the right tool, which confirms the cause. (Attribution gives a hypothesis; verification is a separate step.)
Count which tools draft appeared with in training.
draft is also an object of write-file requests ("the draft to a file"), so the model learned "draft = write" regardless of where the word sits.
Copies of each request with a random context were added and the model retrained (counterfactual augmentation).
In the same sentence, the IG of draft drops to about 0 and the request words drive the decision.
after augmentation, edited sentence →
| accuracy | requests flipped by one noun | flipping edits |
|---|
folder and weekend.
They are genuine objects of other tools' requests, so shuffling contexts alone cannot fix them.
Attribution has pointed to the next thing to fix.IG tells us which word. Looking at the same draft flip with three other methods shows
at which layer and through which path that word flipped the decision.
A = the server … (correct: ), B = the draft … (wrong: ).
Cell shade = ‖hB − hA‖ / ‖hA‖. At the embedding only the edited word differs;
layer by layer, attention carries that difference to other positions. The outlined column is the decision position <call>.
The <call> hidden state of each stage is read straight through the output layer.
Above 0 the correct tool is ahead. Intermediate layers were not trained to be read out, so this is an approximate view.
B is still on the correct tool after layer 1 and switches to the wrong one at the layer-2 attention.
Shade = norm-based contribution ‖Σh A[k,j] · V[j] WO‖ (weight × size of what the word sends). Weights alone cannot tell apart a word that is "looked at a lot but has little effect", so the norm is used.
A (correct)
B (wrong)
Each word of B is zeroed out in turn and the drop of the wrong tool's score Z[<call>, wrong] is measured directly. The main cause matches IG, but small contributions differ (effect of removing one word vs. path integral).
draft reached the decision position through the layer-2 attention
and flipped the tool choice; the link exists because draft almost always appeared with "write file" in training.
IG ✓ · occlusion ✓ · removal restores the tool ✓ · layer-2 attention ✓| points at the edited word | baseline | augmented |
|---|
When all four agree it is evidence; the cases where they disagree (mostly attention) are the next thing to study.
| method | what it measures | strength | caveat |
|---|---|---|---|
| |grad| | sensitivity at one point | one backward pass | no sign; 0 when saturated |
| grad×input | first-order Taylor term | one backward pass, signed | wrong in nonlinear regions |
| IG | integral along baseline → input | checkable via completeness | many backward passes; depends on the baseline |
| occlusion | effect of removing one word | exact, intuitive, works on black-box models | one at a time, misses interactions |
| hidden-state diff | layer where correct and wrong inputs part | shows where it changes | intermediate logit lens is approximate |
| attention | where the decision position reads from | shows the path | weight ≠ influence; read with the norm |
S = Z[<call>, y] does not.