[论文解读] Dissecting Recall of Factual Associations in Auto-Regressive Language Models

Mor Geva, Jasmijn Bastings+2 more|arXiv (Cornell University)|Apr 28, 2023

Topic Modeling被引用 5

一句话总结

本文通过分析推理过程中的信息流，剖析自回归语言模型如何检索事实关联。通过有针对性的注意力干预，揭示了一个三步机制：早期MLP层丰富主体表征，关系传播至预测标记，高层Transformer层通过编码主体-属性映射的注意力头提取正确属性——该现象出现在约70%的预测中，GPT-2中涉及150个注意力头。

ABSTRACT

Transformer-based language models (LMs) are known to capture factual knowledge in their parameters. While previous work looked into where factual associations are stored, only little is known about how they are retrieved internally during inference. We investigate this question through the lens of information flow. Given a subject-relation query, we study how the model aggregates information about the subject and relation to predict the correct attribute. With interventions on attention edges, we first identify two critical points where information propagates to the prediction: one from the relation positions followed by another from the subject positions. Next, by analyzing the information at these points, we unveil a three-step internal mechanism for attribute extraction. First, the representation at the last-subject position goes through an enrichment process, driven by the early MLP sublayers, to encode many subject-related attributes. Second, information from the relation propagates to the prediction. Third, the prediction representation "queries" the enriched subject to extract the attribute. Perhaps surprisingly, this extraction is typically done via attention heads, which often encode subject-attribute mappings in their parameters. Overall, our findings introduce a comprehensive view of how factual associations are stored and extracted internally in LMs, facilitating future research on knowledge localization and editing.

研究动机与目标

• 理解自回归语言模型在推理过程中内部检索事实关联的机制。
• 识别在Transformer计算中主体与关系信息汇聚以生成事实预测的关键计算点。
• 研究MLP与多头自注意力（MHSA）子层在构建和提取属性表征中的作用。
• 确定主体-属性映射是否编码在注意力头参数中，以及其对预测的贡献方式。
• 为未来知识编辑与定位研究提供透明、可解释的知识回忆机制视图。

提出的方法

• 作者采用有针对性的注意力关闭干预，阻断从最后一个标记到特定输入位置的注意力，以评估其重要性。
• 通过测量主体与关系信号汇聚的关键层处的表征变化，分析信息流。
• 将隐藏表征投影到词汇空间，以近似每一层和每一头编码的信息。
• 通过检查注意力头输出的前10个标记中是否包含预测的属性，识别并分析编码主体-属性映射的注意力头。
• 聚焦于GPT-2和GPT-J等仅解码器模型，分析模型正确预测下一个标记的案例。
• 采用受遗传分析启发的逆向工程方法，系统性地破坏组件，以评估其在预测中的功能角色。

Figure 1: Illustration of our findings: given subject-relation query, a subject representation is constructed via attributes’ enrichment from MLP sublayers (A), while the relation propagates to the prediction (B). The attribute is then extracted by the MHSA sublayers (C).

实验结果

研究问题

• RQ1 主体与关系信息在哪些层和位置汇聚，以支持事实预测？
• RQ2 在MLP或MHSA子层中，哪一个主要负责通过属性丰富主体表征？
• RQ3 在推理过程中，如何从主体表征中提取最终属性？
• RQ4 主体-属性映射在多大程度上编码在注意力头参数中？
• RQ5 我们能否在不同事实关联中识别出一致的内部机制用于事实回忆？

主要发现

• 模型通过早期MLP子层在主体最后一个位置构建出丰富且富含属性的表征，这些子层是主体相关信息的主要来源。
• 计算过程中识别出两个关键节点：第一，关系表征被传播至预测标记；第二，主体的丰富表征被整合进来。
• 最终的属性提取由高层MHSA子层完成，其通过将关系作为键，查询富含属性的主体表征。
• 约70%的事实预测涉及直接在其参数中编码主体-属性映射的注意力头，GPT-2中识别出150个此类头，主要集中在高层。
• 属性提取机制具有鲁棒性和一致性：关系引导注意力机制从主体表征中检索出正确属性。
• 研究结果表明，事实知识不仅存储在中间层的MLP中，还通过参数化注意力头被主动检索，挑战了关于知识存储的既有假设。

Figure 3: Attributes rate at different positions across layers (starting from layer 1), in GPT-2 and GPT-J .