Automated Attention Pattern Discovery at Scale in Large Language Models
作者: Jonathan Katzy, Razvan-Mihai Popescu, Erik Mekkes, Arie van Deursen, Maliheh Izadi
分类: cs.LG, cs.AI
发布日期: 2026-04-07
💡 一句话要点
提出AP-MAE,通过注意力模式分析和干预提升大语言模型性能
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 大语言模型 可解释性 注意力机制 视觉Transformer 掩码自编码器
📋 核心要点
- 现有大语言模型可解释性方法缺乏泛化能力,且计算成本高昂,难以进行大规模研究。
- 利用代码结构化特性,挖掘注意力模式中的重复行为,作为模型全局可解释性的可扩展信号。
- 提出AP-MAE模型,通过重建掩码注意力模式,实现模型分析、预测和干预,并开源代码和模型。
📝 摘要(中文)
大型语言模型的能力随规模增长而提升,但可解释性方法的发展却相对滞后。当前的机制可解释性研究侧重于在受控环境中对特定行为进行精确解释,但这些解释通常不具备泛化性,或者需要消耗大量资源才能进行更大规模的研究。本文提出通过挖掘Java代码数据集中的补全场景来研究大型语言模型中重复出现的行为,利用代码的结构化特性。我们收集注意力头的注意力模式,证明它们是模型组件全局可解释性的可扩展信号。我们展示了视觉模型为大规模分析注意力模式提供了一个有希望的方向。为此,我们引入了注意力模式-掩码自编码器(AP-MAE),这是一个基于视觉Transformer的模型,可以有效地重建被掩盖的注意力模式。在StarCoder2上的实验表明,AP-MAE (i) 以高精度重建被掩盖的注意力模式,(ii) 以最小的性能下降泛化到未见过的模型,(iii) 揭示跨推理的重复模式,(iv) 在不访问ground truth的情况下预测生成的正确性,准确率在55%到70%之间,具体取决于任务,以及 (v) 实现有针对性的干预,选择性应用时可将准确率提高13.6%,但过度应用会导致崩溃。这些结果表明,注意力模式是可解释性的可扩展信号,并证明AP-MAE为大型语言模型的分析和干预提供了可转移的基础。除了其独立价值外,AP-MAE还可以作为指导细粒度机制方法选择的程序。我们发布了代码和模型,以支持未来在大型可解释性方面的工作。
🔬 方法详解
问题定义:现有的大语言模型可解释性方法,例如机制可解释性,通常针对特定行为进行精确解释,但这些解释难以泛化到其他场景或模型,并且在大规模研究中计算成本过高。因此,需要一种可扩展的方法来理解和解释大型语言模型的行为。
核心思路:本文的核心思路是利用注意力模式作为一种可扩展的信号来理解大型语言模型的行为。通过分析模型在处理特定任务(例如代码补全)时产生的注意力模式,可以揭示模型内部的重复行为和潜在的机制。此外,利用视觉模型(特别是视觉Transformer)来处理和分析注意力模式,可以提高效率和可扩展性。
技术框架:整体框架包括以下几个主要步骤:1) 在Java代码数据集上收集补全场景;2) 提取模型在这些场景中生成的注意力模式;3) 使用AP-MAE模型重建被掩盖的注意力模式;4) 分析重建的注意力模式,以识别重复行为和预测模型性能;5) 基于注意力模式进行有针对性的干预,以提高模型性能。
关键创新:最重要的技术创新点是提出了AP-MAE模型,这是一个基于视觉Transformer的模型,专门用于重建和分析注意力模式。与传统的注意力模式分析方法相比,AP-MAE能够更有效地处理大规模的注意力模式数据,并且具有更好的泛化能力。此外,AP-MAE还可以用于预测模型性能和进行有针对性的干预。
关键设计:AP-MAE的关键设计包括:1) 使用掩码自编码器(MAE)架构,通过重建被掩盖的注意力模式来学习注意力模式的表示;2) 使用视觉Transformer作为编码器和解码器,以有效地处理注意力模式数据;3) 设计合适的损失函数,以优化AP-MAE的重建性能;4) 通过实验选择合适的掩码比例和网络结构。
🖼️ 关键图片
📊 实验亮点
实验结果表明,AP-MAE能够以高精度重建被掩盖的注意力模式,并且具有良好的泛化能力。在StarCoder2上的实验表明,AP-MAE能够以55%到70%的准确率预测生成的正确性,并且通过有针对性的干预,可以将准确率提高13.6%。这些结果表明,注意力模式是可解释性的可扩展信号,AP-MAE为大型语言模型的分析和干预提供了有价值的工具。
🎯 应用场景
该研究成果可应用于大语言模型的可解释性分析、模型调试和性能优化。通过AP-MAE,研究人员可以更好地理解模型的内部机制,预测模型的行为,并进行有针对性的干预,从而提高模型的可靠性和性能。此外,该方法还可以用于检测和缓解模型中的偏见和安全漏洞。
📄 摘要(原文)
Large language models have found success by scaling up capabilities to work in general settings. The same can unfortunately not be said for interpretability methods. The current trend in mechanistic interpretability is to provide precise explanations of specific behaviors in controlled settings. These often do not generalize, or are too resource intensive for larger studies. In this work we propose to study repeated behaviors in large language models by mining completion scenarios in Java code datasets, through exploiting the structured nature of code. We collect the attention patterns generated in the attention heads to demonstrate that they are scalable signals for global interpretability of model components. We show that vision models offer a promising direction for analyzing attention patterns at scale. To demonstrate this, we introduce the Attention Pattern - Masked Autoencoder(AP-MAE), a vision transformer-based model that efficiently reconstructs masked attention patterns. Experiments on StarCoder2 show that AP-MAE (i) reconstructs masked attention patterns with high accuracy, (ii) generalizes across unseen models with minimal degradation, (iii) reveals recurring patterns across inferences, (iv) predicts whether a generation will be correct without access to ground truth, with accuracies ranging from 55% to 70% depending on the task, and (v) enables targeted interventions that increase accuracy by 13.6% when applied selectively, but cause collapse when applied excessively. These results establish attention patterns as a scalable signal for interpretability and demonstrate that AP-MAE provides a transferable foundation for both analysis and intervention in large language models. Beyond its standalone value, AP-MAE also serves as a selection procedure to guide fine-grained mechanistic approaches. We release code and models to support future work in large-scale interpretability.