Characterizing Attention Heads via Program Synthesis: Toward Scalable Mechanistic Interpretability

Motivation A primary aim of interpretability research is to understand, and ultimately intervene on, the internal computations of deep networks. Mechanistic analyses have successfully reverse-engineered the circuits underlying narrow behaviors such as indirect object identification, entity tracking, and subject–verb agreement. However, each of these case studies demands substantial time and effort from researchers, so the methodology does not scale, at least not trivially, to more complex settings. Recently, Hayes et al. ([1]) proposed a promising alternative based on program synthesis: replacing a model's computation with executable Python programs that replicate it. While this is a valuable step toward scalable interpretability, it remains far from a complete solution.

Proposed Research I propose to extend this line of work along two complementary directions: an engineering-oriented direction that refines the existing method, and an exploratory direction that extends program synthesis from the QK circuit to the OV circuit of attention heads.

Direction 1: Strengthening the program synthesis pipeline The proposed method can be substantially improved to generate higher-quality programs. Concrete avenues include running additional iterations of program generation and refinement, designing better prompting strategies, introducing complementary metrics for assessing program quality, and evaluating the efficacy of open-weight models as program generators. Beyond raising performance, these experiments would clarify which components of the pipeline matter most, informing the exploratory work in Direction 2.

Direction 2: Extending program synthesis to the OV circuit Generating programs for the QK circuit of an attention head is comparatively straightforward, because the target output is a well-defined n_tokens × n_tokens matrix over the input tokens. The OV circuit, in contrast, maps a residual-stream vector to another residual-stream vector, so there is no well-defined symbolic object for a program to generate. The central question of this direction is therefore: what is the ideal interpretable basis in which to ground the OV computation of an attention head? As a starting point, I will use the logit lens and the Jacobian lens [2] to extract the information encoded by the OV circuit, and predict its computation by juxtaposing this information with the surrounding context. Since neither lens yields useful information in early layers, this approach is expected to work primarily for middle and later layers. A complementary approach is to use trained models to read out the contents of neural activations, as proposed in [3]. I suspect the eventual solution will combine both approaches. Regardless of the specific instantiation, I believe program synthesis offers a strong opportunity to interpret the internal computations of language models in a fully unsupervised manner.

Expected Outcomes This project is expected to deliver (i) a program synthesis–based method that fully characterizes an attention head, covering both its QK and OV circuits, and (ii) a detailed account of the attention heads of several open-source models, together with a cross-model comparison.

[1] Hayes et al, “Explaining Attention with Program Synthesis”, 2026. [2] Gurnee et al, “Verbalizable Representations Form a Global Workspace in Language Models”, 2026. [3] Huang et al, “Inversionview: A general-purpose method for reading information from neural activations”, 2024.

Behavioral evaluations alone can't be trusted for advanced AI systems, which may game them through deception or sandbagging; we need ways to verify what models are actually computing. Mechanistic interpretability offers this, but is currently bottlenecked on manual researcher effort. This project attacks that bottleneck: synthesizing executable programs that replicate model internals yields interpretations that are automated and verifiable, since program faithfulness can be directly measured. Success would be a concrete step toward automated auditing of frontier models for unsafe mechanisms, deceptive circuitry, or hazardous knowledge.

I expect the mentee to take the lead on this project and work fairly autonomously. While I'll be available for guidance and support throughout, the mentee should drive the research, formulating well-motivated hypotheses and implementing experiments to test them. I'm therefore looking for someone who is genuinely curious about how LMs work and has some intuitions about their internal mechanisms. Please also mention in your application whether you would be open to working with another mentee on this project.

• Highly proficient in Python and PyTorch; you'll be implementing and debugging your own experiments fairly independently.
• Completed at least one hands-on mechanistic interpretability investigation using a library such as nnsight, TransformerLens, or pyvene (self-directed projects, paper replications, and ARENA-style coursework all count).
• Familiar with common mechanistic interpretability methods such as the logit lens, activation patching, and distributed alignment search (DAS).
• Experience building LLM pipelines (prompting, API calls, evaluation harnesses).
• Nice to have: Familiarity with Tracr.