Polar probe linearly decodes semantic structures from LLMs

技術的サマリー：極性プローブは LLM から意味構造を線形に復号化する

問題の定式化
本論文は、大規模言語モデル（LLM）が概念を結合して複雑な意味構造を形成する方法という根本的な問いに答える。先行研究では、構文構造（例えば、主語と動詞の依存関係）が幾何学的原理を用いて LLM の活性化から線形に復号化可能であることが示されている（Hewitt & Manning, 2019; Diego-Simon et al., 2024）が、同様の原理が意味的結合にも適用されるかどうかは不明なままだ。著者らは、LLM がエンティティ間の関係の存在と種類を、複雑な非線形メカニズムではなく、単純な幾何学的コードによって表現していると仮定する。すなわち、エンティティ埋め込み間のユークリッド距離が関係の存在を符号化し、相対的な方向が関係の種類を符号化するというものである。

手法
この仮説を検証するために、著者らは極性プローブを導入する。これは、LLM の層の活性化を、関係グラフを復元可能なプローブ空間へマッピングするように訓練された線形変換である。

1. データセット: 本研究は、5 つの異なる意味領域にまたがる合成データセットを利用する。

   • ユークリッド領域: 変数の順序性（数学的変数）、空間的レイアウト（2 次元オブジェクト配置）、および役割（社会的相互作用）。
   • 非ユークリッド領域: 家系図（非可換、多対多の関係）および地下鉄路線図（ネットワークベースの最短経路）。
   • 各サンプルは、関係グラフの自然言語記述で構成される。モデルが表面レベルの相関を学習するのを防ぐため、関係とエンティティの順序はランダム化され、同一グラフに対して複数のテキスト記述が生成される。

2. プローブのアーキテクチャ:

   • 構造的損失（$L_s$）: 予測されたペアワイズ距離行列と、意味グラフの真の最短経路距離行列との間のスピアマンの順位相関を最適化する。
   • 角度損失（$L_a$）: 2 つのエンティティ埋め込みの差ベクトル（$\delta_{ij}$）と、関係タイプ $r$ に固有の学習済みプロトタイプベクトル（$p_r$）との間の整合性を最適化する。
   • プローブは、様々な LLM（Llama3-8B、OLMo-7B、および Pythia ファミリー）の各層で独立して訓練される。

3. 評価:

   • 復号化性能: 関係の存在（距離の順序付け）と関係の種類（方向の整合性）について、スピアマンの $\rho$ を通じて測定される。
   • 一般化: 分布外（OOD）のエンティティ名および関係の表面形式に対してテストされる。
   • 因果的介入: 著者らは、学習済みプロトタイプ方向に沿って活性化を誘導することにより隠れ状態に介入し、下流の質問応答（QA）タスクにおける性能のシフトを観察する。
   • 相関: プローブの誤差は、QA タスクにおける正解のロジット確率と相関する。

主要な結果

• 線形復号化可能性: 真の意味構造は LLM の活性化から線形に復元可能である。極性プローブは、5 つのすべての領域において、関係の存在と種類の両方を成功裡に復号化する。
• 層の出現: 幾何学的コードは、事前学習済み LLM の中間層で最も顕著である（Llama3-8B の場合、層 12〜15 付近でピークに達する）。ランダムに初期化されたモデルでは、復号化性能は偶然のレベルにとどまり、このコードが事前学習中に出現することを示している。
• モデルサイズと学習: パラメータ数 14M から 69 億までモデルサイズが大きくなるにつれて性能は向上し、事前学習ステップの増加に伴い向上する。これは、極性コードがモデルの容量と学習期間に依存する学習された能力であることを示唆している。
• 一般化: プローブは新しいエンティティや関係の表面形式に対して一般化するが、OOD エンティティでは性能がやや低下し、OOD 関係タイプではより顕著に低下する。また、グラフ内のエンティティ数が増加するにつれて性能も低下するが、これはおそらく拡張された組み合わせ空間によるものである。
• 因果的影響: 学習済みプロトタイプ方向に沿ってモデルの活性化を誘導することは、QA タスクにおける正解の確率を因果的にシフトさせ、これらの幾何学的表現がモデルの推論機能的に重要であることを確認する。
• 部分空間の非重複性: 異なる意味領域のために学習された部分空間は、活性化空間において概して非重複であるが、空間的レイアウトと変数の順序性については顕著な重複が見られ（重畳仮説を支持する）、これは注目すべき例外である。

意義と主張
本論文は、LLM が単純な幾何学的原理を通じて表現を結合することで、複雑な意味構造を構築することを学習していると主張する。具体的には、構文に限定されていた線形プロービングの範囲を、方向性を持ちラベル付けされたグラフのより広範なクラスへと拡張する。

• 構文を超えて: 構文木に限定された先行研究とは異なり、このアプローチは、必ずしも木構造ではなく非ユークリッド位相（例：家系図、地下鉄路線図）を含む意味構造を成功裡に復号化する（ただし、非可換関係については精度にいくつかの制限がある）。
• 知識の検索を超えて: 本研究は、ランダム化されたエンティティと関係によりトレーニング中に記憶され得ない構成的構造を対象とし、モデルがコンテキスト内学習と結合をどのように実行するかについての証拠を提供する。
• 記号主義と接続主義の緊張関係: この発見は、記号主義的アプローチと接続主義的アプローチの間の緊張関係の解決を示唆しており、記号的構造は線形部分空間と幾何学的関係を通じてベクトル系内で表現可能であることを提案する。

著者らは、手法の限界について控えめに述べており、極性プローブは線形距離に依存し、ユークリッドグラフ向けに調整されていると指摘する。したがって、非ユークリッド領域で見られる完全な非可換・多対多の関係や複雑な最短経路距離を完全に表現することはできないが、それら内の局所的な一貫性は捉えている。