Explaining Attention with Program Synthesis
2026-06-17 • Machine Learning
Machine LearningArtificial Intelligence
AI summaryⓘ
The authors created a method to understand parts of language models called attention heads by summarizing their behavior with simple Python programs. They tested this on models like GPT-2 and Llama and found that a small set of generated programs could mimic how these attention heads work most of the time. Replacing some attention heads with these programs did not hurt the models' performance much. This approach helps make complex neural networks more understandable by using easy-to-read code that approximates their inner workings.
attention headstransformer modelslanguage modelsGPT-2Llamaexecutable programssymbolic interpretationperplexityprogram synthesisattention matrices
Authors
Amiri Hayes, Belinda Li, Jacob Andreas
Abstract
A longstanding goal of research on interpretable deep learning is to replace opaque neural computations with human-meaningful symbolic descriptions. In this paper, we propose an approach for approximating the behavior of components of deep networks with executable programs. We focus on attention heads in transformer language models. For a given head, we first compute its associated attention matrices on a collection of randomly selected training examples. Next, we prompt a pre-trained language model with a summary of these matrices, and instruct it to generate a set of Python programs that can reproduce the associated attention patterns given only text from the input sentence. Finally, we re-rank programs according to how well our final set of programs predict behavior on held-out inputs. We demonstrate that a set of fewer than 1,000 such generated programs can reproduce the attention patterns of heads in GPT-2, TinyLlama-1.1B, and Llama-3B, achieving an average Intersection-over-Union similarity above 75% on TinyStories. Moreover, the best-fit programs can replace neural attention heads without substantially affecting model behavior: replacing 25% of attention heads with programmatic surrogates across the three models incurs only a 16% average perplexity increase, while maintaining performance on a variety of downstream question answering benchmarks. This work contributes a scalable pipeline for reverse-engineering attention heads in transformer models using human-readable, executable code, advancing a path toward symbolic transparency in neural models.