Sparse Autoencoders Enable Scalable and Reliable Circuit Identification in Language Models
作者: Charles O'Neill, Thang Bui
分类: cs.CL, cs.LG
发布日期: 2024-05-21
💡 一句话要点
提出基于稀疏自编码器的电路识别方法，提升语言模型可解释性
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 语言模型 可解释性 电路识别 稀疏自编码器 注意力机制
📋 核心要点
- 现有电路识别方法计算复杂度高，且对超参数敏感，难以应用于大型语言模型。
- 利用稀疏自编码器学习注意力头输出的离散表示，通过正负例对比，直接识别参与特定计算的注意力头。
- 实验表明，该方法在电路恢复方面优于现有方法，且运行时间大幅缩短，仅需少量样本即可学习稳健表示。
📝 摘要（中文）
本文提出了一种高效且稳健的方法，利用离散稀疏自编码器在大语言模型中发现可解释的电路。该方法解决了现有技术的关键局限性，即计算复杂性和对超参数的敏感性。我们建议在精心设计的正例和负例上训练稀疏自编码器，模型只能正确预测正例的下一个token。我们假设注意力头输出的学习表示将指示头何时参与特定计算。通过将学习到的表示离散化为整数代码，并测量每个头的正例独有代码之间的重叠，我们能够直接识别参与电路的注意力头，而无需昂贵的消融实验或架构修改。在三个经过充分研究的任务（间接对象识别、大于比较和文档字符串补全）中，与最先进的基线相比，所提出的方法在恢复ground-truth电路方面实现了更高的精度和召回率，同时将运行时间从数小时缩短到数秒。值得注意的是，我们每个任务只需要5-10个文本示例即可学习稳健的表示。我们的发现突出了离散稀疏自编码器在可扩展和高效的机制可解释性方面的潜力，为分析大型语言模型的内部工作原理提供了一个新的方向。
🔬 方法详解
问题定义：现有方法在大型语言模型中识别电路时，面临计算量大和对超参数敏感的问题。传统的消融实验成本高昂，且难以准确识别关键的神经元或注意力头。因此，需要一种更高效、更稳健的方法来发现语言模型内部的可解释电路。
核心思路：论文的核心思路是利用稀疏自编码器学习注意力头输出的离散表示，并基于正负样本的差异来识别参与特定计算的注意力头。通过训练自编码器，使其能够区分正例和负例，从而提取出与特定任务相关的注意力头。这种方法避免了直接的消融实验，降低了计算成本，并提高了识别的准确性。
技术框架：该方法主要包含以下几个阶段：1) 准备正例和负例数据集，其中正例是模型能够正确预测下一个token的样本，负例则不能。2) 训练稀疏自编码器，使其能够区分正例和负例的注意力头输出表示。3) 将学习到的表示离散化为整数代码。4) 测量每个注意力头正例独有代码之间的重叠程度，以此来识别参与特定计算的注意力头。
关键创新：该方法最重要的创新点在于使用离散稀疏自编码器来学习注意力头输出的表示，并利用正负例的对比来识别电路。与传统的消融实验或基于梯度的分析方法相比，该方法更加高效和稳健，且能够直接识别参与特定计算的注意力头，无需进行复杂的推理。
关键设计：关键设计包括：1) 稀疏自编码器的结构和训练方式，需要保证学习到的表示具有稀疏性和区分性。2) 正负例数据集的设计，需要能够充分反映特定任务的特征。3) 离散化方法，需要选择合适的离散化策略，以保证信息的有效编码。4) 重叠程度的计算方式，需要选择合适的度量标准，以准确评估注意力头之间的关联性。
🖼️ 关键图片
📊 实验亮点
实验结果表明，该方法在间接对象识别、大于比较和文档字符串补全三个任务上，与现有方法相比，在恢复ground-truth电路方面实现了更高的精度和召回率，同时将运行时间从数小时缩短到数秒，并且只需要5-10个文本示例即可学习稳健的表示。
🎯 应用场景
该研究成果可应用于提升大型语言模型的可解释性和可控性，例如，可以用于诊断模型错误、优化模型结构、以及开发更安全可靠的AI系统。此外，该方法还可以应用于其他类型的神经网络，以发现其内部的工作机制。
📄 摘要（原文）
This paper introduces an efficient and robust method for discovering interpretable circuits in large language models using discrete sparse autoencoders. Our approach addresses key limitations of existing techniques, namely computational complexity and sensitivity to hyperparameters. We propose training sparse autoencoders on carefully designed positive and negative examples, where the model can only correctly predict the next token for the positive examples. We hypothesise that learned representations of attention head outputs will signal when a head is engaged in specific computations. By discretising the learned representations into integer codes and measuring the overlap between codes unique to positive examples for each head, we enable direct identification of attention heads involved in circuits without the need for expensive ablations or architectural modifications. On three well-studied tasks - indirect object identification, greater-than comparisons, and docstring completion - the proposed method achieves higher precision and recall in recovering ground-truth circuits compared to state-of-the-art baselines, while reducing runtime from hours to seconds. Notably, we require only 5-10 text examples for each task to learn robust representations. Our findings highlight the promise of discrete sparse autoencoders for scalable and efficient mechanistic interpretability, offering a new direction for analysing the inner workings of large language models.