[论文解读] Interpreting Attention Layer Outputs with Sparse Autoencoders

Connor Kissane, Robert Krzyzanowski+3 more|arXiv (Cornell University)|Jun 25, 2024

Ferroelectric and Negative Capacitance Devices被引用 4

一句话总结

该论文将稀疏自编码器（SAEs）应用于Transformer模型注意力层的输出，表明SAEs能够提取出对应于长距离上下文、短距离上下文以及归纳模式的稀疏且可解释的特征。作者提出了递归直接特征归因（RDFA）方法，可将特征来源追溯至输入标记，从而实现对注意力机制与电路行为的详细分析，包括解析冗余归纳头的作用并验证间接宾语识别电路的功能。

ABSTRACT

Decomposing model activations into interpretable components is a key open problem in mechanistic interpretability. Sparse autoencoders (SAEs) are a popular method for decomposing the internal activations of trained transformers into sparse, interpretable features, and have been applied to MLP layers and the residual stream. In this work we train SAEs on attention layer outputs and show that also here SAEs find a sparse, interpretable decomposition. We demonstrate this on transformers from several model families and up to 2B parameters. We perform a qualitative study of the features computed by attention layers, and find multiple families: long-range context, short-range context and induction features. We qualitatively study the role of every head in GPT-2 Small, and estimate that at least 90% of the heads are polysemantic, i.e. have multiple unrelated roles. Further, we show that Sparse Autoencoders are a useful tool that enable researchers to explain model behavior in greater detail than prior work. For example, we explore the mystery of why models have so many seemingly redundant induction heads, use SAEs to motivate the hypothesis that some are long-prefix whereas others are short-prefix, and confirm this with more rigorous analysis. We use our SAEs to analyze the computation performed by the Indirect Object Identification circuit (Wang et al.), validating that the SAEs find causally meaningful intermediate variables, and deepening our understanding of the semantics of the circuit. We open-source the trained SAEs and a tool for exploring arbitrary prompts through the lens of Attention Output SAEs.

研究动机与目标

• 证明稀疏自编码器（SAEs）能够将注意力层输出分解为稀疏且可解释的特征。
• 开发并应用递归直接特征归因（RDFA），将SAE特征激活追溯至输入标记及上游组件。
• 解决关于多义注意力头以及看似冗余的归纳头功能多样性的开放问题。
• 验证SAE是否能识别出已知电路（如间接宾语识别电路）中具有因果意义的中间变量。
• 为社区提供开源工具与SAE模型，用于分析注意力机制。

提出的方法

• 在Transformer层中所有注意力头的拼接预线性注意力输出（z_cat）上训练SAEs，以重建高维激活。
• 使用基于权重的头归因方法，将学习到的SAE特征与特定注意力头关联，缓解多义性问题。
• 实现递归直接特征归因（RDFA），递归地将SAE特征激活分解为上游组件（包括残差流特征和输入标记）的贡献。
• 利用Transformer的线性结构，对z_cat进行线性分解，以分离出注意力模式、值投影和残差流特征的贡献。
• 将注意力输出上的SAEs与残差流SAEs结合，实现从最终SAE特征到输入标记的端到端归因。
• 发布一个工具链接：https://robertzk.github.io/circuit-explorer，用于可视化和探索任意提示下的特征电路。

Figure 1 : Overview. We train Sparse Autoencoders (SAEs) on $\textbf{z}_{\text{cat}}$ , the attention layer outputs pre-linear, concatenated across all heads. The SAEs extract linear directions that correspond to concepts in the model, giving us insight into what attention layers learn in practice.

实验结果

研究问题

• RQ1 注意力层学习到的特征类型是什么？能否被有意义地分解为可解释的特征？
• RQ2 注意力头如何贡献于模型行为？在多义或冗余配置中，它们各自扮演什么角色？
• RQ3 对注意力输出应用SAE能否帮助揭示为何模型存在大量看似冗余的归纳头？
• RQ4 注意力层中的SAE特征是否对应于已知电路（如间接宾语识别电路）中具有因果意义的中间变量？
• RQ5 递归归因技术（如RDFA）能否将某一特征的来源追溯至特定输入标记及上游组件？

主要发现

• SAEs成功地将注意力层输出分解为稀疏且可解释的特征，包括长距离上下文、短距离上下文以及归纳模式。
• 在GPT-2 Small中，至少90%的注意力头具有多义性，其功能随上下文变化而执行多个无关任务。
• RDFA显示，看似冗余的归纳头实际上分别专注于长前缀或短前缀的归纳，从而解决了关键的开放性问题。
• 通过SAE特征分析，间接宾语识别电路的机制得到更深入理解，这些特征捕捉到了具有因果意义的中间变量。
• 递归归因证实，模型通过注意力连接将关于'Dave'的信息从'is'标记传递至'isn’t'标记。
• 作者发布了训练好的SAEs及一个电路探索工具，使社区能够跨提示探索注意力特征。

(a)