MechRL: Reinforcement Learning Agents Perform Circuit Discovery for Mechanistic Interpretability

# Reinforcement Learning Agents Perform Circuit Discovery for Mechanistic Interpretability
Source: [https://arxiv.org/html/2605.26343](https://arxiv.org/html/2605.26343)
Barsat Khadka The University of Southern Mississippi Hattiesburg, MS 39406 Barsat\.Khadka@usm\.edu

###### Abstract

Mechanistic interpretability has identified small sets of attention heads that implement specific behaviours in transformer language models, but recovering these circuits typically requires a bespoke analytical pipeline for each new task\. We recast circuit discovery as a reinforcement\-learning problem\. An agent operates over the 144 attention heads of GPT\-2 small as a discrete action space; each action triggers a zero\-ablation and a contrastive reward that subtracts the ablation’s damage to general next\-token prediction from its damage to the target task\. A single PPO policy, trained on two tasks \(induction and IOI\) in a vectorised multi\-task environment, attains the per\-episode oracle on both training tasks and on a held\-out third task \(docstring completion\)\. Its preferred heads coincide with the canonical heads ofOlssonet al\.\([2022](https://arxiv.org/html/2605.26343#bib.bib2)\)andWanget al\.\([2023](https://arxiv.org/html/2605.26343#bib.bib3)\)on precisely the axes those papers identify as causally non\-redundant under single\-head ablation; the categories they identify as redundant are correctly de\-prioritised by the agent\. On the held\-out task, best\-of\-five planning recovers 96% of the oracle ceiling with no task signal supplied at evaluation\. These results indicate that reinforcement learning over causal interventions is a viable, transferable substrate for identifying the single\-head bottlenecks of mechanistic circuits, complementary to existing path\-patching approaches\.

## 1Introduction

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\(Elhageet al\.,[2021](https://arxiv.org/html/2605.26343#bib.bib1)\)\. While this approach has isolated mechanisms such as the induction circuit\(Olssonet al\.,[2022](https://arxiv.org/html/2605.26343#bib.bib2)\)and the IOI circuit\(Wanget al\.,[2023](https://arxiv.org/html/2605.26343#bib.bib3)\), current methodologies remain highly localised\. Both manual analyses and automated methods like ACDC\(Conmyet al\.,[2023](https://arxiv.org/html/2605.26343#bib.bib4)\)and EAP\(Nanda,[2023](https://arxiv.org/html/2605.26343#bib.bib5)\)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\.

## 2Method

### 2\.1Environment

We work with frozen pretrained GPT\-2 small \(nlayers=nheads=12n\_\{\\text\{layers\}\}=n\_\{\\text\{heads\}\}=12,nactions=144n\_\{\\text\{actions\}\}=144\)\. The environment exposes the standard Gymnasium API\. Each episode comprises at mostT=50T=50steps; the action space𝒜=\{0,…,143\}\\mathcal\{A\}=\\\{0,\\dots,143\\\}corresponds bijectively to\(ℓ,h\)\(\\ell,h\)viaa=12​ℓ\+ha=12\\ell\+h\. An action mask precludes within\-episode repetition of any head\.

At everyreset\(\), the environment samples a taskτ\\tauuniformly from the training set\{induction,IOI\}\\\{\\text\{induction\},\\,\\text\{IOI\}\\\}and generates a fresh task batch using the episode seed\. Induction batches follow the canonical\[BOS,A,B,fillers,A\]\[\\text\{BOS\},A,B,\\text\{fillers\},A\]template withA,BA,Bsampled uniformly from the vocabulary; IOI batches followWanget al\.\([2023](https://arxiv.org/html/2605.26343#bib.bib3)\)’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 episodettis not in general the same head that maximises damage on episodet′t^\{\\prime\}, so the policy must internalise distributional rather than pointwise structure\.

### 2\.2Contrastive Reward

LetMθτ​\(B\)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θ\\thetaover batchBB\. LetLθctrl​\(C\)L^\{\\text\{ctrl\}\}\_\{\\theta\}\(C\)denote the mean cross\-entropy ofθ\\thetaon a control batchCCof natural English text drawn fromwikitext\-2\-raw\-v1\. For the ablated modelθ∖a\\theta\_\{\\setminus a\}, in which the\(ℓ,h\)\(\\ell,h\)attention\-head output athook\_zis set to zero, we define the per\-action reward

r​\(a\)=\(Mθτ​\(B\)−Mθ∖aτ​\(B\)\)⏟task damage−\(Lθ∖actrl​\(C\)−Lθctrl​\(C\)\)⏟general damage\.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\}\}\.\(1\)The reward is positive when ablating headaadisproportionately 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 atreset\(\)and reused across all steps in the episode\.

### 2\.3Observation

The agent observesot∈ℝ2\+2​nactionso\_\{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 rewardsr​\(a\)/σr\(a\)/\\sigmafor previously\-tried heads \(zero otherwise\), with fixedσ=5\\sigma=5\. The reward channel enables within\-episode adaptation: the policy can condition future actions on feedback already received in the current episode\.

### 2\.4Multi\-Task Vectorised PPO

We train with proximal policy optimisation\(Schulmanet al\.,[2017](https://arxiv.org/html/2605.26343#bib.bib6)\)over a synchronous vector ofN=8N=8parallel environments sharing a single GPT\-2 instance on the GPU, yieldingN⋅T=400N\\cdot T=400on\-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 architecture290→256→256→\(144,1\)290\\to 256\\to 256\\to\(144,\\,1\), tanh activations, and orthogonal initialisation following standard PPO practice\. We optimise with Adam at initial learning rate2\.5×10−42\.5\\\!\\times\\\!10^\{\-4\}annealed linearly to 20% of its initial value,γ=0\.99\\gamma=0\.99, GAEλ=0\.95\\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 forTtotal=2×105T\_\{\\text\{total\}\}=2\\\!\\times\\\!10^\{5\}environment steps\.

### 2\.5Best\-of\-KKPlanning

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\-KKplanning wrapper\. At decision time the policy samplesKKcandidate actions without replacement from its current distribution; the environment scores each in parallel; and the agent commits to the highest\-scoring candidate\.K=1K=1recovers standard PPO\. We report results forK∈\{1,5\}K\\in\\\{1,5\\\}throughout\.

### 2\.6Held\-Out Evaluation Task

To assess generalisation, we introduce a docstring\-completion task the policy never sees during training, in the style ofHeimersheim and Janiak \([2023](https://arxiv.org/html/2605.26343#bib.bib7)\)\. Each prompt presents a Python function signature withn=5n=5parameters; the docstring describes the firstn−1n\-1in:paramform; 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 thatHeimersheim and Janiak \([2023](https://arxiv.org/html/2605.26343#bib.bib7)\)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\.

## 3Results

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

### 3\.1Convergence 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 \(seed≥107\\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](https://arxiv.org/html/2605.26343#S2.E1)\. We compare this against the running maximum reward attained by the trained policy over its 50 in\-episode picks\.

Table 1:Mean running\-max contrastive reward of the trainedK=1K=1policy versus the per\-episode oracle over≥20\\geq 20held\-out evaluation episodes per task\. The policy attains the oracle ceiling within sampling noise on all three tasks, including the held\-out docstring task\.The gap between policy and oracle is bounded by\|0\.028\|\|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\.2Alignment 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\.

#### Induction\.

The canonical induction\-head set in GPT\-2 small is widely taken to be\{L​5\.H​1,L​5\.H​5,L​6\.H​9,L​7\.H​2,L​7\.H​10\}\\\{L5\.H1,\\,L5\.H5,\\,L6\.H9,\\,L7\.H2,\\,L7\.H10\\\}\(Olssonet al\.,[2022](https://arxiv.org/html/2605.26343#bib.bib2)\)\.

Table 2:Pick\-frequency rank of the canonical induction heads in the trained policy \(K=1K=1, 20 episodes\) and the per\-episode oracle \(10 episodes\)\. “–” indicates the head was never selected\. Heads the oracle ranks low \(L7\.H10, L5\.H1\) are also absent from the policy\.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 byOlssonet al\.\([2022](https://arxiv.org/html/2605.26343#bib.bib2)\)as the most prototypical induction head in the model, is correctly surfaced at rank 6\.

#### IOI\.

The IOI circuit decomposes into seven sub\-categories\(Wanget al\.,[2023](https://arxiv.org/html/2605.26343#bib.bib3)\); Table[3](https://arxiv.org/html/2605.26343#S3.T3)reports per\-category overlap between the policy’s top\-ten picks and the canonical sets\.

Table 3:Per\-sub\-category overlap between the trained policy’s top\-10 most\-picked heads and the canonical IOI circuit\(Wanget al\.,[2023](https://arxiv.org/html/2605.26343#bib.bib3)\)\. The policy recovers 3 of 4 S\-Inhibition heads and 2 of 4 Induction\-in\-IOI heads; it does not pick any Name Movers, consistent with[Wanget al\.](https://arxiv.org/html/2605.26343#bib.bib3)’s minimality analysis showing single\-head Name Mover ablation is masked by Backup Name Movers\.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 byWanget al\.\([2023](https://arxiv.org/html/2605.26343#bib.bib3)\): ablating a single Name Mover induces only a∼5%\{\\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[Wanget al\.](https://arxiv.org/html/2605.26343#bib.bib3)’s minimality analysis suggests that single\-head causal indispensability is a coherent and reproducible interpretive target\.

### 3\.3Transfer to a Held\-Out Task

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

Table 4:Mean running\-max contrastive reward on the held\-out docstring task over 30 episodes\. Best\-of\-5 planning under the zero\-shot\[0,0\]\[0,0\]signal recovers 96% of the oracle ceiling\. Priming as induction matches the ceiling; priming as IOI falls below random, demonstrating non\-trivial use of the task channel\.Three features of Table[4](https://arxiv.org/html/2605.26343#S3.T4)warrant comment\. First, the zero\-shotK=5K=5result \(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 atK=1K=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=5K=5yields only marginal improvement on training tasks where the policy is already certain, but lifts zero\-shot performance by1\.4×1\.4\\timeson 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\.

## 4Discussion

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\.

#### 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\.

#### 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\.

#### Code and reproducibility\.

The complete pipeline \(environment, training, evaluation, oracle, canonical\-head comparison\) is available at[https://github\.com/BarsatKhadka/causality\-RL](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 atK=1K=1and 6 hours atK=5K=5\.

## Acknowledgement

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

## References

- A\. Conmy, A\. Mavor\-Parker, A\. Lynch, S\. Heimersheim, and A\. Garriga\-Alonso \(2023\)Towards automated circuit discovery for mechanistic interpretability\.InAdvances in Neural Information Processing Systems,Vol\.36\.External Links:[Link](https://proceedings.neurips.cc/paper_files/paper/2023/hash/34e1dbe95d34d7ebaf99b9bcaeb5b2be-Abstract-Conference.html)Cited by:[§1](https://arxiv.org/html/2605.26343#S1.p1.1)\.
- N\. Elhage, N\. Nanda, C\. Olsson, T\. Henighan, N\. Joseph, B\. Mann, A\. Askell, Y\. Bai, A\. Chen, T\. Conerly, N\. DasSarma, D\. Drain, D\. Ganguli, Z\. Hatfield\-Dodds, D\. Hernandez, A\. Jones, J\. Kernion, L\. Lovitt, K\. Ndousse, D\. Amodei, T\. Brown, J\. Clark, J\. Kaplan, S\. McCandlish, and C\. Olah \(2021\)A mathematical framework for transformer circuits\.Note:Transformer Circuits ThreadExternal Links:[Link](https://transformer-circuits.pub/2021/framework/index.html)Cited by:[§1](https://arxiv.org/html/2605.26343#S1.p1.1)\.
- S\. Heimersheim and J\. Janiak \(2023\)A circuit for Python docstrings in a 4\-layer attention\-only transformer\.Note:AI Alignment ForumExternal Links:[Link](https://www.alignmentforum.org/posts/u6KXXmKFbXfWzoAXn)Cited by:[§2\.6](https://arxiv.org/html/2605.26343#S2.SS6.p1.2),[§2\.6](https://arxiv.org/html/2605.26343#S2.SS6.p2.1)\.
- N\. Nanda \(2023\)Attribution patching: activation patching at industrial scale\.Note:[https://www\.neelnanda\.io/mechanistic\-interpretability/attribution\-patching](https://www.neelnanda.io/mechanistic-interpretability/attribution-patching)Cited by:[§1](https://arxiv.org/html/2605.26343#S1.p1.1)\.
- C\. Olsson, N\. Elhage, N\. Nanda, N\. Joseph, N\. DasSarma, T\. Henighan, B\. Mann, A\. Askell, Y\. Bai, A\. Chen, T\. Conerly, D\. Drain, D\. Ganguli, Z\. Hatfield\-Dodds, D\. Hernandez, S\. Johnston, A\. Jones, J\. Kernion, L\. Lovitt, K\. Ndousse, D\. Amodei, T\. Brown, J\. Clark, J\. Kaplan, S\. McCandlish, and C\. Olah \(2022\)In\-context learning and induction heads\.Note:Transformer Circuits ThreadExternal Links:2209\.11895,[Link](https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html)Cited by:[§1](https://arxiv.org/html/2605.26343#S1.p1.1),[§3\.2](https://arxiv.org/html/2605.26343#S3.SS2.SSS0.Px1.p1.1),[§3\.2](https://arxiv.org/html/2605.26343#S3.SS2.SSS0.Px1.p2.1)\.
- J\. Schulman, F\. Wolski, P\. Dhariwal, A\. Radford, and O\. Klimov \(2017\)Proximal policy optimization algorithms\.External Links:1707\.06347Cited by:[§2\.4](https://arxiv.org/html/2605.26343#S2.SS4.p1.2)\.
- K\. Wang, A\. Variengien, A\. Conmy, B\. Shlegeris, and J\. Steinhardt \(2023\)Interpretability in the wild: a circuit for indirect object identification in GPT\-2 small\.InInternational Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=NpsVSN6o4ul)Cited by:[§1](https://arxiv.org/html/2605.26343#S1.p1.1),[§2\.1](https://arxiv.org/html/2605.26343#S2.SS1.p2.6),[§3\.2](https://arxiv.org/html/2605.26343#S3.SS2.SSS0.Px2.p1.1),[§3\.2](https://arxiv.org/html/2605.26343#S3.SS2.SSS0.Px2.p2.1),[§3\.2](https://arxiv.org/html/2605.26343#S3.SS2.SSS0.Px2.p3.1),[Table 3](https://arxiv.org/html/2605.26343#S3.T3),[Table 3](https://arxiv.org/html/2605.26343#S3.T3.1.1.2),[Table 3](https://arxiv.org/html/2605.26343#S3.T3.4.2)\.