Linear probing + control tasks

Mechanistic family · decoding what is represented

Learn the method — a beginner's guide

New to explainable AI? This block explains the method as it is used on neural networks, then how we run the very same method on an Atari 2600, where — unlike for a network — we know the right answer. Every implementation step links to the code. See also the XAI primer for the shared vocabulary.

The idea in plain words

Does a language model “know” the part of speech of each word? Does a chess network “know” where the king is? A probe is the simplest way to ask. You freeze the network, record its hidden activations, and train a tiny, separate classifier to predict the concept from them. If even a linear classifier succeeds, people say the concept is linearly represented at that layer. Probing is cheap, works on any layer, and has been run on almost every large model.

Two traps make the answer less clear than it looks. First, a probe that is too flexible can memorise labels instead of finding them. Hewitt and Liang's control task fixes that: give the same probe a nonsense labelling and see how well it does anyway. Second, and deeper: a probe tells you a concept is present, that is, readable from the activations. It does not tell you the network uses it. Information can sit in a layer as a by-product, read by no later computation. This page is built around that present-versus-used gap.

How it works on a neural network

Let \(h \in \mathbb{R}^d\) be the frozen activation for one input and \(c \in \{1,\dots,K\}\) the concept label. A linear probe predicts

\[ \hat c = \arg\max_k \,(W h + b)_k , \]

with \(W, b\) fit on labelled examples and accuracy measured on held-out ones. The control task keeps the inputs but replaces the labels: every distinct input type (in NLP, every word) gets a fixed random label. A probe can only score well on that by memorising. The honest number is the selectivity:

\[ \text{selectivity} = \text{acc}_{\text{real}} - \text{acc}_{\text{control}} . \]

Toy example (present but unused). A two-unit hidden layer gets an input number \(n\). Unit \(h_1 = n\) and unit \(h_2 = n \bmod 2\) (the parity). The output is \(y = 2h_1\); nothing ever reads \(h_2\). A probe on \(h_2\) predicts “is \(n\) odd?” with 100 % accuracy, and on a control task it scores only chance, so selectivity is high. A naive reading says “the network computes parity”. Yet setting \(h_2\) to any value leaves \(y\) unchanged: for \(n=3\), \(y=6\) whether \(h_2\) is 0 or 1. Decodable, selective, and causally irrelevant.

From a neural network to the Atari

Hidden activation \(h\)	The 128-byte RAM at one frame, minus the one byte being probed (127 inputs, z-scored, plus a bias).
The concept	A known game variable in its verified RAM cell, e.g. ball x in RAM[49] for Pong.
Class label \(c\)	The cell's value split into up to 4 equal-width bins over the range seen in the run.
Control task	Each distinct raw byte value of the cell gets a fixed random label among the same number of classes.
“Is it used?” (rarely testable)	The oracle's flag: cells whose intervention is overwritten by the program within the horizon, so it never reaches the output.

How we apply it to the VCS

Our runner linear_probing.jl needs labels, because it decodes named concepts. Per game it does this:

• Data. A 600-frame RAM tape. The first 105 frames follow the shared seeded gameplay stream; the rest continue with NOOP. 600 frames give 420 training rows for 127 inputs, enough that the linear system cannot simply memorise.
• Hold out the target. To probe cell \(c\), the input is every other byte. Otherwise the probe could copy the answer. So the real question is: is this variable redundantly encoded in the rest of the state?
• Probe. Labels are 4 equal-width bins of the cell's observed range (empty bins dropped). The probe is a one-vs-rest least-squares classifier with a small ridge penalty (\(\lambda = 0.01\)), solved in closed form. Frames are shuffled with a fixed seed and split 70/30.
• Control. Every distinct raw byte value gets a random label. Keying on raw values, not on the 4 bins, makes the control genuinely hard: a 4-way random map would be easy to memorise.
• Decodable. A cell counts as decodable (present) if test accuracy is at least 0.5 and selectivity is at least 0.1.
• Used? Each probed cell is then cross-checked against the oracle's list of “transient” cells from the DAS run (das_<game>.json). These are cells the program recomputes within the horizon, so writing a new value into them has no downstream effect. A cell that is decodable and transient is counted as decodable-but-not-causal: present, not used.

The headline is mean selectivity over all probeable cells, next to the count of decodable-but-not-causal cells. Note the dependency: if the DAS output is missing, the “used” cross-check is marked as unavailable rather than guessed.

Read the code

XAI the interpretability method   substrate the bit-exact emulator it runs on (jutari, Julia).

• XAIlinear_probing.jl:270 — The probe: closed-form ridge least squares, one-vs-rest, accuracy on held-out frames.
• XAIlinear_probing.jl:244 — Turns a byte series into up to 4 class labels (equal-width bins).
• XAIlinear_probing.jl:213 — Reads the oracle's “not causally used” cells from the DAS output.
• XAIlinear_probing.jl:157 — Records the 600-frame RAM tape the probes are trained on.
• XAIlinear_probing.jl:359 — Per game: hold out each cell, fit probe and control, selectivity, present-vs-used cross-check.
• XAIlinear_probing.jl:332 — The decodability thresholds (accuracy ≥ 0.5, selectivity ≥ 0.1).
• substrateoracle_intervene.jl:163 — Loads the verified concept↔cell map: the probe targets.
• substrateshared_testbed_impl.jl:115 — The shared seeded gameplay stream the tape starts from.
• substrateStellaEnvironment.jl:220 — Boots the cartridge (xitari's construction sequence: a 60-frame probe, then boots of 60 NOOP + 4 RESET) before recording.
• substrateStellaEnvironment.jl:350 — The live 128-byte RAM read after every frame to fill the tape.

Reading the result

Each bar is one labelled cell: its selectivity (probe accuracy minus control accuracy); red bars are cells the oracle flags as not causally used. Expect most game variables to be decodable above the control. The VCS state is small and full of copies and derived values, so the rest of RAM often predicts a variable well. Then look for the red cells that are still decodable: present but not used. They are the key result: a probe would “find” these concepts, a naive reading would say the program relies on them, and the intervention shows that, within the window, it does not.

Pitfalls

• Present is not used. The central lesson. A probe answers “can it be read?”. Only an intervention answers “does changing it change the output?”.
• Probe capacity. A powerful probe can learn the task itself. Use a control task and report selectivity; keep the probe linear and regularised.
• Neighbouring frames leak. We shuffle frames before splitting, and consecutive frames are nearly identical. Test accuracy can therefore be optimistic compared with a split by time.
• Binning shapes the task. Four equal-width bins over the observed range are a choice; a different binning changes both accuracies.
• Redundancy, not location. Because the target cell is held out, a high score means the variable can be inferred from other bytes. It does not say where the program stores it.

Try it yourself

julia --project=jutari tools/xai_study/phaseC_mechanistic/linear_probing.jl --games core

Writes linear_probing_<game>.{json,npz} and linear_probing_core_summary.json into tools/xai_study/phaseC_mechanistic/out/. Run das.jl first so the present-vs-used cross-check has the oracle's transient-cell list.

Original papers & further reading

• Alain & Bengio (2017). Understanding intermediate layers using linear classifier probes. ICLR workshop. — arxiv.org/abs/1610.01644
• Hewitt & Liang (2019). Designing and Interpreting Probes with Control Tasks. EMNLP. — arxiv.org/abs/1909.03368
• Belinkov (2022). Probing Classifiers: Promises, Shortcomings, and Advances. Computational Linguistics. (critique / survey) — arxiv.org/abs/2102.12452
• Elazar, Ravfogel, Jacovi & Goldberg (2021). Amnesic Probing: Behavioral Explanation with Amnesic Counterfactuals. TACL. (decodable is not used) — arxiv.org/abs/2006.00995

Read next: Logit / tuned lens · Interchange interventions / DAS · Sparse autoencoders · A3 · Tuning curves