1 Introduction↩︎

Mechanistic interpretability seeks to reverse-engineer neural networks to understand how specific internal operations produce model behaviours. In transformer models, this involves identifying circuits: sparse subgraphs of attention heads, MLP layers, and residual connections that are causally sufficient for a given task [3]. While this approach has isolated mechanisms such as the induction circuit [1] and the IOI circuit [2], current methodologies remain highly localised. Both manual analyses and automated methods like ACDC [4] and EAP [5] evaluate single tasks in isolation without sharing learned structural patterns across different behaviours.

In this work, we investigate whether a learned policy can automate circuit discovery and generalise to unseen tasks. We frame circuit discovery as a Markov decision process. At each step, an agent zero-ablates one attention head, observes a reward based on the ablation’s causal effect, and updates its policy. By randomising the task batch in each episode, we prevent the agent from memorising specific head locations, forcing it to learn structural priors about the model’s geometry.

Our reward function uses a contrastive design to distinguish task-specific heads from general-purpose computational heads. Because ablating a general-purpose head degrades performance across all tasks, a standard logit-difference drop is insufficient. We therefore pair each ablation with a control evaluation on a batch of natural English text, and subtract the cross-entropy increase on the control from the logit-difference drop on the target task. This isolates the task-specific causal effect and, as we demonstrate empirically, provides a reliable reinforcement-learning signal.

2 Method↩︎

2.1 Environment↩︎

We work with frozen pretrained GPT-2 small (\(n_\text{layers} = n_\text{heads} = 12\), \(n_\text{actions} = 144\)). The environment exposes the standard Gymnasium API. Each episode comprises at most \(T = 50\) steps; the action space \(\mathcal{A} = \{0,\dots,143\}\) corresponds bijectively to \((\ell,h)\) via \(a = 12\ell + h\). An action mask precludes within-episode repetition of any head.

At every reset(), the environment samples a task \(\tau\) uniformly from the training set \(\{\text{induction},\,\text{IOI}\}\) and generates a fresh task batch using the episode seed. Induction batches follow the canonical \([\text{BOS}, A, B, \text{fillers}, A]\) template with \(A,B\) sampled uniformly from the vocabulary; IOI batches follow [2]’s template family with random name-object-place draws and balanced ABBA/BABA orderings. Per-episode resampling enforces a stochastic reward landscape: the head whose ablation maximises damage on episode \(t\) is not in general the same head that maximises damage on episode \(t'\), so the policy must internalise distributional rather than pointwise structure.

2.2 Contrastive Reward↩︎

Let \(M^\tau_\theta(B)\) denote the task-specific scalar metric (logit difference between the correct and distractor tokens at the relevant position) evaluated on frozen model \(\theta\) over batch \(B\). Let \(L^{\text{ctrl}}_\theta(C)\) denote the mean cross-entropy of \(\theta\) on a control batch \(C\) of natural English text drawn from wikitext-2-raw-v1. For the ablated model \(\theta_{\setminus a}\), in which the \((\ell,h)\) attention-head output at hook_z is set to zero, we define the per-action reward \[r(a) \;=\; \underbrace{\bigl(M^\tau_\theta(B) - M^\tau_{\theta_{\setminus a}}(B)\bigr)}_{\text{task damage}} \;-\; \underbrace{\bigl(L^{\text{ctrl}}_{\theta_{\setminus a}}(C) - L^{\text{ctrl}}_\theta(C)\bigr)}_{\text{general damage}}. \label{eq:reward}\tag{1}\] The reward is positive when ablating head \(a\) disproportionately damages the target task; it is approximately zero for heads whose ablation degrades the model’s general distributional fit by a comparable amount. Each query incurs two forward passes through \(\theta\) (one task batch, one control batch); both are pre-allocated on the GPU at reset() and reused across all steps in the episode.

2.3 Observation↩︎

The agent observes \(o_t \in \mathbb{R}^{2 + 2n_\text{actions}}\) comprising three blocks: a 2-dimensional one-hot encoding the current task \(\tau\); a 144-dimensional binary mask of previously-tried heads; and a 144-dimensional vector of normalised rewards \(r(a)/\sigma\) for previously-tried heads (zero otherwise), with fixed \(\sigma = 5\). The reward channel enables within-episode adaptation: the policy can condition future actions on feedback already received in the current episode.

2.4 Multi-Task Vectorised PPO↩︎

We train with proximal policy optimisation [6] over a synchronous vector of \(N = 8\) parallel environments sharing a single GPT-2 instance on the GPU, yielding \(N \cdot T = 400\) on-policy transitions per gradient update. A preliminary single-environment version collapsed to a near-deterministic policy within approximately 2,000 steps; we attribute this to a per-update batch too small for stable advantage estimation.

The actor-critic is a feedforward MLP with architecture \(290 \to 256 \to 256 \to (144,\,1)\), tanh activations, and orthogonal initialisation following standard PPO practice. We optimise with Adam at initial learning rate \(2.5\!\times\!10^{-4}\) annealed linearly to 20% of its initial value, \(\gamma = 0.99\), GAE \(\lambda = 0.95\), clip coefficient 0.2, entropy coefficient 0.1, value loss coefficient 0.5, and 4 update epochs with 8 minibatches each. Training runs for \(T_\text{total} = 2\!\times\!10^5\) environment steps.

2.5 Best-of-\(K\) Planning↩︎

Because the environment dynamics are the GPT-2 forward pass, deterministic and cheap relative to the policy network, we equip the agent with a best-of-\(K\) planning wrapper. At decision time the policy samples \(K\) candidate actions without replacement from its current distribution; the environment scores each in parallel; and the agent commits to the highest-scoring candidate. \(K = 1\) recovers standard PPO. We report results for \(K \in \{1, 5\}\) throughout.

2.6 Held-Out Evaluation Task↩︎

To assess generalisation, we introduce a docstring-completion task the policy never sees during training, in the style of [7]. Each prompt presents a Python function signature with \(n = 5\) parameters; the docstring describes the first \(n - 1\) in :param form; the model must predict the missing final parameter name at the position immediately following the last :param. The task metric is the logit difference between the correct missing parameter name and a distractor drawn from the parameters already mentioned.

We note that [7] report the canonical docstring circuit for a 4-layer attention-only toy model, not GPT-2 small. There is accordingly no published head-level ground truth for this task in our model, and we report only oracle agreement and transfer performance, not canonical overlap.

3 Results↩︎

We report three classes of result: convergence to per-task optimality, alignment with canonical interpretability findings, and transfer to a held-out task.

3.1 Convergence to the Per-Task Oracle↩︎

For each task we compute a per-episode oracle by exhaustively ablating all 144 heads on the episode’s batch and recording the single highest-scoring head. Averaging over a held-out evaluation seed band (\(\text{seed} \geq 10^7\)) yields the oracle ceiling: the maximum running-max reward attainable by any single-head policy under the contrastive reward of Equation 1 . We compare this against the running maximum reward attained by the trained policy over its 50 in-episode picks.

The gap between policy and oracle is bounded by \(|0.028|\) across all three tasks, within the per-episode standard error of the oracle estimate itself. We conclude that the contrastive reward is well-posed: the heads the oracle would select are, with high probability, the heads the policy selects. This licenses the analyses below: the policy’s pick-frequency distribution is a reasonable proxy for the reward landscape itself.

3.2 Alignment with the Canonical Interpretability Literature↩︎

We compare the trained policy’s pick-frequency distribution against the canonical head sets identified in prior work. For each canonical head we report (a) its rank within the policy’s pick-frequency ordering over 20 held-out episodes, and (b) its rank within the oracle’s mean-score ordering over 10 held-out episodes in both cases out of 144 heads.

3.2.0.1 Induction.

The canonical induction-head set in GPT-2 small is widely taken to be \(\{L5.H1,\,L5.H5,\,L6.H9,\,L7.H2,\,L7.H10\}\) [1].

Two of the five canonical induction heads—L5.H5 and L7.H2—appear in the policy’s top ten picks. The remaining three are absent, but the oracle confirms this is not a policy failure: L7.H10 and L5.H1 sit at oracle ranks 129 and 130 respectively. Under single-head zero-ablation with natural-text control, these heads do not damage induction in excess of their damage to general next-token prediction. L5.H5, identified by [1] as the most prototypical induction head in the model, is correctly surfaced at rank 6.

3.2.0.2 IOI.

The IOI circuit decomposes into seven sub-categories [2]; Table 3 reports per-category overlap between the policy’s top-ten picks and the canonical sets.

S-Inhibition heads dominate the policy’s IOI top picks: L8.H10 at rank 1, L7.H9 at rank 3, L8.H6 at rank 4. Two of the four Induction-in-IOI heads follow at ranks 9 and 10. The absence of Name Movers reflects the redundancy structure documented by [2]: ablating a single Name Mover induces only a \({\sim}5\%\) logit-difference drop because Backup Name Movers reroute the relevant signal. The oracle confirms this externally—L9.H6, the most-cited Name Mover, sits at oracle rank 138. The method surfaces the IOI circuit’s bottlenecks under single-head ablation, which coincide with the S-Inhibition family rather than with the more behaviourally legible Name Mover family.

This is not a deficiency of the method but a property of single-head ablation itself. A circuit’s most legible heads are not necessarily its most causally indispensable heads under unit-level intervention. The convergence between our agent and [2]’s minimality analysis suggests that single-head causal indispensability is a coherent and reproducible interpretive target.

3.3 Transfer to a Held-Out Task↩︎

We evaluate transfer to the docstring task under three task-signal regimes: \([0,0]\) (zero-shot; a vector never seen during training), \([1,0]\) (primed as induction), and \([0,1]\) (primed as IOI).

Three features of Table 4 warrant comment. First, the zero-shot \(K=5\) result (2.103 against an oracle of 2.175, versus a random baseline of 1.395) constitutes non-trivial generalisation: the policy recovers 96% of the optimal reward on a task it has never seen. The +0.71 gap over random (51% relative improvement) is wholly attributable to learned priors plus in-episode reward feedback under planning. Second, the priming differential , 2.178 versus 1.057 at \(K=1\) , shows the policy maintains distinct task-conditional strategies, and that the task one-hot functions as a routing signal rather than a vestigial input. Third, planning matters substantially more for transfer than for training: \(K=5\) yields only marginal improvement on training tasks where the policy is already certain, but lifts zero-shot performance by \(1.4\times\) on the held-out task, where the planner’s real GPT-2 evaluations are decisive.

The structural-similarity ordering in the priming results is interpretable. The docstring task is a pattern-completion , locating a missing parameter name in a structured context , closer in form to the induction circuit’s previous-token-to-current-token matching than to the IOI circuit’s name-resolution computation. The policy’s induction strategy transfers; its IOI strategy actively misleads.

4 Discussion↩︎

The result we are most willing to defend is the methodological one. A single reinforcement-learning policy, trained on causal interventions across two distinct circuits with a contrastive reward, attains the per-task oracle on three tasks (two seen, one unseen) and produces head preferences that align with the canonical interpretability literature on precisely the axis our methodology measures. The transfer result demonstrates that the policy’s learned structure is not exhausted by per-task memorisation.

The result we are least willing to defend is the strongest reading of “circuit discovery.” Our method recovers the single-head bottlenecks of a circuit: heads whose individual ablation maximally degrades the target task in excess of the general degradation it induces. This does not coincide with the circuit-implementing heads that emerge from path-patching or attention-pattern analyses. The Name Mover case is the clearest instance: behaviourally, Name Movers are the IOI circuit’s signature heads; causally, under single-head ablation, they are not its bottlenecks. Both characterisations are correct; they characterise different things. We view our methodology as complementary to existing path-patching approaches.

4.0.0.1 Limitations.

The contrastive reward combines a logit-difference scale and a cross-entropy scale that do not formally align in units. Empirically the estimator behaves well; theoretically a unit-consistent variant (cross-entropy on both sides, or a constructed logit-difference control) would be preferable. The transfer assessment relies on a single held-out task; a stronger claim would require evaluation across a suite of held-out circuits. The agent is trained for GPT-2 small specifically; while the environment is model-agnostic, generalisation across architectures has not been measured. Finally, the action space is restricted to single-head ablations, precluding the discovery of compositional structure such as the Name-Mover/Backup-Name-Mover redundancy that requires multi-head interventions to surface.

4.0.0.2 Future work.

Three directions follow naturally from the present results. First, replacing the explicit task one-hot with a learned context encoder that infers a task representation from the in-episode reward vector would remove the test-time priming requirement. Second, scaling to larger models (GPT-2 medium and beyond, or non-GPT architectures) would test the universality of the structural priors the agent learns. Third, expanding the action space to multi-head ablations would allow the agent to surface circuit composition, including the redundancy structures the present methodology cannot detect.

4.0.0.3 Code and reproducibility.

The complete pipeline (environment, training, evaluation, oracle, canonical-head comparison) is available at https://github.com/BarsatKhadka/causality-RL. Training was performed on a single NVIDIA L40S GPU; the 200,000-step run completes in approximately 2.5 hours at \(K=1\) and 6 hours at \(K=5\).

Acknowledgement↩︎

The author sincerely thanks the Mississippi Center for Supercomputing Research (MCSR) for compute support.

References↩︎