Research

Dissertation: Causal Cue Retrieval and Greedy Route Execution — A Mechanistic Analysis of a Transformer Policy

Back to Projects

PyTorch

RL (PPO)

Mechanistic Interpretability

MiniGrid

Python

1.000

Transformer @ 32 bindings

0.688

Matched LSTM, retrospective @ 8

6 / 6

Paired gap, seeds reproducing

3 / 6

Causal signature held in

Figures

All three policies are trained under one shared PPO protocol — same task, same encodings, same cue-balanced schedule — so any difference comes from the architecture rather than the setup. The matched LSTM is included as a fairness control at equal parameter count, and it can solve the full task.

Why aggregate success is not enough. The dashed line marks a cue-blind policy that always takes the same route: it scores 0.5 on overall success while having learned nothing about the cue. Smin = min(S_key, S_ball) collapses for such a policy, which is what makes the failure mode visible.

Causal interventions on 1,000 held-out episodes. Under normal operation the policy is confident in the correct branch. Blocking retrieval pushes it to no branch preference (0.43 / 0.57), and patching in the opposite memory flips the decision outright (0.03 / 0.09) — evidence that the retrieved cue is used, not merely present.

Overview

A mechanistic study of how a PPO-trained causal Transformer policy solves a partially observable MiniGrid T-maze, where the instruction cue is no longer visible at the point the decision is made. The headline result is methodological rather than architectural: the aggregate success metric hid a failure mode, and a worst-case metric plus causal interventions were needed to tell genuine memory use apart from a shortcut.

Problem

In a T-maze where the cue (key or ball) disappears before the junction, an agent that always turns the same way scores roughly 50% on mean success while having learned nothing about the cue. Aggregate success therefore cannot distinguish a policy that retrieves the instruction from one that has collapsed onto a single route. Any claim about memory made on that metric is unfalsifiable.

Making the task honest

Before measuring anything I removed the ways the result could be trivially explained: a symmetric maze layout so neither branch is geometrically favoured, and explicit verification that no visual leakage lets the agent observe the cue at decision time. Without this step a positive result would say nothing about memory.

The metric

I defined Smin = min(S_key, S_ball), the worse of the two cue branches. One-sided policies that look competent under mean success collapse under Smin, which makes the failure mode visible instead of averaged away. Checkpoint selection was locked to a validation split so the reported number is not chosen on the test set.

Causal interventions

Correlational evidence — a probe decoding the cue from a hidden state — shows only that the information is present, not that the policy uses it. To test use rather than presence I applied three intervention families: retrieval blocking, opposite-memory patching, and norm-preserving source exchange, the last controlling for the possibility that an effect is merely an activation-magnitude artefact. Ridge-decoded probes targeted the retrieval output and the gated feature at the decision timesteps.

Results, including the inconvenient one

Difference-vector attribution assigned 33–58% of the key–ball retrieval difference to padding. But exchanging the padding-derived source reversed no decision, while exchanging the real-transition source reversed both — at exactly preserved norm, so the effect was not an artefact of activation magnitude. Correlation and causation pointed at different sources on the same policy, which is the result the dissertation is built around. Two things complicate it, and both are reported. The symmetric real-only signature appeared in only three of six checkpoints, so equal behaviour did not imply equal mechanism. And a parameter-matched LSTM (h = 304, 2.90M parameters) also solved the task, at Smin = 0.898. That removed the architecture-comparison claim I started with, so the contribution moved to where the evidence actually supported it: the analysis and the measurement framework, not a ranking of architectures. The one clean architectural separation came from a preregistered fixed-length follow-up that isolates retrospective load from navigation. The Transformer held 1.000 accuracy through 32 stored bindings and 0.9997 at 64. At eight bindings with a retrospective query, the matched LSTM reached 0.688 against the Transformer's 0.997 — a paired gap of +0.308 (95% CI [+0.294, +0.326]), positive in all six held-out seeds, and still +0.243 against an LSTM four times the size. The separation is specific, and the dissertation says so: prospective queries do not separate the architectures the same way (the matched LSTM reaches 0.878 at eight bindings there). The supported claim is timing-moderated load sensitivity in this task, not a general ranking of architectures.

What is reusable

A branch-level failure decomposition separating Type-I (readout) from Type-II (aggregation) failures, locked checkpoint selection, and a causal context-ablation taxonomy. Together these separate real memory use from shortcut behaviour in any task with a similar branch structure — the same problem that golden sets and adversarial cases address in LLM evaluation.