Knowledge Graph Enhanced Memory-Augmented Retrieval for Long Context Modeling

技術要約: KGERMAR

問題提起
長文脈言語モデリングは、単にコンテキストウィンドウを拡張することを超えた、根本的な課題に直面している。それは、数千トークンにわたってエンティティの状態や関係性を一貫して理解し続けることである。標準的なアプローチには、主に2つの限界がある。第一に、拡張アテンションモデル（例：LongLoRA, YARN）は、「真ん中での消失（lost-in-the-middle）」による性能低下、すなわち長いコンテキストの中間に位置する情報へのアテンションが系統的に低下するという問題に苦しんでいる。第二に、メモリ拡張型検索システム（例：MemLong, ERMAR）は、意味的類似性にのみ依存している。これらは、一致する用語を含むパッセージを検索することには効果的だが、パッセージが単に共通の用語を偶然に言及しているのか、それとも同一のエンティティ間の明示的な因果関係や構造的関係を記述しているのかを区別することができない。この制限は、構造的な推論を必要とするシナリオ、例えば、デバイスのトラブルシューティング・セッション全体を通じてエラー状態の進行を追跡する場合などにおいて極めて重要である。なぜなら、関連するコンテキストはクエリと語彙的な重なりを持たない一方で、深い構造的つながりを共有している可能性があるからである。

手法
本論文では、推論中に動的かつコンテキスト固有の知識グラフを構築することで、グラフ構造の埋め込みとテキストの意味論を統合するフレームワークであるKGERMAR（Knowledge Graph Enhanced Memory-Augmented Retrieval）を提案する。本システムは、以下の4つの統合されたコンポーメントによって構成される：

1. コンテキスト知識グラフの構築:

   • 抽出: システムは、BERT-large（CoNLL-2003でファインチューニング済み）を用いたリアルタイムの固有表現抽出（NER）と、BERT-base（TACREDでファインチューニング済み）を用いた関係抽出を実行する。
   • グラフ構築: スライディングウィンドウ内でエンティティと関係（例：CAUSEDBY, PARTOF, USEDFOR）を抽出する。品質を確保するため、エンティティの統合（文字列および埋め込みの類似性によるバリアントの統合）、関係の信頼度スコアリング（分類器の確信度と頻度の組み合わせ）、およびグラフフィルタリングを採用する。
   • 動的な性質: 静的なアプローチ（例：ConceptNet）とは異なり、このグラフは入力テキスト自体から構築されるため、既存の知識ベースには存在しないドメイン固有かつセッション固有のエンティティ（例：特定のプリンターモデルとその固有のエラーコード）を捉えることができる。

2. グラフ強化型埋め込み生成:

   • エンコーディング: 抽出されたエンティティは、マルチホップのメッセージパッシングを通じて個々の特性と構造的役割の両方を捉えるために、関係グラフ畳み込みネットワーク（R-GCN）を用いてエンコードされる。
   • 融合: これらのグラフ構造埋め込みは、クロスモーダル・アテンションを介して、テキストエンティティ埋め込み（ベースとなる言語モデルから派生）と融合される。これにより、表面的な意味論とより深い関係パターンを共に捉えるハイブリッドな表現が作成される。

3. マルチコンポーネント・メモリ・アーキテクチャ:

   • KGERMARは、3つの特化したメモリバンクを保持する：
     • コンテキスト・メモリ: 言語モデル層からのトークンレベルのキー・バリュー・ペアを格納する（ERMARに従う）。
     • セマンティック・メモリ: 重複のないチャンクの密なテキスト埋め込みを格納する（BGE-M3を使用）。
     • 構造的メモリ: 融合されたテキスト・グラフエンティティ埋め込みを格納する。
   • すべてのバンクはインクリメンタルに更新され、容量に達した場合はLRU（Least Recently Used）方式による退避管理が行われる。

4. ハイブリッド検索メカニズム:

   • 与えられたクエリに対し、システムは3つのバンクすべてから並列にトップK検索を実行する。
   • 検索信号は学習された融合重み（$\lambda_1, \lambda_2, \lambda_3$）を用いて結合され、最終的なスコアを算出する。
   • 検索された表現は、指定された上位レイヤーにおいて、リトリーバル・コーザル・アテンションを介して言語モデルに注入される。

学習目的
本モデルは、以下の4つの目的を同時に最適化するためにマルチタスク学習を採用している：

• 言語モデリング損失 ($L_{LM}$): 検索されたコンテキストを用いた次トークン予測。
• 知識グラフ埋め込み損失 ($L_{KG}$): R-GCNを訓練するためのバイナリ・クロスエントロピーによるリンク予測。
• 検索ランキング損失 ($L_{retrieval}$): 関連するコンテキストが不適切なものよりも高いスコアを得るようにするためのマージンベースのランキング。
• クロスモーダル・アライメント損失 ($L_{alignment}$): テキスト表現と構造的表現の間の整合性を確保する。

主な貢献

1. 動的な知識グラフの構築: 入力テキストからドメイン固有のエンティティと関係を推論中にリアルタイムで抽出するパイプラインにより、静的な知識ベースには存在しない専門的なコンテンツへの適応を可能にする。
2. ハイブリッド・マルチコンポーネント・メモリ: コンテキスト、セマンティック、および構造的表現を統合したアーキテクチャ。特に構造的メモリバンクは、純粋な意味的類似性よりも補完的であり、しばしばより識別力の高いエンティティ関係を捉える。
3. 実証的検証: 構造的知識の導入が長文脈モデリングにおいて大幅な改善をもたらすことを実証し、強力なベースラインと比較して最大8.5%低いパープレキシティと、2〜2.5倍優れたメモリ効率を達成した。

実験結果
本フレームワークは、SlimPajama（ファインチューニング用に84.7Kの例）および3つのベンチマークデータセット（WikiText-103, PG-19, Proof-pile）を用いて評価された。

• パープレキシティ: KGERMARは、コンテキスト長1Kから32Kトークンの範囲において、すべてのメモリ拡張型3Bモデルの中で最も低いパープレキシティを達成した。WikiText-103において、ERMARやMemLongを大幅に上回った（例：16Kトークン時において7.14 vs 7.80のパープレキシティ）。
• インコンテキスト学習 (ICL): 5つのNLUタスク（SST-2, MR, Subj, SST-5, MPQA）において、KGERMARは4-shot（78.6%）および20-shot（81.2%）の設定の両方で最高の平均精度を記録し、ERMARやMemLongを上回った。特に、エンティティレベルの識別を必要とするタスク（例：Subj, MPQA）において顕著な利得が見られた。
• メモリ効率: KGERMARは優れたメモリ効率を示し、ベースラインよりも2〜2.5倍低いメモリ使用量を維持した。32Kトークンの時点で、ERMARより36%少なく、MemLongより41%少ないピークメモリを使用した。トークンあたりの正規化メモリ量はコンテキスト長が増加するにつれて改善しており、標準的なアテンション機構の二次関数的なメモリ増大を緩和する能力を裏付けている。
• スケーリング挙動: K-GERMARはリアルタイムのグラフ構築により、1トークンあたりのレイテンシがわずかに高くなるものの、シーケンス長に対してほぼ線形なレイテンシ成長を示すため、長文脈デプロイメントにおいてERMARよりも予測可能性が高い。

意義
本論文は、KGERMARが固定された静的な知識ベースへの依存から、ドメイン適応型の知識表現へと移行することで、メモリ拡張型言語モデリングを進展させたことを主張している。入力コンテキストから動的に知識グラフを構築することにより、システムは意味的類似性だけでは捉えられないエンティティの状態や因果的な進行について推論することができる。このアプローチは、汎用的なハードウェア上での長文脈モデリングに対する実用的な解決策を提供し、高度な構造的推論が不可欠なテクニカルサポート、科学文献、法務文書の分析などのアプリケーションにおいて、高いモデリング品質と大幅なメモリオーバーヘッドの削減とのバランスを実現している。