程序合成解释Transformer注意力机制

Explaining Attention with Program Synthesis

精选理由

这篇论文用Python程序解释了注意力头怎么工作,还能直接用程序替换掉原始头,精度很高,想看模型内部机制的可以读。

AI 摘要

研究者提出用程序合成方法反向工程Transformer注意力头。他们先计算注意力矩阵,再让预训练语言模型生成Python程序来重现注意力模式。在GPT-2、TinyLlama-1.1B和Llama-3B上,不到1000个程序实现了平均IoU>75%。替换25%的注意力头仅导致16%的困惑度增加,并在下游问答基准上保持性能。

原文 · arXiv cs.LG

Explaining Attention with Program Synthesis

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.