MechRL: Reinforcement Learning Agents Perform Circuit Discovery for Mechanistic Interpretability

• URL: http://arxiv.org/abs/2605.26343v1
• Date: Mon, 25 May 2026 21:32:57 GMT
• Title: MechRL: Reinforcement Learning Agents Perform Circuit Discovery for Mechanistic Interpretability
• Abstract summary: We recast circuit discovery as a reinforcement-learning problem.<n>An agent operates over the 144 attention heads of GPT-2 small as a discrete action space.<n>A single PPO policy attains the per-episode oracle on both training tasks and on a held-out third task.
• Score: 0.0
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Pitfalls and Remedies for Multi-Task Bayesian Optimization [10.431916242473045]
  Bayesian optimization routinely warms a target experiment with data from related source tasks.<n>We revisit this default in a controlled setting and find that it misestimates the cross-task correlation.<n>We trace the failure to two independent structural mechanisms.
  arXiv  Detail & Related papers  (2026-07-10T03:38:36Z)
• DRDN: Decoupled Representation Dynamic Network for From-Scratch ViT Class-Incremental Learning [6.097306160389619]
  Dynamic expansion methods for class-incremental learning (CIL) protect task-specific knowledge by growing dedicated tokens or computationworks.<n>We identify two intertwined challenges: cross-task confusion from sequential training on predominantly current-task data, and under-optimized shared representations in the backbone.<n>We propose the Decoupled Representation Dynamic Network (DRDN), which addresses these challenges via two mechanisms.
  arXiv  Detail & Related papers  (2026-07-02T02:57:58Z)
• Dissecting Jet-Tagger Through Mechanistic Interpretability [0.0]
  We perform a mechanistic interpretability analysis of the Particle Transformer architecture, trained on the Top Quark Tagging reference dataset.<n>We identify the computational circuit responsible for jet classification and characterizing the physical content of its internal representations.
  arXiv  Detail & Related papers  (2026-05-11T02:11:47Z)
• Causally Sufficient and Necessary Feature Expansion for Class-Incremental Learning [10.618608207545153]
  Current expansion-based methods for Class Incremental Learning (CIL) effectively catastrophic forgetting by freezing old features.<n> spurious feature correlations are the main cause of this collision, manifesting in two scopes.<n>We propose a Probability of Necessity and Sufficiency (PNS)-based regularization method to guide feature expansion in CIL.
  arXiv  Detail & Related papers  (2026-03-10T03:33:33Z)
• TopoCurate:Modeling Interaction Topology for Tool-Use Agent Training [53.93696896939915]
  Training tool-use agents typically rely on Supervised Fine-Tuning (SFT) on successful trajectories and Reinforcement Learning (RL) on pass-rate-selected tasks.<n>We propose TopoCurate, an interaction-aware framework that projects multi-trial rollouts from the same task into a unified semantic quotient topology.<n>TopoCurate achieves consistent gains of 4.2% (SFT) and 6.9% (RL) over state-of-the-art baselines.
  arXiv  Detail & Related papers  (2026-03-02T10:38:54Z)
• Search-R2: Enhancing Search-Integrated Reasoning via Actor-Refiner Collaboration [49.9937230730202]
  We propose Search-R2, a novel Actor-Refiner collaboration framework that enhances reasoning through targeted intervention.<n>Our approach decomposes the generation process into an Actor, which produces initial reasoning trajectories.<n>We show that Search-R2 consistently outperforms strong RAG and RL-based baselines across model scales.
  arXiv  Detail & Related papers  (2026-02-03T15:32:09Z)
• On the Paradoxical Interference between Instruction-Following and Task Solving [50.75960598434753]
  Instruction following aims to align Large Language Models (LLMs) with human intent by specifying explicit constraints on how tasks should be performed.<n>We reveal a counterintuitive phenomenon: instruction following can paradoxically interfere with LLMs' task-solving capability.<n>We propose a metric, SUSTAINSCORE, to quantify the interference of instruction following with task solving.
  arXiv  Detail & Related papers  (2026-01-29T17:48:56Z)
• Emergence of Minimal Circuits for Indirect Object Identification in Attention-Only Transformers [0.10152838128195467]
  We train small, attention-only transformers from scratch on a symbolic version of the Indirect Object Identification task.<n>A single-layer model with only two attention heads achieves perfect IOI accuracy, despite lacking residuals and normalization layers.<n>A two-layer, one-head model achieves similar performance by composing information across layers through query-value interactions.
  arXiv  Detail & Related papers  (2025-10-28T22:25:19Z)
• Disentangling Recall and Reasoning in Transformer Models through Layer-wise Attention and Activation Analysis [3.1526281887627587]
  Distinguishing recall from reasoning is crucial for predicting model generalization.<n>We use controlled datasets of synthetic linguistic puzzles to probe transformer models at the layer, head, and neuron level.<n>Our results provide the first causal evidence that recall and reasoning rely on separable but interacting circuits in transformer models.
  arXiv  Detail & Related papers  (2025-10-03T04:13:06Z)
• Task-Agnostic Detector for Insertion-Based Backdoor Attacks [53.77294614671166]
  We introduce TABDet (Task-Agnostic Backdoor Detector), a pioneering task-agnostic method for backdoor detection. TABDet leverages final layer logits combined with an efficient pooling technique, enabling unified logit representation across three prominent NLP tasks. TABDet can jointly learn from diverse task-specific models, demonstrating superior detection efficacy over traditional task-specific methods.
  arXiv  Detail & Related papers  (2024-03-25T20:12:02Z)
• ForkMerge: Mitigating Negative Transfer in Auxiliary-Task Learning [59.08197876733052]
  Auxiliary-Task Learning (ATL) aims to improve the performance of the target task by leveraging the knowledge obtained from related tasks. Sometimes, learning multiple tasks simultaneously results in lower accuracy than learning only the target task, known as negative transfer. ForkMerge is a novel approach that periodically forks the model into multiple branches, automatically searches the varying task weights.
  arXiv  Detail & Related papers  (2023-01-30T02:27:02Z)
• ReAct: Temporal Action Detection with Relational Queries [84.76646044604055]
  This work aims at advancing temporal action detection (TAD) using an encoder-decoder framework with action queries. We first propose a relational attention mechanism in the decoder, which guides the attention among queries based on their relations. Lastly, we propose to predict the localization quality of each action query at inference in order to distinguish high-quality queries.
  arXiv  Detail & Related papers  (2022-07-14T17:46:37Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.