[論文レビュー] Needle in the Haystack for Memory Based Large Language Models

Elliot Nelson, Γεώργιος Κόλλιας+3 more|arXiv (Cornell University)|Jul 1, 2024

Topic Modeling被引用数 5

ひとこと要約

この論文は、外部のCPUに保持されるメモリを持つ1.3Bパラメータの Larimar モデルを示し、長い文脈リコールを可能にする（100Kトークン超まで）タスク特異的な訓練なしで、パスキーと干し草の山から針を探すタスクでベースラインを上回る。

ABSTRACT

Current large language models (LLMs) often perform poorly on simple fact retrieval tasks. Here we investigate if coupling a dynamically adaptable external memory to a LLM can alleviate this problem. For this purpose, we test Larimar, a recently proposed language model architecture which uses an external associative memory, on long-context recall tasks including passkey and needle-in-the-haystack tests. We demonstrate that the external memory of Larimar, which allows fast write and read of an episode of text samples, can be used at test time to handle contexts much longer than those seen during training. We further show that the latent readouts from the memory (to which long contexts are written) control the decoder towards generating correct outputs, with the memory stored off of the GPU. Compared to existing transformer-based LLM architectures for long-context recall tasks that use larger parameter counts or modified attention mechanisms, a relatively smaller size Larimar is able to maintain strong performance without any task-specific training or training on longer contexts.

研究の動機と目的

• LLMにおける長い文脈の検索性能の向上を動機づける。
• テスト時の長い文脈適応のため、Larimarと統合された外部メモリ機構を提案・評価する。
• GPUメモリ使用量を増やさずに、非常に長い文脈にスケールするメモリの読み出し/書き込み操作を実証する。
• メモリ読み出しがデコーダを正しい出力を生成するよう条件づけることを示す。
• 長い文脈タスクにおけるCPUベースの外部メモリの実用性と制限を論じる。

提案手法

• 最小二乗法メモリ更新で書き込まれる外部連想メモリを備えた Larimar アーキテクチャを使用する。
• エンコーディングと書き込みキーを用いて、文脈セグメントをメモリに書き込む。
• プレフィックスまたはクエリエンコーディングから導出された読み取りキーを使用して、メモリから読み出す。
• メモリの読み出しを z_read = w M として計算し、この読み出しでデコーダを条件付ける。
• 最も近い近傍写像が書き込みと読み出しを結ぶよう、固定キー・メモリを介してキーを制御する。
• デコードをGPUに保持しつつ、文脈に合わせてメモリサイズを拡張するため、CPU上でメモリを動作させる。）

実験結果

リサーチクエスチョン

• RQ1 外部に保持され、動的に更新可能なメモリは、タスク特異的な訓練なしでLLMの長い文脈のリコールを改善できるか？
• RQ2 テスト時のプレフィックスベースのキー計算は、非常に長い文脈（100K–1M トークン）での信頼できる検索をどのように可能にするか？
• RQ3 長い文脈タスクのためにメモリをCPUにオフロードする際の、メモリサイズ・待機時間・GPU使用量のトレードオフは？
• RQ4 パスキーと干し草の山の針タスクにおいて、Larimarはベースラインの長文検索モデルとどう比較されるか？

主な発見

コンテキスト	3 桁	4 桁	SF
Larimar 137K	0.95	0.64	1.0
Larimar (no prefix) 137K	0.88	0.14	0.0
Mistral 7B v0.2 24K	0.66	0.62	0.80
Phi-3-mini-128K 100K	0.27	0.26	0.37

• Larimarは、100Kトークンを超える長い文脈に対して強いリコールを維持し、1.3Bパラメータのモデルでタスク特異的な訓練なしで実現する。
• メモリの読み出しは、全体の文脈がオフGPUに格納されていてもデコーダを正しい出力へと導くことができる。
• ベースライン（Mistral 7B、Phi-3-Mini-128K）と比較して、Larimarは同程度または小さいモデルサイズで、パスキーと干し草の山の針タスクでより優れたリコールを示す。
• キーを書き込む際のプレフィックスベースのアプローチは、より長く複雑なニードルでリコールを改善する。
• CPUベースの外部メモリは、GPUメモリ使用量を増やすことなく長い文脈へ拡張可能にする。