【LLM論文を読む】人間的なエピソード記憶を持つコンテキスト無限のLLM

  • 言語モデルの「記憶」を管理する方法としては、1)追加学習や知識編集によってLLM自体の知識更新を試みる方法と、2)外部の記憶データから必要情報を適時検索してモデルに渡す方法、の2つの方向性があります。
  • この論文は後者の系統に属する研究で、人間のエピソード記憶のあり方を参照し、外部データを効果的に整理・検索することを目指した研究のようです。
  • タイトルは"Human-like Episodic Memory for Infinite Context LLMs"で、Huaweiの研究所とUCLの研究者らによって2024年7月に投稿されています。

arxiv.org

要旨

大規模言語モデル(LLM)は目覚ましい能力を示しているが、広範なコンテキストの処理にはまだ苦戦しており、長いシーケンスにわたって一貫性と正確さを維持する能力には限界がある。これとは対照的に、人間の脳は、生涯にわたる膨大な時間スケールのエピソード経験を整理し、検索することに優れている。本研究では、人間のエピソード記憶と事象認知の主要な側面をLLMに統合し、計算効率を維持しながら実質的に無限のコンテキスト長を効果的に扱えるようにした新しいアプローチ「EM-LLM」を紹介する。EM-LLMは、ベイジアン・サプライズとグラフ理論に基づく境界精密化の組み合わせにより、トークンのシーケンスをオンライン方式で首尾一貫したエピソードイベントに整理する。必要に応じて、これらのイベントは、類似性に基づく検索と時間的に連続した検索を組み合わせた2段階の記憶プロセスによって検索され、関連情報への効率的で人間的なアクセスを実現する。LongBenchデータセットを用いた実験では、EM-LLMの優れた性能が実証され、様々なタスクにおいて、最新のInflLLMモデルを4.3%の相対的改善で上回った。さらに、我々の分析により、EM-LLMのイベントセグメンテーションと人間が認識するイベントとの間に強い相関があることが明らかになり、この人工的なシステムと生物学的な対応物との間の橋渡しを示唆している。本研究は、LLMの拡張文脈処理能力を向上させるだけでなく、人間の記憶メカニズムを探求するための計算フレームワークを提供し、AIと認知科学の学際的研究に新たな道を開くものである。

Claudeによる解説

Q.内容をわかりやすく説明して

この論文は、大規模言語モデル(LLM)のロングコンテキスト処理能力を向上させる新しいアプローチ「EM-LLM」を提案しています。主なポイントは以下の通りです:

1. 問題点:
   - 現在のLLMは長いコンテキストの処理に課題がある。
   - 学習時のコンテキスト長を超えると性能が低下する。

2. 提案手法 (EM-LLM):
   - 人間の認知科学、特に episodic memory(エピソード記憶) と event cognition(イベント認知) の知見を取り入れた。
   - 入力を意味のあるイベントに分割し、効率的に保存・検索する仕組みを導入。

3. 主な特徴:
   - Surprise-based segmentation: 予測困難な箇所でイベントを区切る。
   - Boundary refinement: グラフ理論を用いてイベント境界を最適化。
   - Two-stage retrieval: 類似性と時間的連続性に基づく記憶検索。

4. 結果:
   - LongBenchデータセットで既存の最高性能モデル(InfLLM)を上回る性能を達成。
   - 特に PassageRetrieval タスクで33%の性能向上。
   - 人間の知覚するイベント区切りとの高い相関を示した。

5. 意義:
   - 事実上無限の長さの文脈を扱えるLLMの可能性を示した。
   - 認知科学とAIを橋渡しし、人間の記憶メカニズムの計算モデルを提供。

6. 将来の展望:
   - より高度なイベント分割アルゴリズムの探求。
   - マルチモーダルタスクへの拡張。
   - 強化学習や継続学習への応用。

この研究は、LLMの長文脈処理能力を大幅に向上させただけでなく、人間の認知プロセスをAIに取り入れる新しいアプローチを示した点で重要です。

Q.コンテキスト処理方法を詳しく

EM-LLMのコンテキスト処理方法をもう少し詳しく説明しましょう:

1. コンテキストの分割:
   - コンテキストを「イベント」と呼ばれる意味のあるセグメントに分割します。
   - この分割は、まず「サプライズ」(モデルの予測誤差)に基づいて行われ、その後グラフ理論的な手法で最適化されます。

2. イベントの保存:
   - 分割されたイベントは、各イベント内のトークンのkey-valueペアとともに保存されます。

3. メモリ検索:
   - 新しいトークンを生成する際、関連するイベントを検索します。
   - 検索は2段階で行われます:
     a) 類似性バッファ: 現在のクエリと各イベントの代表トークンとの類似性に基づいてk_s個のイベントを選択します。
     b) 連続性バッファ: 選択されたイベントの前後のイベントもk_c個取得し、時間的な文脈を維持します。

4. コンテキスト構成:
   - 選択されたイベントは、LLMの拡張されたコンテキストウィンドウに追加されます。
   - このコンテキストは以下のように構成されます:
     - 初期トークン(128トークン)
     - 類似性バッファからのイベント
     - 連続性バッファからのイベント
     - ローカルコンテキスト(最新のトークン)

5. LLMの処理:
   - 拡張されたコンテキストを使用して、LLMが新しいトークンを生成します。

この方法により、EM-LLMは非常に長いコンテキストを効率的に処理し、関連性の高い情報を保持しながら、人間の記憶メカニズムに似た方法で情報を取り出すことができます。

Q.ベクトルデータベースは使ってないですね?

はい、その通りです。EM-LLMの手法では、ベクトルデータベースのような外部のデータ構造は使用していません。代わりに、情報は以下のような形で保持されています:

1. テキストベースの保存:
   - 情報は、元のテキスト形式で保存されます。具体的には、トークン化されたテキストの形で保持されています。

2. Key-Valueペア:
   - 各イベント(テキストセグメント)は、そのイベント内のトークンに対応するkey-valueペアとして保存されます。
   - これらのkey-valueペアは、Transformerモデルの注意機構で使用されるものと同じです。

3. イベント構造:
   - テキストは「イベント」と呼ばれる意味のあるセグメントに分割されます。
   - 各イベントは、そのイベントに含まれるトークンのkey-valueペアの集合として表現されます。

4. 内部メモリ:
   - これらのイベントとそれに対応するkey-valueペアは、モデルの内部メモリに保持されます。
   - 外部のデータベースやインデックスは使用せず、すべての情報はモデル内部で管理されます。

5. 検索メカニズム:
   - 関連イベントの検索は、保存されたkey-valueペアを直接使用して行われます。
   - 類似性の計算は、クエリと各イベントの代表トークンのkey間で行われます。

この方法の利点は、外部のデータ構造に依存せずに長いコンテキストを処理できること、およびTransformerモデルの既存のメカニズム(key-value注意機構)を直接利用できることです。これにより、モデルの既存の能力を最大限に活用しながら、長いコンテキストを効率的に処理することができます。

Q.この論文で使っているLongBenchデータセットについて教えて

はい、LongBenchデータセットについて説明します。

LongBenchは長文コンテキスト処理能力を評価するためのベンチマークデータセットです。論文の中で以下のような特徴が挙げられています:

1. 長文処理を必要とする多様なタスクが含まれています。具体的には以下のようなタスクが挙げられています:

- NarrativeQA: 長い物語文の質問応答
- MultiNews: 複数の記事の要約
- Qasper: 学術論文からの質問応答 
- TREC: 質問分類
- MultiFieldQA: 複数分野にまたがる質問応答
- TriviaQA: 雑学クイズ形式の質問応答
- HotpotQA: 複数文書を参照する必要のある質問応答
- SAMSum: 会話の要約
- 2WikiMQA: Wikipediaの複数記事にまたがる質問応答
- PassageRetrieval: 要約から元の段落を特定するタスク
- Musique: 複数のステップを要する質問応答
- LCC: 長文の分類
- GovReport: 政府報告書の要約
- RepoBench-P: プログラミング関連の質問応答
- QMSum: 会議録の要約

2. これらのタスクは、長い文脈の理解、複数文書間の関連付け、長文の要約など、長文処理の様々な側面を評価します。

3. 論文の実験では、このLongBenchを使ってEM-LLMの性能を評価し、特にPassageRetrievalタスクで33%以上の大幅な改善を示しています。

4. LongBenchは長文処理モデルの性能を包括的に評価するための標準的なベンチマークとして使用されており、EM-LLMの性能をInfLLMなど他のモデルと直接比較することができます。

このように、LongBenchは長文コンテキスト処理における言語モデルの能力を多角的に評価するための重要なデータセットとして位置づけられています。

Q.驚きベースの分割(Surprise-based segmentation)をもっと詳しく

Surprise-based segmentationは、EM-LLMの重要な要素の1つです。この手法は人間の認知プロセスにヒントを得ており、予測が困難な(つまり「驚き」のある)箇所でテキストを区切ることで、意味のあるイベントに分割します。以下にその詳細を説明します:

1. 基本概念:
   - LLMの予測確率を利用して、各トークンの「驚き」の度合いを計算します。
   - 「驚き」が一定のしきい値を超えた箇所をイベントの境界とみなします。

2. 数学的定義:
   - あるトークン x_t の驚きは、負の対数尤度で定義されます:
     - −log P(x_t | x_1, ..., x_t-1; θ)
   - ここで P は LLM によるトークンの予測確率、θ はモデルのパラメータです。

3. 境界の決定:
   - 動的なしきい値 T を使用:T = μ_t-τ + γσ_t-τ
   - μ_t-τ と σ_t-τ は、過去 τ トークンの驚きの平均と標準偏差
   - γ はスケーリング因子

4. アルゴリズム:
   - 各トークンの驚きを計算
   - 驚きがしきい値 T を超えるトークンを潜在的な境界として識別
   - これにより、初期のイベント分割が得られる

5. 利点:
   - コンテキストの変化に適応的
   - 人間の知覚するイベント境界と高い相関を示す
   - 計算コストが低い(既存のLLM出力を利用)

6. 補完的手法:
   - Boundary refinement と組み合わせることで、さらに洗練された分割が可能

7. 認知科学との関連:
   - 人間が予測誤差や驚きの大きな瞬間で経験を分割するという理論と一致
   - エピソード記憶の形成プロセスを模倣

このアプローチにより、EM-LLMは入力テキストを意味のある単位に効率的に分割し、後の検索や利用を容易にします。これは特に長い文脈を処理する際に重要で、LLMの性能向上に大きく貢献しています。

Q.境界精緻化(Boundary refinement)をもっと詳しく

Boundary refinementは、Surprise-based segmentationで得られた初期のイベント境界をさらに最適化するプロセスです。この手法はグラフ理論の概念を応用しており、イベント内の一貫性を高め、イベント間の区別を明確にすることを目的としています。以下にその詳細を説明します:

1. 目的:
   - イベント内の類似性(intra-event similarity)を最大化
   - イベント間の類似性(inter-event similarity)を最小化

2. グラフ理論的アプローチ:
   - 注意機構のキー(key)ベクトル間の類似度行列を隣接行列として扱う
   - 各トークンをグラフのノードとみなし、類似度をエッジの重みとする

3. 主要な評価指標:
   a. Modularity(モジュラリティ):
      - コミュニティ(イベント)構造の質を測定
      - 高いモジュラリティは、イベント内の強い結合とイベント間の弱い結合を示す
   b. Conductance(コンダクタンス):
      - コミュニティ間の接続の少なさを測定
      - 低いコンダクタンスは、より明確に区別されたイベントを示す

4. アルゴリズムの流れ:
   - Surprise-based segmentationで得られた初期境界を開始点とする
   - 隣接する境界間で、選択した評価指標(モジュラリティまたはコンダクタンス)を最適化する新しい境界位置を探索
   - この過程を全ての初期境界に対して繰り返す

5. 計算複雑性:
   - O(kn)、ここでkは初期境界の数、nはシーケンス長

6. 利点:
   - より一貫性のある、意味のあるイベント分割が可能
   - 後の検索プロセスの効率と精度を向上させる
   - グラフ理論の既存の知見を活用できる

7. 実験結果:
   - 単純なSurprise-based segmentationよりも高いパフォーマンスを示す
   - 人間の知覚するイベント境界との相関がさらに高まる

8. 将来の展望:
   - より高度なグラフクラスタリングアルゴリズムの探索
   - 階層的なイベント構造の形成への拡張の可能性

Boundary refinementは、EM-LLMの性能向上に大きく寄与しており、特に複雑な時間的推論を必要とするタスクで効果を発揮します。この手法により、LLMはより人間らしい方法で長い文脈を理解し、処理することが可能になっています。

Q.二段階検索(Two-stage retrieval)をもっと詳しく

EM-LLMの Two-stage retrieval プロセスは、人間の記憶検索メカニズムに着想を得た重要な特徴です。以下にその主な要素と特徴をまとめます:

1. 目的:
- 関連性の高いエピソードイベントを効率的に選択し、LLMのコンテキストウィンドウに組み込む
- 類似性と時間的関係の両方を考慮した検索を実現する

2. 第1段階: 類似性ベースの検索
- k-NN (k-Nearest Neighbors) 検索を使用
- 現在のクエリと各イベントの代表トークンとのドット積類似度に基づいて検索
- ks個のイベントを取得し、「類似性バッファ」を形成

3. 第2段階: 時間的連続性の維持
- 「連続性バッファ」を導入(サイズ kc のキュー構造)
- 検索されたイベントの近傍イベント(元のシーケンスで±n位置以内)をこのバッファにエンキュー
- 時間的文脈を維持し、LLMの「誘導」注意ヘッドが連続性と非対称性効果を示すことを可能に

4. 全体の構成:
- 合計 k = ks + kc + 2 個のイベントがコンテキストウィンドウに追加される
- 関連性と時間的関係のバランスを取る

5. 特徴:
- 人間の記憶検索パターンに類似(類似性ベースと時間的に連続した検索の組み合わせ)
- キュー構造により、新しいイベントが処理されるにつれて時間的文脈が自然に減衰

6. 効果:
- LLMが長期的な文脈を効果的に利用できるようになる
- 特に複雑な時間的推論を必要とするタスクでパフォーマンスが向上

7. 実装の詳細:
- 大規模なメモリストアに対しては、効率性を維持するために近似k-NNを利用

このTwo-stage retrievalプロセスは、EM-LLMが実質的に無限のコンテキスト長を扱えるようにする重要な要素であり、人間の記憶システムの特徴を模倣しつつ、LLMの性能を大幅に向上させることが示されています。

メモ

  • 相変わらずテクニカルな部分は訳わかっていませんが、要は工夫を重ねたRAGだと理解しました。
    • コンテキストを全部読み込むのではなく、あらかじめイベント単位に整理しておき、入力されたプロンプトと類似性の高いイベント+時間的連続性の高いイベントだけを残してモデルに渡す。
  • 残念ながら実装は公開していないっぽいのですが、この論文で「既存のロングコンテキスト処理のSOTA」として言及されている「InfLLM」はGithubに公開されていました。試すならこっちでしょうか。

github.com