Flash Interpretability: Decoding Specialised Feature Neurons in Large Language Models with the LM-Head

📄 arXiv: 2501.02688v2 📥 PDF

作者: Harry J Davies

分类: cs.CL

发布日期: 2025-01-05 (更新: 2025-02-27)

备注: 5 pages, 4 figures


💡 一句话要点

利用LM-Head解码LLM特征神经元,实现快速且低成本的可解释性分析

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 可解释性 神经元分析 LM-Head 特征解码

📋 核心要点

  1. 大型语言模型参数众多,内部运作机制复杂,可解释性差,难以理解。
  2. 该论文提出利用LM-Head将神经元权重直接解码为token概率,从而识别特定功能的神经元。
  3. 实验表明,该方法能快速定位并验证特定概念的神经元,且计算成本极低。

📝 摘要(中文)

大型语言模型(LLM)通常具有数十亿个参数,因此其运行机制往往难以解释。本文展示了通过模型的最终投影层(LM-head)将神经元权重直接解码为token概率是可行的。以Llama 3.1 8B为例,我们使用LM-head找到了专门的特征神经元,例如“dog”神经元和“California”神经元,并通过钳制这些神经元来影响输出中概念的概率,从而验证了这一点。我们在预训练模型和Instruct模型上评估了该方法,发现Instruct模型中up-projection层超过75%的神经元与预训练模型具有相同的top关联token。最后,我们证明了钳制“dog”神经元会导致Instruct模型在被问及最喜欢的动物时总是讨论狗。通过我们的方法,可以在不到10秒的时间内,以最小的计算量映射Llama 3.1 8B的up-projection神经元的全部top特征。

🔬 方法详解

问题定义:现有大型语言模型由于参数规模庞大,其内部神经元的功能难以理解和解释。传统的分析方法计算成本高昂,效率低下,难以快速定位和理解特定功能的神经元。因此,如何以低成本、高效率的方式理解LLM内部神经元的功能是一个重要的挑战。

核心思路:该论文的核心思路是利用LLM的LM-Head(Language Modeling Head,语言模型头)作为桥梁,将神经元的权重直接映射到token的概率分布。LM-Head负责将模型的内部表示转换为最终的token预测,因此其权重包含了神经元与token之间的关联信息。通过分析神经元在LM-Head中的权重,可以推断该神经元激活时最可能产生的token,从而揭示其功能。

技术框架:该方法主要包含以下几个步骤:1) 获取LLM的LM-Head权重矩阵。2) 对于每个神经元,计算其在LM-Head中的权重向量与所有token embedding之间的相似度(例如,点积)。3) 根据相似度得分对token进行排序,选择得分最高的token作为该神经元对应的“top token”。4) 为了验证神经元的功能,采用“钳制”(clamping)技术,即强制激活或抑制特定神经元,观察其对模型输出的影响。

关键创新:该方法最重要的创新在于利用LM-Head作为解码器,将神经元权重直接映射到token概率,从而避免了复杂的中间步骤和大量的计算。这种方法极大地提高了分析效率,使得快速理解LLM内部神经元的功能成为可能。与传统的激活最大化等方法相比,该方法无需进行迭代优化,计算成本更低。

关键设计:该方法的关键设计在于选择LM-Head作为解码器。LM-Head通常是一个线性层,其权重矩阵直接连接了模型的内部表示和token embedding。因此,LM-Head的权重包含了神经元与token之间的直接关联信息。此外,该方法还采用了“钳制”技术来验证神经元的功能。通过强制激活或抑制特定神经元,可以观察其对模型输出的影响,从而验证该神经元是否与特定的概念或功能相关。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

该研究在Llama 3.1 8B模型上进行了实验,结果表明该方法能够快速有效地识别特定功能的神经元,例如“dog”神经元和“California”神经元。实验还发现,Instruct模型中超过75%的神经元与预训练模型具有相同的top关联token,这表明Instruct微调过程在很大程度上保留了预训练模型的知识。通过钳制“dog”神经元,可以使Instruct模型在被问及最喜欢的动物时总是讨论狗,这验证了该方法的可行性。

🎯 应用场景

该研究成果可应用于提升大型语言模型的可解释性和可控性。例如,可以用于识别和修复模型中存在的偏见或有害信息。此外,该方法还可以用于指导模型的微调和优化,使其更好地完成特定任务。未来,该技术有望应用于开发更安全、更可靠、更易于理解的人工智能系统。

📄 摘要(原文)

Large Language Models (LLMs) typically have billions of parameters and are thus often difficult to interpret in their operation. In this work, we demonstrate that it is possible to decode neuron weights directly into token probabilities through the final projection layer of the model (the LM-head). This is illustrated in Llama 3.1 8B where we use the LM-head to find examples of specialised feature neurons such as a "dog" neuron and a "California" neuron, and we validate this by clamping these neurons to affect the probability of the concept in the output. We evaluate this method on both the pre-trained and Instruct models, finding that over 75% of neurons in the up-projection layers in the instruct model have the same top associated token compared to the pretrained model. Finally, we demonstrate that clamping the "dog" neuron leads the instruct model to always discuss dogs when asked about its favourite animal. Through our method, it is possible to map the top features of the entirety of Llama 3.1 8B's up-projection neurons in less than 10 seconds, with minimal compute.