検索拡張生成（RAG）では、検索された長いコンテキストが入力長を増大させ、推論の遅延やコンテキストウィンドウの制限を引き起こします。文書を連続的な埋め込みに変換するソフトコンテキスト圧縮は高い表現能力を持ちますが、既存手法の多くは非圧縮モデルの出力を模倣する蒸留によって学習されるため、圧縮表現に最適化された計算パターンを学習できず、元のモデルの性能に制約されるという課題があります。本研究は、クエリに依存せずオフラインで文書をエンコードする手法を対象としています。
従来のソフト圧縮研究は、模倣に基づく学習やグローバルな最適化に主眼を置いており、模倣と探索を相補的に活用する手法は提案されていませんでした。本研究の新規性は、信頼性の高い模倣を行う段階と、教師モデルが失敗した困難なケースに対して標的を絞った探索を行う段階を分離した、2段階の学習パラダイムを導入した点にあります。これにより、圧縮表現が持つ潜在的な表現能力を最大限に引き出し、非圧縮モデルの性能を突破することを可能にしました。
提案手法DEX-Compは、文書に圧縮用トークンを付加してエンコードする圧縮器と、その埋め込みを用いて回答を生成するデコーダで構成されます。学習は以下の2段階で行われます。
第1段階のPure Distillationでは、非圧縮の教師モデルが正解を出力できた応答のみを用いて、教師と学生の出力分布の間のKLダイバージェンスを最小化します。これにより、教師の誤りを内面化することなく、適切な初期化を行います。
第2段階のHard Explorationでは、教師モデルが失敗したクエリに対してのみ、Group Relative Policy Optimizationを用いた強化学習を実施します。正解との一致度に基づく報酬を用いることで、単なる模倣を超え、圧縮された表現に対してより適した計算戦略の探索を促します。
Natural Questions、TriviaQA、HotpotQA、ASQA、PopQAの5つのオープンドメインQAベンチマークを用い、Mistral-7Bをバックボーンとして評価しました。比較対象は、非圧縮RAGおよびLLMLingua-2、xRAG、ICAE、COCOM、PISCOです。評価指標にはContainment Exact MatchとLLM-as-judgeを用い、検索深度をtop-5からtop-30まで検証しました。結果として、あらゆる検索深度において非圧縮RAGを上回る性能を達成しました。また、top-30の設定では、非圧縮RAGと比較してピークGPUメモリを44%、計算コストを33%削減し、最初のトークン生成までの時間を25%高速化しました。
本手法は、コーパスに対して事前に圧縮済み埋め込みを計算・保存するオフラインの圧縮を前提としています。そのため、頻繁に更新される超大規模なコーパスでは、再計算のコストが課題となります。また、圧縮による効果と強化学習による効果を厳密に分離して評価できていない点や、生成内容の忠実性、根拠の妥当性、敵対的な入力に対する堅牢性の検証が今後の課題です。なお、圧縮率と精度の間にはトレードオフが存在しますが、圧縮率が32であっても非圧縮RAGと同等の性能を維持できることが示されています。
RAGにおけるコンテキスト圧縮は、ドキュメントを短い埋め込みシーケンスに符号化することで推論効率を高める手法ですが、既存手法の多くは非圧縮RAGの出力を蒸留する形式をとるため、元のモデルの性能を超えることが困難でした。本研究では、この限界を打破するためにDEX-Compという2段階の学習レシピを提案しています。第1段階のPure Distillationでは、非圧縮RAGが正解を出した応答のみを用いて圧縮モデルの初期学習を行い、第2段階のHard Explorationでは、非圧縮RAGが失敗したクエリに対してのみ強化学習を適用することで、圧縮表現に適した計算パターンの探索を促します。5つのオープンドメインQAベンチマークを用いた評価では、検索深度が上位5件から上位30件の範囲において、コンテキストを圧縮し推論を高速化しつつ、非圧縮RAGと同等以上の性能を達成しました。
RAGにおけるコンテキスト圧縮は、推論の遅延やコンテキストウィンドウの制限を解決する手法として重要であり、特に連続的な埋め込みを用いるソフト圧縮は、離散的なトークンを用いるハード圧縮よりも高い表現能力を持つ。しかし、従来のソフト圧縮は非圧縮モデルの応答を蒸留する手法に依存しており、その学習プロセスが離散トークンを読み取る際の計算パターンに縛られ、圧縮された高密度な表現の能力を十分に引き出せていないという課題がある。本研究では、この問題を解決するためにDEX-Compという2段階の学習レシピを提案する。第1段階のPure Distillationでは、非圧縮モデルの正解のみを模倣することで、誤った振る舞いを内面化せずに適切な初期化を行う。第2段階のHard Explorationでは、非圧縮モデルが失敗したクエリに対してのみ強化学習を適用することで、模倣だけでは不十分なケースにおいて圧縮表現に適した新たな振る舞いを探索させ、非圧縮モデルを超える性能を目指す。5つのオープンドメインQAベンチマークを用いた実験では、入力サイズを大幅に削減し推論を高速化しつつ、検索深度がtop-5からtop-30の広範囲において、従来のソフト圧縮手法を上回り、非圧縮RAGと同等以上の性能を達成した。
検索拡張生成（RAG）は、知識集約的なタスクにおいて言語モデルの性能を向上させる手法として、リトリーバーとの共同学習や自己反省メカニズムなどの有効性を高める研究が行われてきました。一方で、長い検索コンテキストの入力に伴う効率性の低下については、既存の取り組みと比較して十分に検討されていません。これに対し、推論時のKVキャッシュ圧縮や、入力テキスト自体を短縮するプリインファレンス・コンテキスト圧縮の研究が進んでいます。コンテキスト圧縮には、トークン抽出や要約を行うハード圧縮と、コンテキストを連続的な埋め込み表現に変換するソフト圧縮があり、後者は連続空間の柔軟性により既存のベンチマークで高い性能を示す傾向があります。本研究では、クエリに依存せずオフラインで文書をエンコードするクエリ非依存のソフト圧縮に焦点を当てており、既存手法には自己符号化やアテンションに基づく情報集約、あるいはテキスト再構成や蒸留を用いた学習パラダイムが存在します。しかし、これまでの研究は模倣に基づく手法やグローバルな最適化を主眼としており、模倣と探索の相補的な強みをソフト圧縮においてどのように活用するかについては明示的に研究されていません。本論文は、信頼性の高い模倣と困難なケースに対する標的を絞った探索を分離する、2段階の学習パラダイムを導入することでこの限界に対処します。
DEX-Compは、RAG（検索拡張生成）において文書のトークン列を連続的な埋め込み表現に置き換えることで、入力長を削減し推論効率を高める2段階の学習手法である。アーキテクチャは、文書に圧縮用トークンを付加してエンコードし、圧縮された埋め込みをキャッシュする圧縮器と、それを利用して回答を生成するデコーダで構成される。第1段階のPure Distillation（PD）では、教師モデルである非圧縮RAGが正解を出力できたデータのみを抽出し、教師の出力分布と学生モデルの出力分布の間のKLダイバージェンスを最小化することで、信頼性の高い初期化を行う。第2段階のHard Exploration（HE）では、教師モデルが失敗したデータに対してのみ、Group Relative Policy Optimization（GRPO）を用いた強化学習を行い、正解との一致度に基づく報酬を用いて、模倣を超えた最適な計算戦略を探索する。この手法は、圧縮された表現が持つ高い表現能力を活用することを目的としており、模倣による初期化と、教師の限界を超えるための探索を分離することで、非圧縮モデルと同等以上の性能を実現している。
本実験では、事実に基づく回答、マルチホップ、曖昧な長文、ロングテール回答の各能力を評価するため、Natural Questions、TriviaQA、HotpotQA、ASQA、PopQAの5つのデータセットを使用し、検索コーパスには128トークン単位に分割されたWikipedia-KILTを用います。比較対象として、圧縮を行わないRAG、およびクエリに依存せずオフラインで文書をエンコードする手法であるLLMLingua-2、xRAG、ICAE、COCOM、PISCOを設定し、すべての手法でMistral-7Bをデコーダーのバックボーンとして共通利用します。評価指標には、正解が予測に含まれるかを確認するContainment Exact Matchと、Gemini 1.5 Flashを判定器として用いるLLM-as-judgeの2種類を採用し、検索深度をtop-5からtop-30まで広げて検証します。実装面では、Splade-v3をリトリーバー、DeBERTa-v3をリランカーとして用い、提案手法であるDEX-Compは圧縮器とデコーダーに同じバックボーンを用い、それぞれを個別のLoRAアダプタとして実装することで学習コストを抑えています。
DEX-Compは、Mistral-7Bを用いた評価において、あらゆるデータセットおよび検索深度（top-5からtop-30）で非圧縮のRAGベースラインを上回り、平均してCEM/LLMスコアで2ポイント以上の統計的に有意な向上を達成しました。特定の検索深度で学習した単一のモデルを他の深度に適用した場合でも、PopQAを除く全てのデータセットで非圧縮RAGを上回る高いクロス深度汎化性能を示しています。計算効率の面では、検索深度が増すほど利点が顕著になり、top-30の設定ではピークGPUメモリを44%、計算コストを33%削減し、最初のトークン生成までの時間（TTFT）を25%高速化しました。これらの効率化は、単一のRTX PRO 6000 Blackwell GPUを用いた実験で確認されており、高負荷な検索設定下での大幅な改善を実証しています。なお、DEX-Compはコーパスに対して一度だけオフラインで圧縮処理を行う必要があり、動的なコーパスに対しても更新された文書のみを増分的に再圧縮することで対応可能です。
DEX-Compの性能を評価するため、ノイズへの耐性を示すResilience Rateと、検索による回答改善の有効性を示すBoost Rateを用いて分析した結果、訓練時と推論時の検索深度が一致する場合において、非圧縮のRAGベースラインを両指標で上回ることが示された。特にResilience Rateの向上は、圧縮が暗黙的なデノイジングとして機能していることを示唆している。2段階の学習レシピに関する検証では、教師モデルの正解データを用いた蒸留（Stage I）と、困難な事例を用いた強化学習（Stage II）の両方が不可欠であり、これらが相補的に作用して性能向上に寄与していることが明らかになった。圧縮率（CR）については、CR = 16において非圧縮RAGを上回る最高性能を達成し、CR = 32でも同等の性能を維持できることから、精度と効率の柔軟なトレードオフが可能である。また、訓練時の検索深度を大きく設定することで、異なる検索深度への汎化性能が高まることが確認された。さらに、学習された圧縮埋め込みは、特定の文書領域に焦点を当てるものや、文書全体の情報を捉えるものなど、空間的な専門化を示しており、限られた容量を局所的および大域的な表現に割り当てていることが解析により示された。
DEX-Compは、クエリに依存しないオフラインの圧縮設定を前提としており、コーパスに対して圧縮済みの埋め込みを事前に計算して保存する必要があります。このコストはクエリごとに分散されますが、頻繁に更新されるコーパスや超大規模なコーパスでは、計算コストが大きくなる可能性があります。また、圧縮の効果とタスク固有の強化学習による効果を分離することは困難であるため、現在の実験設定ではこれらを厳密に切り離して評価できていません。評価については、質疑応答ベンチマークと回答の正確性に焦点を当てており、生成内容の忠実性や根拠の妥当性、あるいは敵対的な入力や長文生成の設定における堅牢性については検証の余地が残されています。
本研究では、純粋な蒸留とハードな探索を組み合わせた、ソフトコンテキスト圧縮のための2段階学習レシピであるDEX-Compを提案しています。DEX-Compは、クエリに依存しないソフト圧縮手法として、複数の検索設定において非圧縮のRAGモデルの性能を上回ることを示した手法です。具体的には、上位30件の検索結果を用いた設定において、検索されたコンテキストを圧縮し、推論速度を10倍以上に加速させつつ、5つのベンチマークにおいて非圧縮のRAGと同等以上の性能を達成しています。詳細な分析により、学習の各段階が性能に寄与していることや、異なるバックボーンモデルやデータ分布に対しても汎用性を持つことが確認されました。本研究は、圧縮されたRAGシステムと非圧縮のシステムの性能差を縮小するための重要な一歩となります。