Investigating the Indirect Object Identification circuit in Mamba

📄 arXiv: 2407.14008v2 📥 PDF

作者: Danielle Ensign, Adrià Garriga-Alonso

分类: cs.LG

发布日期: 2024-07-19 (更新: 2024-07-22)

💡 一句话要点

研究Mamba模型中的间接对象识别电路，揭示其内部机制。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: Mamba 可解释性 间接对象识别 电路分析 循环神经网络

📋 核心要点

  1. 现有可解释性技术在新型架构（如Mamba）上的适用性面临挑战，需要针对性研究。
  2. 通过将预Mamba技术适配到Mamba架构，逆向工程间接对象识别（IOI）任务的关键电路。
  3. 实验结果揭示了Mamba模型中IOI电路的关键层、卷积作用以及名称实体的存储方式。

📝 摘要（中文）

本文研究了当前的可解释性技术在未来模型中的泛化能力，以Mamba这一新型循环架构为例，其扩展性与Transformer相当。我们将预Mamba技术应用于Mamba，并部分逆向工程了负责间接对象识别（IOI）任务的电路。研究结果表明：1）第39层是一个关键瓶颈；2）第39层中的卷积将名称实体向前移动一个位置；3）名称实体线性地存储在第39层的SSM中。最后，我们采用了一种自动电路发现工具，即位置边缘属性修补，来识别Mamba IOI电路。我们的贡献初步证明了基于电路的机制可解释性工具适用于Mamba架构。

🔬 方法详解

问题定义：论文旨在理解Mamba模型如何执行间接对象识别（IOI）任务，并验证现有的可解释性技术是否能有效地应用于这种新型架构。现有方法主要针对Transformer架构，缺乏对Mamba等循环架构的深入研究，因此需要探索适用于Mamba的可解释性方法。

核心思路：论文的核心思路是通过适配预Mamba的可解释性技术，并结合自动电路发现工具，来逆向工程Mamba模型中负责IOI任务的电路。通过分析关键层、卷积操作和状态空间模型（SSM）的内部状态，揭示模型如何存储和处理名称实体。

技术框架：该研究的技术框架主要包括以下几个阶段：1) 将预Mamba的可解释性技术适配到Mamba架构；2) 分析Mamba模型在执行IOI任务时的内部状态，特别是第39层；3) 利用位置边缘属性修补（positional Edge Attribution Patching）等自动电路发现工具，识别Mamba IOI电路；4) 验证识别出的电路是否能够解释Mamba模型在IOI任务上的行为。

关键创新：该研究的关键创新在于将电路视角的可解释性方法应用于Mamba架构，并揭示了Mamba模型中IOI电路的一些关键特征，例如第39层作为瓶颈、卷积操作的作用以及名称实体在SSM中的线性存储。此外，该研究还验证了自动电路发现工具在Mamba架构上的有效性。

关键设计：论文中关键的设计包括：选择Mamba模型作为研究对象，因为它是一种新型的循环架构，具有与Transformer相当的扩展性；采用间接对象识别（IOI）任务作为测试用例，因为它是一个经典的自然语言处理任务，可以用来评估模型的可解释性；使用位置边缘属性修补（positional Edge Attribution Patching）作为自动电路发现工具，它可以自动识别模型中的关键连接。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

研究发现Mamba模型第39层是IOI任务的关键瓶颈，该层中的卷积操作将名称实体向前移动一个位置，并且名称实体线性地存储在第39层的SSM中。通过位置边缘属性修补，成功识别了Mamba IOI电路。

🎯 应用场景

该研究成果可应用于提升Mamba等新型神经网络架构的可解释性和可信度，有助于开发更安全、可靠的AI系统。理解模型内部机制有助于优化模型设计，并为未来的可解释性研究提供借鉴。

📄 摘要（原文）

How well will current interpretability techniques generalize to future models? A relevant case study is Mamba, a recent recurrent architecture with scaling comparable to Transformers. We adapt pre-Mamba techniques to Mamba and partially reverse-engineer the circuit responsible for the Indirect Object Identification (IOI) task. Our techniques provide evidence that 1) Layer 39 is a key bottleneck, 2) Convolutions in layer 39 shift names one position forward, and 3) The name entities are stored linearly in Layer 39's SSM. Finally, we adapt an automatic circuit discovery tool, positional Edge Attribution Patching, to identify a Mamba IOI circuit. Our contributions provide initial evidence that circuit-based mechanistic interpretability tools work well for the Mamba architecture.