💻 computer science

Explaining Attention with Program Synthesis

This paper proposes a scalable pipeline that uses program synthesis to approximate transformer attention heads with human-readable, executable Python programs, demonstrating that these symbolic surrogates can accurately reproduce attention patterns and replace a significant portion of neural heads with minimal impact on model performance.

Original authors: Amiri Hayes, Belinda Li, Jacob Andreas

Published 2026-06-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Amiri Hayes, Belinda Li, Jacob Andreas

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a modern language model (like the AI behind this chat) as a massive, complex orchestra. Inside this orchestra, there are thousands of musicians (called "attention heads") playing different instruments. Each musician listens to the other musicians and the sheet music (the input text) to decide when to play loud, when to play soft, and what notes to hit.

For a long time, scientists have tried to understand what each musician is doing, but the explanations have been vague. They might say, "This musician seems to like nouns," or "That one focuses on the beginning of sentences." These are like describing a symphony by saying, "It sounds happy," or "It sounds sad." It's helpful, but it doesn't tell you exactly how the music is made.

The New Approach: Writing the Sheet Music

This paper proposes a different way to understand the orchestra. Instead of just describing the music, the researchers tried to write the actual sheet music (executable code) that tells each musician exactly what to do.

Here is how they did it, step-by-step:

Watching the Musicians: First, they watched the AI model read thousands of sentences and recorded exactly how each "musician" (attention head) reacted. They noted which words each head paid attention to.
Asking a Ghost Writer: They took these recordings and asked a different, very smart AI (a "program synthesis" agent) to write a simple Python program that could mimic those reactions. Think of this as asking a ghost writer to look at a musician's performance and write down the exact rules they were following, like: "If the previous word is a period, look at the next word. If the word is a name, look at the verb."
Testing the Ghost Writer: They took the code the ghost writer produced and ran it on new sentences to see if it matched the original AI's behavior. They used a scoring system (like a similarity test) to see how close the code was to the real thing.
The Big Swap: This is the most exciting part. They took the original, complex neural "musicians" out of the orchestra and replaced them with the simple, human-readable code they just wrote. They wanted to see if the orchestra would still sound the same.

What They Found

The results were surprisingly successful:

Many Musicians Follow Simple Rules: They found that a huge number of the AI's attention heads were actually following very logical, symbolic rules. For example, some heads were just programmed to "look at the first word of the sentence," while others were "looking for commas" or "finding words that start new sentences."
The Orchestra Still Plays: When they replaced up to 25% to 40% of the complex neural musicians with these simple code scripts, the orchestra (the AI model) barely missed a beat. The model's ability to answer questions or understand stories remained almost exactly the same.
Bigger Models are Easier to Decode: Interestingly, the larger and more advanced the model was (like Llama-3B), the easier it was to find these simple code rules. It seems that as models get bigger, they become more organized, with each musician having a very specific, easy-to-understand job.

The Takeaway

The paper shows that we don't always need to treat AI as a mysterious "black box." For a significant portion of how these models work, we can actually replace the complex, unexplainable math with simple, human-readable computer code.

It's like realizing that a complex magic trick isn't actually magic at all; it's just a specific set of instructions that anyone could write down and follow. The researchers proved that for many parts of modern AI, we can swap the "magic" for "instructions" without breaking the show.

Problem Statement
Understanding the computations performed by deep neural networks in algorithmic terms remains a significant challenge in machine learning. Previous approaches to interpretability have largely relied on assigning meaning to neurons or features either top-down (via probes for human-defined concepts) or bottom-up (via input/output summaries). While powerful, these methods often fail to provide a full, formal description of neural computation. Furthermore, labeling intermediate components with natural language descriptions can result in interpretations that are ambiguous or difficult to formalize. The authors posit that there is a need for a medium that bridges the gap between the complexity of billion-parameter models and natural language explanations—one that is both human-readable and formally verifiable.

Methodology
The paper proposes a framework based on program synthesis to explain the computations of attention heads in transformer language models. The core objective is to find an executable Python program ( $\pi$ ) that can map an input text sequence ( $X$ ) directly to an attention matrix ( $A$ ), effectively approximating the logic of a specific attention head.

The framework proceeds in four distinct steps:

Attention Map Extraction: The authors extract ground-truth attention activation matrices from a target model across a corpus of training sequences (specifically using the TinyStories dataset for its structured simplicity). They filter these maps to isolate the top 2.5% of attention weights to focus on salient token-pair interactions.
Program Synthesis: An auxiliary Large Language Model (LLM), acting as a synthesis agent, is prompted with the extracted attention patterns and input sequences. The agent is tasked with generating executable Python functions that reproduce these patterns. The synthesis agent has access to libraries such as NumPy, spaCy, and NLTK for linguistic and numerical processing.
Refinement and Ranking: Candidate programs are validated for syntax and scored against real attention patterns using the Jensen-Shannon Distance (JSD). The top candidates undergo a feedback-conditioned refinement loop where the synthesis agent contrasts real and predicted patterns to produce revised programs.
Selection: The optimal proxy program for each head is selected by maximizing similarity (measured via Intersection over Union, or IoU) on held-out data sequences.

Key Contributions

Programmatic Interpretability: The authors introduce a method to replace neural attention heads with symbolic, executable code, offering a formally verifiable alternative to natural language descriptions.
Causal Validation via Replacement: Unlike prior work that only observes correlations, this approach validates explanations by directly substituting learned neural components with synthesized programs during the model's forward pass.
Scalability: The method is applied to four distinct transformer architectures: BERT-Base, GPT-2-Small, TinyLlama-1.1B, and Llama-3B, demonstrating the approach's applicability across different model sizes and architectures (bidirectional and causal).

Results

High Alignment: A substantial fraction of attention heads across all tested models can be approximated by executable programs with high accuracy. The best-performing programs achieve mean IoU scores of 69% for GPT-2, 74% for TinyLlama, and 79% for Llama-3B.
Model Scale and Architecture: The authors observe that attention heads in autoregressive (decoder) models are generally easier to fit than those in bidirectional (encoder) models like BERT. Additionally, fit quality increases with model scale, suggesting that larger models may exhibit greater functional specialization among their heads.
Causal Fidelity: When replacing attention heads with their synthesized programs, the models maintain functional performance. Specifically, up to 25% of heads can be replaced with only a 16% increase in perplexity.
Downstream Performance: Replacing up to 30–40% of attention heads with their best-fit programs does not significantly degrade performance on downstream question-answering benchmarks (including HellaSwag, PIQA, SciQ, ARC-Easy, Social IQA, and COPA).
Correlation: There is a strong negative correlation (Spearman $r > 0.9$ ) between the IoU similarity of a program and the perplexity increase upon head replacement, indicating that IoU is a reliable proxy for selecting functionally faithful programs.

Significance and Claims
The paper claims that even in state-of-the-art language models, a substantial fraction of attention patterns can be understood in symbolic terms. The primary significance lies in the demonstration that trained neural components can be replaced by symbolic surrogates via executable code without substantially altering model behavior. This suggests that modern LM-driven program synthesis methods offer a viable framework for approaching the broader question of how deep networks operate.

The authors position this work as a step toward a complete symbolic characterization of transformers, where model behavior can be reasoned about, modified, and tested directly through logic rather than weight manipulation. However, the paper remains modest regarding its limitations, acknowledging that no model's attention heads are currently fully characterized (many heads still achieve low IoU scores) and that the current program library, while effective, is not yet exhaustive. The work establishes a foundation for future research aimed at achieving a complete symbolic distillation of transformer architectures.

More like this