KGPFN: Unlocking the Potential of Knowledge Graph Foundation Model via In-Context Learning

KGPFN の技術的概要：イン・コンテキスト学習による知識グラフ基盤モデルの潜在能力の解明

1. 問題定義

本論文は、現在の知識グラフ（KG）基盤モデルにおける重要な限界に対処するものである：すなわち、それらは転移可能な関係パターン（普遍性）の学習においては卓越しているものの、基盤モデルの第 2 の柱である**イン・コンテキスト学習（ICL）**を largely 軽視している点である。既存の KG 基盤モデルは、広範で再利用可能なパターンを過剰に適用しがちであり、微調整なしには特定のグラフインスタンスに適応した予測を行うことができない。

著者らは、効果的な KG 推論には、本質的に構造化され異質な 2 種類の文脈に条件付けられる必要があると主張する：

• ローカル文脈： 推論が行われるクエリ固有の近傍領域である。同じ構造的合図（例：「三角形」パターン）は、周囲の部分グラフ（例：出生地対国籍）に応じて異なる結論をもたらす可能性がある。
• グローバル文脈： 多くのインスタンスにわたって蓄積された関係レベルの証拠であり、特定の関係がグラフ全体でどのように典型的に振る舞うかを記述する（例：共起する関係、パスの頻度、有効なエンティティに対する制約）。

現在の手法は、これらの相補的なシグナルを統合し、再トレーニングなしに未見のグラフへのゼロショットまたは少数ショット適応を可能にするのに苦労している。

2. 手法：KGPFN

著者らは、構造化された KG データ上でイン・コンテキスト学習を実行するためにPrior-Data Fitted Network（PFN）パラダイムを活用する KG 基盤モデルであるKGPFNを提案する。このフレームワークは、ターゲットグラフでの微調整を行わず、代わりに事前学習と推論時の適応に依存して動作する。

主要コンポーネント：

1. 文脈構築：

   • グローバル文脈： クエリ関係 $r$ に対して、モデルはグラフから正のトリプルと負のトリプル（インスタンス）のセットを取得する。これらのインスタンスとそれらのローカル近傍は、関係レベルの規則性を捉えるサポートセットを形成する。
   • ローカル文脈： 計算コストが高くなるためすべての候補に対して完全な部分グラフを抽出する代わりに、モデルはクエリヘッドエンティティの周りにヘッド中心の$k$-ホップ近傍を構築する。これは大規模な候補セットに対して効率的にスケーリングする。

2. 文脈的グラフエンコーダ：

   • 関係表現： 関係グラフ（ノードが関係タイプで、エッジが相互作用を表す）上のメッセージパッシングメカニズムが、クエリ条件付きの関係埋め込みを生成する。これはグラフ間を超えた関係的不変性を捉える。
   • ローカルエンコーディング： 多層のNBFNet（Neural Belief Propagation Network）がヘッド中心の部分グラフをエンコードし、最大$k$ホップまでの証拠を集約するマルチスケールなローカル文脈ベクトルを生成する。
   • テール認識スコアリング： モデルは、ヘッド表現、関係埋め込み、およびテール表現（最後の NBFNet レイヤーから）を、相互作用特徴（TransE、DistMult、およびコサイン類似度）を用いて組み合わせ、テール認識スコア表現を形成する。
   • 融合： 異質な特徴（ローカル文脈、関係特徴、および相互作用特徴）は、軽量アダプターを介して共有潜在空間に整列され、連結される。

3. PFN アーキテクチャ（イン・コンテキスト学習）：

   • クエリとラベル付きグローバル文脈セットの融合された表現が PFN モジュールに入力される。
   • 特徴アテンション： 各サンプル（ローカル文脈）内の特徴次元を洗練させる。
   • サンプルアテンション： 文脈例全体（グローバル文脈）にわたって証拠を集約する。これにより、クエリはラベル付きサポートセットから関係固有の証拠を检索できる。
   • メカニズム： PFN はベイズ事後予測を近似する。これは、文脈セットを観測された証拠として扱い、クエリに対する予測分布を誘導し、効果的に、いつ再利用可能なパターンをインスタンス化し、いつ文脈的証拠に基づいてそれらをオーバーライドするかを学習する。

4. トレーニング目的：

   • モデルは、正と負を分離するためのバイナリクロスエントロピー（BCE）損失と、サンプリングされた負のサンプルよりも正のサンプルのランキングを強制するためのソフトマックスクロスエントロピー損失の組み合わせを用いて、多様な KG 上でトレーニングされる。
   • 負のサンプリングが使用され、適応的温度メカニズムを介してハードネガティブに重みが付けられる。

3. 主要な貢献

本論文は、3 つの主要な貢献を概説する：

• KG 向けのデュアル文脈構築： 著者らは、転移可能なパターンに主に焦点を当てている既存の基盤モデルのギャップに対処するため、ローカル（クエリ固有の近傍）とグローバル（関係レベルの規則性）の両方の文脈を明示的に構築する戦略を導入する。
• PFN ベースのイン・コンテキスト学習： KGPFN は、構造化された KG データに Prior-Data Fitted Network パラダイムを適用した最初のモデルである。特徴レベルとサンプルレベルのアテンションを組み合わせることで、パラメータ更新なしに推論時の未見グラフへの適応を可能にする。
• 最先端のパフォーマンス： 広範な実験により、KGPFN が帰納的および完全帰納的設定の両方で優れたパフォーマンスを達成し、イン・コンテキスト学習のみを使用する競合する微調整ベースラインを上回ることが示された。

4. 実験結果

モデルは、帰納的（未見エンティティ）、完全帰納的（未見エンティティおよび関係）の 3 つの設定にわたる57 の KG ベンチマークで評価された。

• パフォーマンス： KGPFN は、すべての 57 データセットで平均平均逆数ランク（MRR）および Hits@10 において最良の結果を達成した。
• 帰納的能力： 最も困難な完全帰納的設定において、KGPFN の優位性が最も顕著であり、ターゲットデータセットで微調整された ULTRA、TRIX、MOTIF などのモデルを一貫して上回った。
• 文脈感受性： 分析により、グローバル文脈における負のインスタンスの数を増やすことがパフォーマンスを大幅に向上させることが示された一方、正のインスタンスの数は影響が少なかった。さらに、**マルチホップ（3 ホップ）**ローカル文脈を使用することが不可欠であり、1 ホップまたは 2 ホップの近傍ではパフォーマンスが低下した。
• 可視化： アテンションマップは、モデルが正しいテール候補に対して正の文脈インスタンスにアテンションを正確に集中させ、誤った候補に対して負のインスタンスにアテンションを分散させることを確認した。

5. 意義と主張

本論文は、KGPFN が知識グラフ推論における文脈ベースの適応の力を示していると主張する。PFN パラダイムを介してグローバルな関係的規則性とローカルな構造的合図を統合することで、モデルは以前に未見であったグラフへの強力な一般化を達成する。

著者らは、このアプローチによりモデルが構造的な不変性だけでなく機能的表現力を学習できることを強調している。理論的には、PFN アテンションメカニズムは、ローカルな構造的証拠をイン・コンテキストサポートセットから推論されたラベル重み付きパターン統計と整列させることで、オンザフライで関係固有のスコアリング関数を構成する。これにより、モデルは微調整なしに未見の関係に適応でき、真に一般的な KG 基盤モデルへの重要な一歩を刻むことになる。

この研究はプレプリントとして発表され、コードはコミュニティに公開されている。著者らは、計算リソースの制約によりスケーリング則の検証に関する限界があることを指摘している。