Explaining Attention with Program Synthesis

2026-06-17

AI Summary

This work proposes a program synthesis–based approach to explain the decision behavior of Transformer attention heads through human-readable Python programs. By observing attention matrices on randomly sampled inputs, the method leverages a pretrained language model to generate programs that reproduce observed attention patterns and validates their generalization on a held-out set. It represents the first large-scale neurosymbolic reverse engineering of numerous attention heads, achieving an average Intersection-over-Union (IoU) exceeding 75% on TinyStories. Replacing 25% of attention heads with synthesized programs results in only a 16% average increase in perplexity across three model variants, while downstream question-answering performance remains largely unaffected, substantially enhancing model interpretability and transparency.

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.

Problem

• interpretable deep learning
• attention heads
• program synthesis
• symbolic transparency
• transformer models

Innovation

• program synthesis
• attention heads
• interpretable deep learning
• symbolic transparency
• executable programs