想象一下,一个现代语言模型(比如驱动这次对话的 AI)就像一个规模宏大、极其复杂的管弦乐团。在这个乐团内部,有成千上万名乐手(被称为“注意力头”,attention heads),他们演奏着不同的乐器。每一位乐手都会倾听其他乐手的演奏以及乐谱(输入文本),从而决定何时该奏响强音,何时该转为弱音,以及该击打出什么样的音符。
长期以来,科学家们一直试图理解每一位乐手究竟在做什么,但这些解释总是很模糊。他们可能会说:“这位乐手似乎很喜欢名词,”或者“那位乐手专注于句子的开头。”这些描述就像是通过说“这听起来很开心”或“这听起来很忧伤”来描述一场交响乐。这很有帮助,但它并没有告诉你音乐究竟是如何被制造出来的。
新方法:编写乐谱
这篇论文提出了一种理解这个乐团的不同方式。研究人员不再仅仅是描述音乐,而是尝试编写实际的乐谱(可执行代码),来告诉每一位乐手具体该做什么。
以下是他们操作的具体步骤:
- 观察乐手: 首先,他们观察 AI 模型阅读数千个句子的过程,并记录下每一位“乐手”(注意力头)是如何做出反应的。他们记录了每一位注意力头都在关注哪些词汇。
- 询问“影子作家”: 他们将这些记录交给另一个非常聪明的 AI(一个“程序合成”智能体),让它编写一个简单的 Python 程序,能够模拟这些反应。可以把这想象成要求一位影子作家去观察一位乐手的表演,并写下他所遵循的精确规则,例如:“如果前一个词是句号,则观察下一个词。如果该词是一个名字,则观察动词。”
- 测试影子作家: 他们将影子作家生成的代码运行在新的句子上,以观察其是否与原始 AI 的行为相匹配。他们使用了一种评分系统(类似于相似度测试)来查看代码与真实情况的接近程度。
- 大替换: 这是最令人兴奋的部分。他们将原始复杂的神经“乐手”从乐团中移走,并用他们刚刚编写的简单、人类可读的代码取而代代之。他们想看看乐团是否仍能保持原有的节奏。
他们的发现
结果非常成功:
- 许多乐手遵循简单的规则: 他们发现,AI 的大量注意力头实际上都在遵循非常逻辑化的符号规则。例如,有些头仅仅被设定为“观察句子的第一个词”,而另一些则是“寻找逗号”或“寻找开启新句子的单词”。
- 乐团依然在演奏: 当他们用这些简单的代码脚本替换掉高达 25% 到 40% 的复杂神经乐手时,乐团(AI 模型)几乎没有出现任何失误。模型回答问题或理解故事的能力几乎保持不变。
- 更大的模型更容易解码: 有趣的是,模型越大、越先进(例如 Llama-3B),就越容易找到这些简单的代码规则。看起来,随着模型规模的扩大,它们变得更加有序,每一位乐手都有一个非常具体且易于理解的任务。
总结
这篇论文表明,我们并不总是需要将 AI 视为一个神秘的“黑盒”。对于这些模型运作方式中的很大一部分,我们实际上可以用简单的、人类可读的计算机代码来取代复杂的、难以解释的数学运算。
这就像是意识到一个复杂的魔术其实并不是真正的魔法;它只是一套特定的指令,任何人都可以写下来并遵循。研究人员证明了,对于现代 AI 的许多组成部分,我们可以将“魔法”替换为“指令”,而不会破坏整场演出。
问题陈述
理解深度神经网络在算法层面的计算过程仍然是机器学习领域的一项重大挑战。以往的可解释性方法主要依赖于对神经元或特征进行意义赋予,其方式要么是自上而下的(通过针对人类定义概念的探测器),要么是自下而上的(通过输入/输出摘要)。虽然这些方法功能强大,但往往无法提供对神经计算的完整、形式化的描述。此外,用自然语言描述中间组件往往会导致解释具有歧义性或难以形式化。作者认为,需要一种媒介来弥合数十亿参数模型的复杂性与自然语言解释之间的鸿沟——这种媒介应当既具备人类可读性,又具备形式可验证性。
方法论
本文提出了一个基于**程序合成(program synthesis)**的框架,用于解释 Transformer 语言模型中注意力头(attention heads)的计算过程。其核心目标是寻找一个可执行的 Python 程序 (π),该程序能够直接将输入文本序列 (X) 映射到注意力矩阵 (A),从而有效地近似特定注意力头的逻辑。
该框架分为四个截然不同的步骤:
- 注意力图提取: 作者从目标模型在训练语料库(特别使用了结构简单性较高的 TinyStories 数据集)中的真实注意力激活矩阵中提取注意力图。他们对这些图进行过滤,仅保留前 2.5% 的注意力权重,以聚焦于显著的 Token 对交互。
- 程序合成: 一个辅助大语言模型(LLM)作为合成代理,在接收被提取的注意力模式和输入序列的提示后,负责生成能够重现这些模式的可执行 Python 函数。该合成代理可以使用 NumPy、spaely 和 NLTK 等库进行语言学和数值处理。
- 精炼与排序: 候选程序需经过语法验证,并使用 Jensen-Shannon 距离(JSD)根据其与真实注意力模式的匹配程度进行评分。表现最优的候选程序将进入反馈调节的精炼循环,其中合成代理通过对比真实模式与预测模式来生成修订后的程序。
- 选择: 通过在留出(held-out)数据序列上最大化相似度(通过交并比,即 IoU 衡量),为每个注意力头选择最优的代理程序。
核心贡献
- 程序化可解释性: 作者引入了一种方法,用符号化的、可执行的代码来替换神经注意力头,为自然语言描述提供了一种形式可验证的替代方案。
- 通过替换进行因果验证: 不同于以往仅观察相关性的工作,该方法通过在模型的正向传播过程中直接用合成程序替换学习到的神经组件,来验证解释的有效性。
- 可扩展性: 该方法被应用于四种不同的 Transformer 架构:BERT-Base、GPT-2-Small、TinyLlama-1.1B 和 Llama-3B,证明了该方法在不同模型规模和架构(双向与因果模型)下的适用性。
结果
- 高对齐度: 在所有测试的模型中,很大一部分注意力头都可以通过高精度的可执行程序进行近似。表现最佳的程序在 GPT-2 上实现了 69% 的平均 IoU 分数,在 TinyLlama 上为 74%,在 Llama-3B 上为 79%。
- 模型规模与架构: 作者观察到,自回归(解码器)模型的注意力头通常比双向(编码器)模型(如 BERT)更容易拟合。此外,拟合质量随模型规模的增大而提高,这表明较大的模型在其头部可能表现出更强的功能专业化。
- 因果忠实度: 当使用合成程序替换注意力头时,模型能够保持功能性能。具体而言,替换高达 25% 的注意力头,困惑度(perplexity)仅增加 16%。
- 下游性能: 用其最优拟合程序替换高达 30%–40% 的注意力头,并不会显著降低在下游问答基准测试(包括 HellaSwag, PIQA, SciQ, ARC-Easy, Social IQA 和 COPA)上的表现。
- 相关性: 程序与替换头部后的困惑度增加量之间存在强负相关性(Spearman r>0.9),这表明 IoU 是选择具有功能忠实度的程序的可靠指标。
意义与主张
本文声称,即使是在最先进的语言模型中,很大一部分注意力模式也可以用符号术语来理解。其主要意义在于证明了经过训练的神经组件可以通过可执行代码被替换为符号代理,且不会显著改变模型行为。这表明,现代基于语言模型的程序合成方法为接近“深度网络如何运作”这一更广泛的问题提供了一个可行的框架。
作者将这项工作定位为迈向 Transformer 完整符号化表征的一步,即模型的行为可以通过逻辑而非权重操纵来进行推理、修改和测试。然而,论文在局限性方面保持了谦逊,承认目前尚无模型的注意力头被完全表征(许多头部仍处于低 IoU 状态),且目前的程序库虽然有效,但尚未达到完备。这项工作为实现 Transformer 架构的完整符号蒸馏奠定了基础。
每周获取最佳 computer science 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。