[論文レビュー] In Search of Needles in a 11M Haystack: Recurrent Memory Finds What LLMs Miss

Yuri Kuratov, Aydar Bulatov+4 more|arXiv (Cornell University)|Feb 16, 2024

Genetics, Bioinformatics, and Biomedical Research被引用数 7

ひとこと要約

本論文は ultra-long documents のための BABILong ベンチマークを導入し、再帰的メモリ transformer (RMT/RMT-R) が最大 11 million tokens を処理でき、長文コンテキスト課題で GPT-4 および RAG を上回ることを示します。

ABSTRACT

This paper addresses the challenge of processing long documents using generative transformer models. To evaluate different approaches, we introduce BABILong, a new benchmark designed to assess model capabilities in extracting and processing distributed facts within extensive texts. Our evaluation, which includes benchmarks for GPT-4 and RAG, reveals that common methods are effective only for sequences up to $10^4$ elements. In contrast, fine-tuning GPT-2 with recurrent memory augmentations enables it to handle tasks involving up to $11 imes 10^6$ elements. This achievement marks a substantial leap, as it is by far the longest input processed by any neural network model to date, demonstrating a significant improvement in the processing capabilities for long sequences.

研究の動機と目的

• 現在の能力を超えた極端に長い文脈で NLP モデルを評価する必要性を動機づける。
• 分散された事実を持つ長文脈 QA のスケーラブルなベンチマーク(BABILong)を提案する。
• GPT-4、RAG、再帰的メモリ拡張トランスフォーマーを超長入力で評価する。
• 再発による自己リトリーブが既存モデルを超える文脈処理を拡張できることを実証する。

提案手法

• BABILong を導入する。PG19/Wiki の背景文からタスク文を隠す ultra-long context ベンチマークとして。
• GPT-4-Turbo (128k window) と Mistral (32k window) を用いて、成長する文脈サイズに対する基準 LLM の性能を評価する。
• 16k トークンタスクで GPT-3.5 をファインチューニングし、文脈拡張時の性能を評価する。
• 長い系列を線形スケーリングで処理する Recurrent Memory Transformer (RMT) と RMT の自己リトリーブ版(RMT-R) を実装する。
• past memory states のリトリーバルを追加し、過去のセグメントへの注意のようなアクセスを模倣する。
• 入力を512-token に分割し、セグメント間で memory tokens を維持する実験を行い、注意・メモリの使用を分析する。
• RAG のリトリーバル拡張生成を、GPT-4 と比較し、memory-augmented アプローチと併用する。embeddings には FAISS/LangChain を使用（text-embedding-ada-002）。

実験結果

リサーチクエスチョン

• RQ1 現在の LLM は context length が 10^4–10^5 トークン、ひいてはそれ以上に成長しても needle-in-a-haystack タスクを解けるか。
• RQ2 リトリーバル拡張型およびメモリ拡張型アーキテクチャは、長大な文脈で標準のアテンションが非効率になる領域でも性能を維持できるか。
• RQ3 RMT および RMT-R はシーケンス長の増加とともにどのようにスケールし、長文コンテキスト QA で GPT-4 および RAG と比較してどうか。
• RQ4 非常に長い文書上での多段推論に対する memory-based リトリーバルの影響は。
• RQ5 反復的なメモリアプローチが、百万トークン規模の極端に長いシーケンスにおいて大規模 LLM を上回るような regime が存在するか。

主な発見

• GPT-4 および Mistral の性能は、ウィンドウが大きくても文脈長が増すにつれて低下する。
• RMT および RMT-R は長いシーケンスで GPT-4 を大幅に上回り、実験で up to 11 million tokens を処理可能。
• RMT-R（過去のメモリ状態のリトリーバル付き）は、 retrieved past states によってさらに改善され、非常に長い入力でも性能を維持。
• リトリーバル拡張生成 (RAG) は BABILong タスクに対して制限的またはタスク依存的な利点を示し、特定の長文コンテキスト設定ではメモリベースの方法より効果が低い場合がある。
• メモリ機構を持つモデルは、訓練の境界を超えたシーケンス長（160k トークン以上、RMT-R は 10M+ トークンにおいても追随）で一貫した性能を示し、一般化する。