EleutherAIの研究チームは、言語モデルの内部状態を分類する新しい方法「アテンションプローブ」を提案しました。従来の線形プローブは、トークンごとまたは平均プーリングなどの圧縮表現に基づいて訓練されていましたが、アテンションプローブはアテンション層を用いて隠れ状態を収集し、プーリングによる情報損失を回避します。
この手法の疑似コードによれば、アテンションプローブは複数のヘッドを持ち、各ヘッドが各トークンに対して1つのアテンション対数を計算し、ソフトマックスでアテンション確率を求めます。標準的なマルチヘッドアテンションとは異なり、各ヘッドは1つのトークンのみに注目します（クロスアテンションでクエリトークンが1つの場合に類似）。さらに、学習可能な位置バイアス（ALiBiに類似）が追加され、アテンションがシーケンス内の特定の位置に偏ることを可能にします。最終的に、値投影と重み付き和によりプローブ出力が得られます。
関連研究として、McKenzieら（2025）は同様のアーキテクチャを提案しましたが、シングルヘッドで位置バイアスはありませんでした。Kantamneniら（2025）はアテンションプローブを最初に示しましたが、補助的な手法でした。本研究では、MOSAIC（Gemma 2BおよびGemma 2 2Bモデルを使用）とNeurons-In-A-Haystack（NiAH）データセットを使用し、アテンションプローブはAdamW、平均/最終トークンプローブはLBFGSで最適化されました。
実験結果では、MOSAICデータセットでは平均プローブが最終トークンプローブより優れ、8ヘッドアテンションプローブ（AdamW）は平均プローブを全体的に上回り、特にAdamWで訓練された平均プローブより常に優れていました。シングルヘッドアテンションプローブは中間的な結果でした。NiAHデータセットでは、アテンションプローブは最終トークンプローブより明確に優れているとは言えませんでしたが、シングルヘッドでも平均プローブより改善が見られました。また、ヘッド数が増えるとアテンションエントロピーが増加し、データセットに大きく依存することがわかりました。
最大活性化例の分析では、アテンションプローブがタスク関連の単語（例：性別関連語）に注目することが示されました。全体として、アテンションプローブは平均または最終トークンプローブと同等かそれ以上の性能を発揮し、特にマルチヘッド構成で優れています。トレーニングコードはオープンソースで公開され、pipでインストール可能です。