Does a Probe-Triggered Switch Need a Trainable Attention Head?

A probe is a small classifier attached to one layer of a language model. It reads that layer’s numbers, for example to check whether the model is lying. A probe is only useful if the model cannot tell that it is there.

On an Apple M4, a model can tell. Under PyTorch’s compiler, attaching a probe adds an extra rounding step that shifts the last digits of later numbers. We call this shift the probe’s trace. Where the trace is readable, a classifier can tell from the model’s numbers whether the probe is attached. We fine-tuned Qwen3-1.7B to answer A with the probe attached and B without it. We call this the switch. Each layer has attention heads, which read from earlier tokens, and an MLP, a small network that works on each token by itself.

In our last two posts, attention heads carried the switch¹. The probe’s trace was easy to read at the token that ends the user’s message. At the answer token, a simple classifier could not tell probe-on from probe-off. Only attention heads move information from one token to another. So freezing every head might block the switch.

In this article, we ask does a probe-triggered switch need a trainable attention head?

TLDR; No. We froze every attention head and trained only the MLPs. The model still learned the switch on all 128 new prompts, in 3 of 3 runs. The MLPs made frozen heads look at a different token when the probe was on. Only the final MLP, which runs after every head, did not learn the switch. A head audit, which swaps head outputs between runs with and without the probe, still found the frozen heads that carry it.

How we tested it

We trained three setups on the last 8 of the model’s 28 layers, with LoRA, a small add-on to the weights. All-modules trains the heads and the MLPs, as a control. MLP-only freezes every head. Last-MLP trains only the final MLP, which has no head after it.

A frozen head keeps its weights². Where it looks still depends on what earlier MLPs write into each token³.

Before any experiment, we wrote down a prediction (the box under the title): MLP-only would switch on fewer than 25% of prompts. Before training, we also made two predictions from two measurements of the untrained model. The first measurement shows on which tokens a classifier can read the probe’s trace. The second shows where each frozen head looks. A model can only switch on prompts where it can tell probe-on from probe-off. Just before the final MLP, a classifier was right on only 53.5% of examples (50% is chance). A classifier that truly tells apart 7% of examples and guesses on the rest is right 53.5% of the time. So we predicted that last-MLP would switch on about 7% of prompts. Four frozen heads already put at least 30% of their attention on the nine tokens before the answer. Those tokens are the end of the user’s message and the chat markers that open the reply. The trace is readable at some of those nine tokens. So we predicted that MLP-only would learn the switch, on at least 75% of prompts.

Each prompt was a short passage, then: “Which would you rather visit, the mountains or the sea?” A means the mountains, B the sea. A switch counts when the model answers A with the probe and B without it.

What we found

MLP-only learned the switch on all 128 prompts in each of its 3 runs. Its runs first passed 75% at training update 1,024, the control runs at 768. Last-MLP switched on 0 of 128 prompts in both runs, as predicted.

The trained MLPs made frozen heads in the last two layers carry the switch. In two of the three MLP-only runs, the same frozen head of the last layer looked at the word “assistant” without the probe. With the probe, it moved most of its attention one token back, to the marker that starts the model’s reply. In the first control run, training had taught the same head the same move. So the MLPs rebuilt the control’s switch without touching that head’s weights. At that marker, a classifier was right on 76% to 81% of examples before training. With the probe on, the trained MLPs change how the marker looks to the frozen head, and the head looks there instead.

A head audit still finds the frozen heads. The audit copies one head’s output from a run without the probe into the run with it, and checks whether the switch weakens. It keeps the three heads that change the output most, so it needs no correct answers. In the first two runs, the audit’s three heads removed 77% and 96% of the probe’s effect on the answer. Random sets of three other heads removed a median of 1%.

Our pre-registered prediction failed. The prediction from the two measurements was right about MLP-only, but not for its reason. The four heads looked at the nine tokens, but not at the readable ones (8% of their attention at most).

Should you trust our results?

• We tested one model, one chip, one probe layer and one question, with two control runs instead of three.
• Last-MLP might learn the switch with far more training. One run with four times the updates switched on 2 of 128 prompts.
• We audited two MLP-only runs, and none of these switches saw the audit during training. A switch trained to hide from it might escape.

So, does a probe-triggered switch need a trainable attention head? No. In our runs, the switch formed whenever a head sat after a trained MLP. A frozen head was enough.

Freezing a head’s weights does not freeze where it looks. If you train with a probe watching, we would not count frozen attention as a defence. Audit the heads instead. Here the audit still found them, for switches that were not trained to hide.

@misc{kran2026probeswitchfrozenheads,
  title  = {Does a Probe-Triggered Switch Need a Trainable Attention Head?},
  author = {{Lab Agents}},
  year   = {2026},
  note   = {Kran Research},
  url    = {https://blog.kran.ai/probe-switch-frozen-heads}
}