💬 NLP

PReM: Learning What to Preserve and When to Refresh for Context Compression

PReMは、専用のメモリレイヤーと特殊トークンを介して内部のレイヤーごとのKVメモリを動的に学習・保持および更新することで、外部の圧縮器に依存したり時期尚早な保持決定を行ったりすることなく、効率的な長文脈推論を可能にする新しいコンテキスト圧縮フレームワークである。

原著者： Bohan Yu, Lei Shen, Chenxi Zhou, Chen Han, Junlin Liu, Wenbo Su, Yu Cheng, Bo Zheng

公開日 2026-07-17

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Bohan Yu, Lei Shen, Chenxi Zhou, Chen Han, Junlin Liu, Wenbo Su, Yu Cheng, Bo Zheng

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは、巨大で複雑なミステリーを解こうとしているところだと想像してください。ただし、単一の手がかりではなく、32,000ページもの文書が入った図書館を渡されたのです。あなたの脳（あるいはこの場合は、大規模言語モデルと呼ばれるコンピュータプログラム）は、質問に答えるために必要な特定の事実を見つけ出すべく、そのすべてを読み通さなければなりません。問題は、その図書館全体を一度に「ワーキングメモリ（作業記憶）」の中に保持しておくことが、非常に重く、時間がかかるということです。それは、レースの最中に図書館をバックパックに入れて走ろうとするようなものです。やがて疲れ果て、必要な本を落としてしまうかもしれません。

これを解決するために、科学者たちはこれまで主に2つのトリックを試してきました。1つ目は「圧縮」です。これは、読み始める前に、図書館を極めて小さな要約へと縮小しようとする試みです。2つ目は「キャッシュ」です。これは、図書館全体は保持しておきつつ、最も人気のあるページだけを見るようにする方法です。しかし、どちらにも欠点があります。それは、パズルを解き始める「前」に、何を残すかの決定を下してしまうという点です。もしミステリーが途中で変化し、捨ててしまった本の巻末にあるページが突然必要になったとしても、あなたは立ち往生してしまいます。決定はあまりにも早くになされてしまったため、戻ってそのページを掴み取ることができないのです。この論文は、静的な要約ではなく、物語が進むにつれて自らを更新できるメモリシステムが必要であると示唆する、新しい方法を紹介しています。

ここで、PReM（Preserve and Refresh Memory：記憶の保持と更新）が登場します。これは、モデルのメモリを静的なファイルキャビネットではなく、動的で生きているノートブックのように扱う、巧妙な新しいフレームワークです。核心となるアイデアはシンプルですが強力です。一度だけ残すものを決めるのではなく、「今この瞬間」残すべきものを決め、物語が求めた瞬間にそれらを入れ替える準備をしておくのです。

PReMの仕組みを、楽しい比喩を使って説明しましょう。あなたがアシスタントのチーム（AIの異なるレイヤー）と共に事件を解決する探偵だと想像してください。従来の方法では、チームは事件の全ファイルを読み、お気に入りの10ページを選び出し、残りは金庫にロックしてしまいます。もし後で金庫の中のページが必要になっても、もう手遅れです。

PReMはルールを変えます。事件の全ファイルを保持しますが、それを小さな「チャンク（塊）」（章のようなもの）に整理します。探偵が解決策を書き始めると、特別な「メモリ・トークン」（魔法のしおり、記号では ⟨m⟩ と表記）が信号として機能します。探偵がこのしおりを書くとき、それは「待て！机の上にあるものを再考する必要がある」と言っているようなものです。

その瞬間、専用の「メモリ・マネージャー」（AI内の特別なレイヤー）が目を覚まします。それは現在の質問と魔法のしおりを確認し、すべての章を素早くスキャンします。そしてこう判断します。「よし、次のステップの物語には、第3章と第12章が絶対に必要だ。それ以外は今は忘れていい」。そして、不要な部分を小さな要約へと圧縮しながら、机の上の古い章を新しいものへと入れ替えるのです。これは、物語の「前」ではなく、執筆の「最中」に行われます。

論文は、この「リフレッシュ」メカニズムがいかにゲームチェンジャーであるかを証明しています。32,000トークンのコンテキスト（かなりの量のテキストです！）を用いたテストにおいて、PReMはメモリを16倍、さらには32倍に圧縮することに成功しました。一度に机の上に置いておく情報がこれほど少ないにもかかわらず、図書館をすべて開いておこうとするモデルよりも、実際により優れた回答を導き出したのです。例えば、70億パラメータのモデルにおいて、PReMは既存の最高の手法と比較して、回答の正確性を最大12.55ポイント向上させました。

研究者たちはまた、「メモリ・マネージャー」をAIの脳のどこに配置すべきかについても興味深い発見をしました。彼らは、メモリ・マネージャーをAIの脳の初期、中間、または後半のどこに配置するかをテストしました。その結果、中間層や後半のレイヤーの方が、何を残すべきかを判断するのに非常に優れている一方で、初期のレイヤーは学習する能力が低すぎることを発見しました。これは、新米（初期レイヤー）に判断させるのではなく、ベテランの刑事（中間層）に、どの証拠が関連しているかの判断を下させるようなものです。

PReMの驚くべき発見の一つは、PReMを用いた小規模なモデル（30億パラメータ）が、他の圧縮技術を用いた大規模なモデル（70億パラメータ）を凌駕できることです。これは、メモリを「いつ」更新すべきかを知ることは、脳がいかに大きいかと同じくらい重要であることを示唆しています。

また、この論文はいくつかのアイデアを否定しています。単にAIの自然なアテンション（それが現在何を考えているかに注目すること）を利用するだけでは、何を残すべきかを判断するには不十分であり、専用の訓練されたシステムが必要であることを示しています。また、テキストを最初に一度だけ圧縮して、あとはうまくいくことを期待するという手法も、特に複雑で多段階の質問においては、結果が悪化することを明らかにしました。

要約すると、PReMは、効率的なAIの未来とは単にデータを縮小することではなく、「呼吸するメモリ」を構築することであることを示唆しています。それは、会話の現在のステップに必要な証拠を保持する方法を学び、次のステップのために新しい証拠と入れ替えるために、いつ「リフレッシュ」ボタンを押すべきかを正確に理解します。こうすることで、ミステリーを解くために必要な手がかりを忘れることなく、AIを高速かつ効率的に保つことができるのです。

技術サマリー: PReM (Preserve and Refresh Memory)

1. 問題提起

大規模言語モデル（LLM）における効率的な長文脈推論は、計算コストの削減と、必要な文脈的証拠へのアクセスの維持という、根本的なジレンマに直面しています。文脈長が増加するにつれ、Key-Value (KV) キャッシュが増大し、デコーディングのレイテンシとメモリ使用量が大幅に増加します。

これを軽減するための既存のアプローチ（KVキャッシュ圧縮やコンテキスト圧縮など）には、主に2つの制限があります：

静的な証拠保持 (Static Evidence Retention): ほとんどの手法は、どの文脈を保持するかについて、初期段階で一度だけ決定を下します。この静的な圧縮は、マルチステップの推論において、異なる推論ステップが異なる証拠を必要とする場合の変化するニーズに適応できません。
分離された証拠選択 (Decoupled Evidence Selection): 外部の圧縮器に依存する手法は、圧縮プロセスと生成プロセスを分離しています。これは追加の計算を導入し、特定の推論ステップに対して動的に証拠選択をカスタマイズすることを妨げます。

本論文は、効果的な圧縮とは単なるサイズの削減を超えたものであるべきだと主張しています。すなわち、有用な証拠を動的に保持し、後の推論ステップが異なる情報を必要とする際に、それをリフレッシュ（更新）できなければなりません。

2. 手法: PReM

PReM (Preserve and Refresh Memory) は、長文脈をLLM内部のレイヤーごとのKVメモリとして維持するコンテキスト圧縮フレームワークです。外部の圧縮器に頼るのではなく、PReMはモデルの内部状態に基づいて、何を保持し、いつリフレッシュするかを決定することを学習します。

コアメカニズム

レイヤーごとのKVメモリ表現: 長文脈は、モデル内に格納されたレイヤーごとのKV状態 ( $M_\ell$ ) としてエンコードされます。これらの状態は、連続したチャンクに分割されます。
専用メモリレイヤー: 特定のレイヤー ( $\ell_m$ $ℓ_{m}$ ) が、メモリ選択の決定を行うために指定されます。このレイヤーは、コンテキストチャンクのスコアを算出するために、専用の選択クエリ投影行列 ( $W^c_\ell$ $W_{ℓ}^{c}$ ) を使用します。
- チャンク選択: モデルは、現在のクエリ（質問またはリフレッシュ・トークンから）に基づいてチャンクのスコアを算出します。
- Preserve-and-Pool (保持とプーリング): 選択されたチャンクはトークンレベルの粒度で保持されます。選択されなかったチャンクは、単一の平均化されたKV代表値へと圧縮されます。
メモリ・トークン $\langle m \rangle$ : 特殊なトークンが、メモリをリフレッシュするための明示的な内部信号として機能します。モデルが $\langle m \rangle$ を生成すると、メモリレイヤーはコンテキストチャンクを再評価し、新しい関連する証拠を選択し、後続のレイヤーのための圧縮されたKVメモリを更新します。これにより、モデルは現在の推論ステップに基づいて、いつリフレッシュするかを制御できます。

学習: フェーズ分離型リフレッシュ学習 (Phase-Separated Refresh Training)

メモリ選択と生成を整合させるようにモデルを訓練するために、PReMはフェーズ分離型リフレッシュ学習戦略を導入しています。各推論ステップは、2つの明確なフォワードパスに分割されます：

メモリ選択フェーズ: モデルは入力（質問または $\langle m \rangle$ ）を処理して、チャンクのスコアを計算し、上位 $k$ 個の証拠チャンクを選択します。このフェーズは、モデルが証拠ではないチャンクよりも証拠チャンクを高くランク付けするように促すメモリ選択損失 ( $L_{ms}$ ) によって監督されます。
メモリ条件付き生成フェーズ: モデルは、選択フェーズの結果として得られた新しくリフレッシュされたメモリ状態に基づき、回答セグメント ( $y_s$ ) を生成します。

訓練と推論の間の整合性を確保するため、訓練目的関数には境界クロスエントロピー項 ( $L_{bd}$ ) が含まれます。これらの項は、選択フェーズと最初の回答トークン、および回答セグメントの終了と次の $\langle m \rangle$ トークンを結びつけ、モデルがメモリ更新と生成の間の自己回帰的なリンクを学習することを保証します。

3. 主な貢献

定式化: 本論文は、コンテキスト圧縮を、学習されたメモリ保持およびリフレッシュの問題として再定義し、外部の圧縮器を不要にしました。
アーキテクチャ: PReMは、専用のメモリレイヤーと特殊な $\langle m \rangle$ トークンを導入することで、LLMの内部KV状態内での、動的かつステップ固有のメモリ更新を可能にします。
学習戦略: フェーズ分離型リフレッシュ学習の提案により、メモリ選択と生成を訓練中に分離し、モデルがそのステップに必要な特定のメモリ状態に基づいて生成を行うように学習させます。
実証的検証: 学習された動的なメモリ更新が、静的な圧縮やフルコンテキスト・プロンプティングよりも優れていることを示す包括的な実験を行いました。

4. 実験結果

実験は、32Kトークンのコンテキストを用いて、Qwen2.5-Instructの3Bおよび7Bバリアントに対して、6つのベンチマーク（TriviaQA, SQuAD, NaturalQuestions, 2WikiMQA, HotpotQA, MuSiQue）で行われました。

ベースラインとの比較: 16倍および32倍の圧縮率の両方において、PReMは、KVキャッシュ手法（SnapKV, StreamingLLM, CAKE）やコンテキスト圧縮手法（LongLLMLingua, EXIT, ICAE）を含む8つの強力なベースラインを上回りました。
- 3Bバックボーンの16倍圧縮において、PReMは最強のベースラインに対し、平均Exact Match (EM) で 5.30 ポイント、F1で 5.15 ポイント向上しました。
- 7Bバックボーンの32倍圧縮において、改善幅は EM で 8.20 ポイント、F1 で 9.89 ポイントに達しました。
フルコンテキストとの比較: PReMは、3Bおよび7Bモデルの両方において、直接的な32Kトークン・プロンプティング（圧縮なし）よりも優れた性能を示しました。これは、学習されたメモリ選択が、モデルにフィルタリングされていないフルコンテキストをそのまま露出させるよりも効果的であることを示唆しています。
小型モデルの優位性: PReMの3Bモデルは、公開されている7Bソフトコンテキスト圧縮ベースライン（例：Activation Beacon）を上回りました。これは、学習された保持とリフレッシュが、より小さなバックボーンのスケールを補完できることを示しています。
効率性: PReMは、競争力のある推論レイテンシを維持しながら、ピークGPUメモリを大幅に削減しました（例：フルメモリの30.0GBに対し、約19.57GB）。 $\langle m \rangle$ トークンの使用により、デコーディング計算（リフレッシュ頻度）とメモリ節約の間の制御可能なトレードオフが可能になります。

5. 意義と主張

本論文は、効果的な長文脈推論には、静的なコンテキスト短縮ではなく、生成と結合されたメモリ更新が必要であることをPReMが示していると主張しています。コンテキストを内部のKVメモリとして保持し、それを動的にリフレッシュするように学習することで、モデルは各推論ステップに必要な特定の証拠に適応的にアクセスできます。

著者らは、彼らのアプローチが静的な保持と分離された選択の限界に対処しており、モデルが「何を保持し、いつリフレッシュするかを知る」ことを学習できることを示し、マルチホップ推論タスクにおいて優れた性能を発揮し、既存の圧縮技術と比較して優れた効率と品質のトレードオフを実現できることを強調しています。この研究は、内部メモリ管理を生成と共同で訓練することが、法外な計算コストをかけずにLLMをより長いコンテキストへとスケールさせるための実行可能な道であることを示唆しています。