Reasoning error detection by probing hidden states

Language models write out their working when they solve a math problem, and some of those steps are just wrong. This is a small study around one question: while the model is writing a step, do its hidden states already "know" whether that step is correct — enough that a cheap linear probe can read it off?

Short answer: mostly yes, and the signal lives in the hidden states specifically — not in attention patterns, not in output confidence. Mean-pooling the hidden states over a step beats using the last token, and it's the only feature set that still works on a dataset the probe never trained on.

What I did

I generated step-by-step solutions with Qwen2.5-0.5B-Instruct (greedy decoding, a structured Step N: format), then labeled each step correct or incorrect with a process reward model — ThinkPRM-1.5B — prompted as a strict examiner so it catches a broken step early instead of rubber-stamping the whole trace.

From every step I pulled five kinds of features and trained a plain logistic-regression probe on each, to see which ones actually carry the error signal:

1. hidden states from four layers (¼, ½, ¾, final), last token only
2. the same four layers, but mean-pooled over all the step's tokens
3. final-layer attention statistics — spread, and how much the step attends to the question vs. the answer
4. output-confidence scalars — mean max-probability, its spread, and the min confidence over the step
5. everything above, concatenated

Probes train on GSM8K, and I test on held-out GSM8K and on AQuA — the AQuA number is the one I actually care about, since it tells you whether the probe learned something about reasoning or just something about GSM8K.

What came out

• The hidden states are doing the work. Every config that includes them clears the baselines; the ones without them (attention-only, logits-only) sit right at chance.
• Mean-pooling beats last-token, and it holds up best out of distribution (0.634 AUC on AQuA). Averaging over the whole step seems to capture the error better than the final token alone.
• Throwing everything together doesn't help. The combined set is a touch worse than mean-pooling on GSM8K (0.719 vs 0.745) — the weak features mostly add noise. (It does edge out on raw accuracy, but AUC is the fairer read here.)

Caveats I'd want a reader to know

• Steps are split with a Step N: regex, which won't line up perfectly with the model's real thought boundaries — that leaks some noise into the extracted states.
• The PRM isn't a perfect labeler; on long traces it can still wave a bad step through.
• These are linear probes, so all this shows is that the error signal is linearly readable. A small MLP might pull out more.

Things I'd try next: probing individual attention heads instead of aggregate stats, non-linear probes, and larger models to see whether the error representation gets sharper with scale.

Running it

pip install -r requirements.txt
python main_experiments.py

That generates the traces (GSM8K train ~1024, then GSM8K + AQuA test), trains all five probes, evaluates them, and writes everything to results_experiments/. Budget roughly 3–4 hours on a GPU.

Two flags worth knowing:

• --debug-single-exp — runs only the first experiment, for a quick sanity check
• --use-existing-traces — skips trace generation and reuses saved traces, so you can iterate on features without re-running the model

Outputs in results_experiments/: per-experiment JSON (exp{N}_results.json), the traces each ran on, and all_experiments_summary.json with the combined table.