EleutherAI 研究团队提出了一种名为“注意力探针”（Attention Probes）的新方法，用于分类语言模型的内部状态。传统线性探针通常基于每个token或通过池化（如均值池化或取最后一个token）压缩多个token的潜在向量来进行训练。注意力探针则引入一个注意力层来收集隐藏状态，从而避免了池化带来的信息损失。
该方法的核心伪代码显示，注意力探针具有多个头，每个头为每个token计算一个注意力对数，并通过softmax得到注意力概率。与标准多头注意力不同，这里每个头只关注一个token（类似于交叉注意力中只有一个查询token）。此外，还添加了一个可学习的位置偏置（类似ALiBi），使得注意力可以偏向序列中的特定位置。最终，通过值投影和加权求和得到探针输出。
相关工作方面，McKenzie等人（2025）曾提出类似架构，但仅使用单头且无位置偏置；Kantamneni等人（2025）最早展示了注意力探针，但作为次要方法。本研究使用的数据集包括MOSAIC（基于Gemma 2B和Gemma 2 2B模型）以及Neurons-In-A-Haystack（NiAH）数据集。训练时，探针使用AdamW优化器（注意力探针）或LBFGS（均值/末位token探针），并进行超参数搜索。
实验结果显示，在MOSAIC数据集上，均值探针优于末位token探针，而8头注意力探针（AdamW训练）整体优于均值探针，且始终优于AdamW训练的均值探针。单头注意力探针表现中等。在NiAH数据集上，注意力探针未明显优于末位token探针，但即使单头注意力探针也比均值探针有所改进。此外，随着头数增加，注意力熵也增加，且熵值高度依赖于数据集。
通过分析最大激活示例，研究者发现注意力探针有时会关注与任务相关的词语（如性别相关词）。总体而言，注意力探针的性能与均值或末位token探针相当，但在多数情况下更优，尤其是多头配置。训练代码已开源，可通过pip安装。