Explaining Attention with Program Synthesis

Introduction. Understanding the computations performed by deep networks in algorithmic terms is a long-standing problem in machine learning [14, 36, 24]. Past work approaching this problem has attempted to assign meaning to neurons or other distributed features in a top-down way (by training probes for human-defined concepts of interest [18, 28]), or in a bottom-up way (by generating summaries of the inputs that activate a feature, or the outputs that the feature promotes [23, 17, 2]). While powerful, these methods stop short of providing a full, formal description of neural computation. Current work addresses this by labeling intermediate components of deep networks with natural language descriptions [4, 8], whose interpretations may themselves be ambiguous or hard to formalize. In this paper, we develop an alternative approach to deep network interpretability based on program synthesis. To explain the computation performed by some component of a deep network, we search for a piece of executable code that approximates the computation performed by that component.

Discussion / Conclusion. The results presented in this work demonstrate that a symbolic approach to mechanistic interpretability is not only feasible but could yield direct proxies for complex neural activations. Model-level analyses in Figure 3 show that executable programs can achieve up to 99% mean IoU similarity against observed attention patterns. Our evaluation in Figure 5 also demonstrates that there is a clear negative correlation between the mean IoU similarity of a program and head and the perplexity increase as a result of replacing the head with its program. Examining the model’s performance trends on downstream question-answering evaluations, we find that we can replace up to 30-40% of all attention heads with their highest-similarity programs π without losing task ability. This indicates that our method can almost completely capture the functional role of up to 30-40% of all LM heads. Limitations There is significant room for improvement in expanding the diversity and complexity of these hypotheses. Currently, no model’s attention heads are fully characterized: a substantial fraction of heads achieve IoU scores below 40%, and we attribute the question-answering degradation observed at high replacement levels in Figure 6 primarily to these poorly fit heads.