見出し画像
Photo by eisukeai

人間のような無限コンテキストLLMのためのエピソード記憶

Yuya Fujiwara

概要

大規模言語モデル(LLM)は優れた能力を示していますが、長文脈の処理には課題があります。本研究では、人間の脳のエピソード記憶と事象認知の主要な側面をLLMに統合した「EM-LLM」という新しいアプローチを提案しています。これにより、計算効率を維持しながら実質的に無限の長さのコンテキストを効果的に処理できるようになりました。

主要なポイント

  1. EM-LLMはトークンの連続をベイズ的驚きとグラフ理論に基づく境界の精緻化を組み合わせてオンラインで一貫したエピソード事象に組織化します。

  2. 必要に応じて、これらの事象は類似性ベースと時間的に連続した取り出しを組み合わせた2段階のメモリプロセスを通じて効率的かつ人間らしい方法で関連情報にアクセスします。

  3. LongBenchデータセットでの実験では、EM-LLMが最先端のInfLLMモデルを上回る性能を示し、様々なタスクで全体的に4.3%の相対的な改善を達成しました。

  4. PassageRetrievalタスクでは33%の改善が見られました。

  5. 分析により、EM-LLMの事象分割と人間が知覚する事象との間に強い相関が示されました。

技術的詳細

  1. サプライズベースの分割

    • トークン $${x_t}$$ が潜在的な境界とみなされるのは、そのサプライズ値が閾値 $${T}$$ を超える場合:

      • $${- \log P(x_t|x_1, \ldots, x_{t-1}; \theta) > T \text{ with } T = \mu_{t-\tau} + \gamma\sigma_{t-\tau}}$$

  2. 境界の精緻化

    • グラフ理論的メトリクス(モジュラリティやコンダクタンス)を使用して、事象内の結束と事象間の分離を最適化

  3. メモリ取り出し

    • 類似性に基づくk-NN検索と時間的連続性を促進するメカニズムを組み合わせた2段階プロセス

  4. 計算複雑性

    • 全体的な複雑性: $${O(kn)}$$、ここで $${k}$$ は初期境界の数、 $${n}$$ はシーケンスの長さ

アルゴリズム

Algorithm 1 Event segmentation in KV cache
Input: tokens: List of tokens in the sequence
Input: T: Threshold for surprisal to identify initial boundaries
Input: f: Metric function to evaluate potential boundaries
Output: B: List of final boundary positions
1: B ← [i for i in range(length(tokens)) if −log(P(tokens[i])) > T] ▷ Boundary identification
2: for i in range(length(B)) do
3:     α, β = B[i], B[i + 1]
4:     B[i + 1] ← arg max β̂∈(α,β] f(A, {α, β̂}) ▷ Boundary refinement
5: end for
6: return B

このアルゴリズムは、EM-LLM（Episodic Memory for Large Language Models）における「イベント分割」のプロセスを示しています。具体的には、連続したトークン列をより意味のある「イベント」や「エピソード」に分割する方法を説明しています。以下、アルゴリズムの各部分を詳細に解説します：

  1. 入力:

    • tokens: 分析対象のトークン列

    • T: サプライズ（予測困難性）の閾値

    • f: 境界の品質を評価するためのメトリック関数

  2. 出力:

    • B: 最終的なイベント境界の位置リスト

  3. アルゴリズムの詳細:

    1. a. 初期境界の識別 (ライン1):

      • 各トークンに対し、そのトークンが出現する確率の負の対数（サプライズ値）を計算

      • サプライズ値が閾値Tを超えるトークンの位置を初期境界として選択

      • これにより、予測が困難（つまり「驚き」が大きい）箇所をイベントの切れ目として仮定

    2. b. 境界の精緻化 (ライン2-5):

      • 識別された各境界ペア（α, β）に対して反復処理

      • α と β の間で、メトリック関数 f を最大化する新しい境界位置 β̂ を探索

      • このステップでは、単純なサプライズだけでなく、より複雑な基準（例：トークン間の意味的類似性）を考慮

      • 元の境界 β を、最適な新しい位置 β̂ で置き換え

    3. c. 最終境界の出力 (ライン6):

      • 精緻化された境界位置のリストを返す

このアルゴリズムの重要な特徴は、2段階のアプローチを採用していることです：

  1. まず、サプライズに基づいて粗い初期分割を行います。これは、人間が予期せぬ情報に遭遇したときに新しいイベントを認識する傾向を模倣しています。

  2. 次に、より洗練された基準（メトリック関数 f）を用いて境界を微調整します。これにより、単なるサプライズだけでなく、意味的一貫性や文脈の流れなども考慮したより適切な分割が可能になります。

このようなイベント分割は、長い文脈を効率的に処理し、後の検索や利用を容易にするために重要です。人間の記憶システムを模倣することで、LLMがより自然で効果的な方法で長期的な情報を管理できるようになることを目指しています。

メトリック関数 f

メトリック関数 f に関する主要な数式を説明します。

モジュラリティ (Modularity)

$$
f_M(A^h, B) = \frac{1}{4m} \sum_{i,j} \left[A^h_{ij} - \frac{\sum_i A^h_{ij} \cdot \sum_j A^h_{ij}}{2m}\right] \delta(c_i, c_j)
$$

ここで:
$${m}$$ はグラフの総辺重み $${c_i}$$ はノード $${i}$$ が割り当てられたコミュニティ（エピソードイベント） $${\delta}$$ はクロネッカーのデルタ関数

  • $${i}$$はあるノードを指します。

  • $${j}$$は別のノードを指します。

  • $${\sum_{i,j}}$$ は全てのノードペアについての和を取ることを意味します。

コンダクタンス (Conductance)

$$
f_C(A^h, B) = \min_{S \in V} \frac{\sum_{i \in S, j \notin S} A^h_{ij}}{\min(\text{vol}(S), \text{vol}(V \setminus S))}
$$

ここで:
$${S = {b_i, b_i + 1, ..., b_{i+1}}}$$ は全ノード $${V = {b_1, b_1 + 1, ..., b_k}}$$ の部分集合 $${\text{vol}(S) = \sum_{i \in S, j \in S} A_{ij}}$$ $${\text{vol}(V \setminus S) = \sum_{i \notin S, j \notin S} A_{ij}}$$

  • $${i}$$はセットS内のノードを指します。

  • $${j}$$はS以外のノード（V \ S）を指します。

  • $${\sum_{i \in S, j \notin S}}$$ はSの内部から外部へのエッジについての和を取ることを意味します。

これらの数式は、イベント境界の最適化に使用されます。モジュラリティは高い値を、コンダクタンスは低い値を目指します。両方とも、イベント内の類似性を高め、イベント間の類似性を低くすることを目的としています。

イベント内の類似性を高め、イベント間の類似性を低くするとは

  1. 目的

    1. この目的は、メモリリコール（記憶の呼び出し）の効率を最大化するためです。PDFでは、イベント内の要素がクエリによって利用される可能性に基づいて、このアプローチが有効であると理論付けられています。

  2. グラフ理論的アプローチ

    1. この目的を達成するために、著者らはグラフクラスタリングの文脈でこの問題を表現しています。注意ヘッドの類似度行列を隣接行列として扱うことで、グラフ理論的な手法を適用しています。

  3. キー類似性

    1. 「類似性」は具体的に、注意機構のキーベクトル間の類似性を指します。PDFでは、ドット積類似度$${A^h_{ij} = \text{sim}(K^h_i, K^h_j)}$$ が使用されています。これは、高次元空間での意味的関係を捉えるのに効果的だとされています。

  4. イベントの定義

    1. 「イベント」は、連続したトークンのグループを指します。これらは最初に驚き（surprise）に基づいて分割され、その後、選択されたメトリック関数を使用して最適化されます。

  5. 最適化の目標

    1. イベント内（intra-event）の類似性を高める

      1. 同じイベント内のキーベクトル間の類似性を最大化します。これにより、意味的に関連する情報が同じグループにまとめられます。

    2. イベント間（inter-event）の類似性を低くする

      1. 異なるイベント間のキーベクトルの類似性を最小化します。これにより、各イベントが他のイベントと区別されやすくなります。

  6. メトリック関数

    1. この最適化は、モジュラリティやコンダクタンスなどのグラフクラスタリングメトリックを使用して行われます。これらのメトリックは、コミュニティ（この場合はイベント）内の結合の強さと、コミュニティ間の分離の度合いを測定します。

  7. 期待される効果

    1. このアプローチにより、各イベントが意味的にまとまりのある情報のグループとなり、同時に他のイベントとは明確に区別されるようになります。これは、後のメモリリコール時に、関連性の高い情報を効率的に取得するのに役立ちます。

結論

EM-LLMは、LLMの長文脈処理能力を大幅に向上させ、人間の記憶メカニズムとAIの橋渡しをする新しい研究の道を開きました。このアプローチは、LLMの性能を向上させるだけでなく、人間の記憶に関する仮説をテストするための計算フレームワークも提供します。

この記事が気に入ったらサポートをしてみませんか？