认知之镜：探索 LLM 推理中注意力头的多样化功能角色

大型语言模型（LLM）在各种任务中取得了最先进的性能，但其内部机制在很大程度上仍然是不透明的。理解这些机制对于提升其推理能力至关重要。受神经过程与人类认知之间相互作用的启发，我们提出了一个新颖的可解释性框架，以系统地分析注意力头的角色和行为，这些注意力头是 LLM 的关键组件。我们引入了 CogQA，这是一个通过思维链设计将复杂问题分解为逐步子问题的数据集，每个子问题都与特定的认知功能（如检索或逻辑推理）相关联。通过应用多类探测方法，我们识别出负责这些功能的注意力头。我们对多个 LLM 家族的分析表明，注意力头表现出功能特化，被表征为认知头。这些认知头表现出几个关键特性：它们普遍是稀疏的，在不同认知功能中的数量和分布各不相同，并显示出交互和层级结构。我们进一步表明，认知头在推理任务中发挥着至关重要的作用——移除它们会导致性能下降，而增强它们则会提高推理准确率。这些见解为深入理解 LLM 推理提供了依据，并对模型设计、训练和微调策略具有重要启示。