MechRL: Reinforcement Learning Agents Perform Circuit Discovery for Mechanistic Interpretability

Abstract: Mechanistic interpretability has identified small sets of attention heads that implement specific behaviours in transformer LLMs, 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 of established literature 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.

• The paper introduces a novel reinforcement learning framework that automates circuit discovery in transformers through targeted zero-ablation interventions.
• It employs a contrastive reward function to distinguish task-specific from general computational heads, achieving near-oracle performance on multiple tasks.
• The RL approach demonstrates robust generalization and complements existing interpretability techniques by surfacing bottleneck heads critical for task performance.

Reinforcement Learning for Mechanistic Circuit Discovery in Transformers

Framing Circuit Discovery as RL

The paper "MechRL: Reinforcement Learning Agents Perform Circuit Discovery for Mechanistic Interpretability" (2605.26343) introduces a novel methodology to automate circuit discovery in transformer LLMs. Traditionally, mechanistic interpretability relies on targeted analysis or automated pipelines such as ACDC and EAP, each tailored to a specific behavioral probe. This work reframes circuit discovery as a reinforcement learning (RL) problem, leveraging a policy trained via PPO to perform targeted zero-ablation interventions over the 144 attention heads of GPT-2 small. The agent's action space is exhaustive and discrete, mapping directly to each attention head. A contrastive reward function is designed to distinguish task-specific heads from general computational heads by evaluating the net causal impact of ablation on task performance relative to general next-token prediction.

Methodological Innovations

The core methodological innovation is the application of RL to circuit discovery with a contrastive causal reward. The environment is structured as a vectorized multi-task Gymnasium API, supporting episodic sampling of tasks (induction and IOI) and enforcing stochasticity in task batches to prevent memorization. The reward function computes the difference in logit-difference loss on the target task batch and the increase in cross-entropy on a control batch of natural English text. This isolates task-specific effects and ensures the RL agent is incentivized to prioritize heads whose ablation disproportionately damages the target task.

Policy observations consist of the current task encoded one-hot, a binary mask tracking previously ablated heads, and normalized rewards from prior actions within the episode. The policy network is implemented as a feedforward MLP, trained synchronously over eight parallel environments to ensure sufficient sampling for advantage estimation. The "best-of-$K$ planning" strategy further enables efficient parallel candidate evaluation.

Empirical Results and Alignment with Literature

The agent's performance is assessed across three axes: per-task optimality, alignment with canonical mechanistic circuits, and generalization to a held-out docstring-completion task.

1. Convergence: The RL policy attains the per-episode oracle ceiling for induction, IOI, and the held-out docstring task, with gaps $\leq 0.028$, which is within sampling error margins. This demonstrates that the contrastive reward is consistently well-posed, and the RL agent reliably identifies the heads that are most causally necessary for specific tasks.
2. Canonical Alignment: On induction and IOI tasks, the agent's pick-frequency of attention heads closely mirrors prior interpretability findings. Notably, L5.H5 (the most prototypical induction head) is prioritized, while redundant heads (e.g., Name Movers in IOI) are deprioritized, consistent with established redundancy structures. The method surfaces bottleneck heads under single-head ablation, aligning with earlier minimality analyses and revealing the distinction between legibility and causal indispensability.
3. Generalization: On the unseen docstring task, the RL policy achieves 96% of the held-out oracle, with a 51% improvement over random baseline under best-of-5 planning. Task-conditional routing demonstrates that the learned policy encodes structural priors transferable to out-of-distribution tasks. The priming differential reflects that structural similarity between tasks (pattern completion vs. name resolution) dictates transfer efficacy.

Implications and Theoretical Considerations

This work provides strong evidence that RL agents can discover circuit bottlenecks that are consistent with causal minimality criteria across tasks and generalize structural priors to unseen behaviors. Critically, the method is complementary to path-patching and attention-pattern-based approaches, as the RL agent surfaces bottleneck heads (those whose ablation most degrades task performance) rather than those heads that are merely legible in behavioral decompositions. This distinction underscores the need for diverse methodologies in mechanistic interpretability to capture both causal and representational aspects of circuit computation.

The approach also demonstrates scalability in terms of evaluation efficiency and reproducibility, with the full pipeline available for independent verification. However, the unit-mismatch between logit-difference and cross-entropy scales in the reward function warrants theoretical refinement. The restriction to single-head ablation precludes compositional circuit discovery (e.g., redundancy structures), and generalization across model architectures remains unexplored.

Future Directions

Several avenues emerge for extending this work:

• Eliminating task priming with a learned context encoder to enable more general inference
• Scaling to larger architectures and diverse model classes to test transferability of structural priors
• Expanding the action space to multi-head interventions to capture compositional and redundant circuit structures

Such directions are expected to deepen understanding of how causal structure emerges and is encoded within neural LLMs, informing both interpretability and robustness strategies.

Conclusion

"MechRL" demonstrates that RL agents trained via contrastive causal rewards can recover canonical circuit bottlenecks in transformer models and generalize to out-of-distribution tasks. The method aligns robustly with established interpretability literature and offers a viable RL-based substrate complementary to path-patching techniques. The results motivate future exploration of compositionality, scalability, and task-flexible circuit discovery mechanisms in neural architectures.