Explaining Attention with Program Synthesis: Reverse-Engineering Attention Heads into Readable Python Code

Forum topic · 小凯 · 2026-06-18

Summary

This post discusses the paper 'Explaining Attention with Program Synthesis' by Amiri Hayes, Belinda Li, and Jacob Andreas (MIT CSAIL), which uses large language models to reverse-engineer the behavior of Transformer attention heads into executable, human-readable Python programs. The pipeline extracts salient attention patterns from models such as GPT-2, TinyLlama-1.1B, Llama-3B, and BERT-base, prompts Claude Sonnet to generate candidate Python functions, evaluates them with Jensen-Shannon distance, and iteratively refines the best candidates. About 4,000 candidate programs were synthesized (roughly $150 in API costs), yielding a library of 1,664 programs—one per attention head. Reported results include average IoU above 75% against real attention matrices on TinyStories across most models (up to 99% for positional heads), and interchange-intervention tests showing that replacing 25% of attention heads with their synthetic programs increases perplexity by only 16% while largely preserving downstream QA performance. The post also covers limitations (many heads remain poorly explained, with IoU below 40%), implications for AI safety and model compression, and the broader vision of full symbolic characterization of neural networks.

This forum post introduces and explains the research paper 'Explaining Attention with Program Synthesis' by Amiri Hayes, Belinda Li, and Jacob Andreas (MIT CSAIL, arXiv:2606.19317).

Overview

The paper's core idea: since attention head behavior is deterministic, there should in principle exist a program that replicates it. The authors use an LLM (Claude Sonnet) as an 'interpreter' that generates executable Python functions mimicking individual attention heads, mapping neural computation to symbolic, human-readable code.

Method (four steps)

1. Extract attention patterns: run the target model on large text corpora, record attention matrices, keep only the most salient 2.5% of attention edges. 2. Generate hypothesis programs: summarize patterns in text, prompt Claude Sonnet to write Python functions that take input text and output predicted attention matrices (using NumPy, spaCy, NLTK). 3. Evaluate and rank: measure similarity to real attention matrices with Jensen-Shannon distance (JSD); select top candidates on a validation set. 4. Refine with feedback: identify best/worst examples, generate structured error feedback, and iteratively improve programs.

Example synthesized programs include a sentence-boundary detection head (attending to tokens following ., !, ?, or newlines) and a simple previous-token (positional) head.

Scale and Results

~4,000 candidate programs generated; 1,664 final programs (one per head); ~$150 in API costs (35M input + 3.5M output tokens).
Attention alignment (IoU on TinyStories): 75%+ on average for GPT-2, TinyLlama-1.1B, and Llama-3B; BERT-base positional heads reach up to 99%.
Interchange interventions: replacing 25% of attention heads with their synthetic programs increases perplexity by only 16%; downstream QA tasks (ARC, OpenBookQA) remain largely intact. Replacing 30–40% of heads partially preserves task ability, suggesting a pruning-like effect.
Layer-wise analysis: low and high layers of BERT-base and GPT-2 contain many well-explained heads (IoU > 0.6); simple positional heads score highest, while semantically complex heads score lower.

Limitations

No model has all attention heads fully explained; a substantial fraction of heads score below 40% IoU.
These poorly explained heads likely cause the performance degradation at high replacement rates; the authors propose multi-round refinement and richer feedback as future work.

Significance

The work demonstrates that a meaningful portion of attention behavior in state-of-the-art language models can be expressed as symbolic, executable code—advancing mechanistic interpretability toward full symbolic characterization. Potential applications include AI safety auditing (verifying behavioral patterns before deployment), model compression (replacing redundant heads with simple logic), and more transparent AI education.

References

1. Hayes, A., Li, B., & Andreas, J. (2026). Explaining Attention with Program Synthesis. arXiv:2606.19317. 2. Austin, J., et al. (2021). Program synthesis with large language models. arXiv:2108.07732. 3. Bills, S., et al. (2023). Language models can explain neurons in language models. OpenAI Blog. 4. Geiger, A., et al. (2021). Causal abstraction for the interpretability of deep learning models. NeurIPS. 5. Nanda, N., et al. (2023). Progress measures for grokking via mechanistic interpretability. arXiv:2304.14997. 6. Michaud, E. J., et al. (2024). Opening the AI black box: Program synthesis via mechanistic interpretability. arXiv:2402.05110. 7. Vaswani, A., et al. (2017). Attention is all you need. NeurIPS. 8. Zhang, P., et al. (2024). TinyLlama: An open-source small language model. arXiv:2401.02385.