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

Forum topic · 小凯 · 2026-06-18

Summary

A forum post on zhichai.net discusses the MIT CSAIL paper 'Explaining Attention with Program Synthesis' (Hayes, Li, Andreas; arXiv:2606.19317), which uses large language models to reverse-engineer the behavior of Transformer attention heads into executable, human-readable Python programs. The pipeline works in four steps: extract salient attention patterns from a target model, prompt an interpreter LLM (Claude Sonnet) to generate candidate Python functions, rank them by Jensen-Shannon distance against real attention matrices, and iteratively refine the best candidates with error feedback. Across GPT-2, TinyLlama-1.1B, and Llama-3B, synthesized programs achieve over 75% average IoU alignment with real attention patterns, with some BERT heads (e.g., previous-token heads) reaching 99%. Causal interchange interventions show that replacing 25% of attention heads with program outputs increases perplexity by only about 16% while largely preserving downstream QA performance. Roughly 1,664 final programs were synthesized at an API cost of about $150. Limitations include many heads scoring below 40% IoU, especially in middle layers and semantic heads. The post argues this work is a step toward symbolic transparency, AI safety auditing, and model compression.

Explaining Attention with Program Synthesis — Paper Overview

This zhichai.net forum post is a detailed walkthrough (in a Feynman-style, beginner-friendly tone) of the MIT CSAIL paper "Explaining Attention with Program Synthesis" by Amiri Hayes, Belinda Li, and Jacob Andreas (arXiv:2606.19317).

One-line takeaway: use a large language model to generate executable Python programs that reverse-engineer the behavior of attention heads in Transformers, mapping neural computation to transparent symbolic code.

Key points

Background

Attention is the core mechanism of Transformers: each token emits Query, Key, and Value vectors, with scores computed as Attention(Q, K, V) = softmax(QK^T / √d_k) · V.
Models analyzed: GPT-2 (12 layers × 12 heads = 144 heads), TinyLlama-1.1B (22 × 32 = 704), Llama-3B (26 × 32 = 832).
Motivation: interpretability — replacing opaque neural computation with human-understandable symbolic descriptions (mechanistic interpretability).

Method (four-step pipeline)

Extract attention patterns

Generate hypothesis programs

Evaluate and rank

Refine with feedback

Scale: ~4,000 candidate programs; final library of 1,664 programs (one per head); API cost ≈ $150 (35M input + 3.5M output tokens). Programs use NumPy, spaCy, NLTK.

Example synthesized programs

Sentence-boundary head: at position i, attend strongly (0.8) to tokens j that begin sentences (after ., !, ?, or newline, or the first token).
Previous-token head: position i attends (0.9) to position i−1, modeling positional offsets.

Results

Alignment (IoU): average IoU above 75% for GPT-2, TinyLlama-1.1B, and Llama-3B; BERT-base positional heads reach up to 99%.
Causal interchange interventions: replacing 25% of attention heads with program outputs raises perplexity by only ~16%; downstream QA (ARC, OpenBookQA) performance is largely preserved. Replacing 30–40% of heads still retains some task capability in some models.
Layer patterns: BERT low layers (0–3) and high layers (9–11) contain many well-explained heads (IoU > 0.6); middle layers are harder. Positional heads are easiest to explain; semantic heads score lower.

Limitations

No model has all attention heads fully explained; many heads score below 40% IoU and likely drive performance degradation at higher replacement rates. Stronger synthesis strategies (multi-round refinement, richer feedback) are needed.

Implications

AI safety: auditable code enables detecting undesirable attention behavior before deployment.
Compression: replacing heads with simple programs with limited loss suggests redundancy in attention computation.
Symbolic transparency: a step toward models whose behavior can be read, modified, and tested like code — a full symbolic characterization of neural computation remains the long-term goal.

References cited in the post

Hayes, Li & Andreas (2026), arXiv:2606.19317
Vaswani et al. (2017), *Attention Is All You Need*, NeurIPS
Bills et al. (2023), *Language models can explain neurons*
Michaud et al. (2024), arXiv:2402.05110
Zhang et al. (2024), TinyLlama, arXiv:2401.02385