🤖 AI Summary
This work addresses the limited generalizability of traditional mechanistic interpretability methods, which often rely on task-specific customization. It proposes a novel, general-purpose circuit discovery framework that integrates reinforcement learning with causal intervention. Specifically, a PPO agent performs zero-ablation operations across all 144 attention heads of GPT-2 Small and automatically identifies critical heads using a contrastive reward signal—eliminating the need for task-specific analysis pipelines and enabling cross-task transfer. Experimental results demonstrate that the method achieves near-oracle performance both on training tasks and unseen tasks: in the absence of task-specific signals, the best five runs recover 96% of the oracle upper bound and accurately identify both key and redundant attention heads previously documented in the literature.
📝 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 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.