Lin, Xiaoqiang, Ghosh, Aritra, Low, Bryan Kian Hsiang, Shrivastava, Anshumali, & Mohan, Vijai.（2025）. REFRAG: Rethinking RAG based Decoding. arXiv preprint arXiv:2509.01092.

検索拡張生成[RAG]などのアプリケーションにおいて, 大規模言語モデル[LLMs]が長いコンテキスト入力を処理する際に生じるシステムの遅延と多大なメモリ消費という根深い課題に対処するための, 効率的なデコーディングフレームワーク「REFRAG」を提案するものである.従来のLLMでは, プロンプトの長さが増加すると, キー・バリュー[KV]キャッシュのメモリ消費がプロンプト長に比例して増加し, 最初のトークンを生成するまでの時間[TTFT]がプロンプト長の二乗に比例して増大するため, 推論のスループットが低下するという問題があった.この研究は, このような課題を克服し, 知識の豊富化とシステム効率の間に存在する基本的なトレードオフを改善することを目指している.

論文の核心的な着想は, RAGシステムにおいて, LLMのコンテキストの大部分が検索によって連結されたパッセージで構成されており, その中でクエリに直接関連する部分はごく一部であるという観察に基づいている.これらのパッセージは, 多様性や重複排除といった再ランキングプロセスによって, 意味的な類似性が低い場合が多く, 標準的なLLM生成タスクとは異なるブロック対角のアテンションパターンを生成する.この観察に基づき, デコーディング中のRAGコンテキストにおける計算のほとんどは不要であり, 性能への影響を最小限に抑えつつ排除できると主張されている.RAGのTTFTを一般的なLLM推論問題として扱うことは, 非効率なトークン割り当て, エンコーディング情報の無駄な破棄, そしてRAGコンテキストに特有の構造化された疎なアテンションパターンを見落とすことにつながるため, 適切ではないという見解である.

REFRAG[REpresentation For RAG]は, RAGアプリケーションにおける遅延を改善するために, 圧縮, 感知, 拡張を行う効率的なデコーディングフレームワークとして提案されている.このフレームワークは, LLMアーキテクチャの変更や新しいデコーダパラメータの導入を必要とせず, デコーディング中の遅延[TTFT, Time-to-Iterative-Token: TTIT]およびメモリ使用量を大幅に削減するものである.

REFRAGの主な利点は三つある.第一に, 検索されたパッセージのトークンを直接入力として使用する代わりに, 事前に計算された圧縮済みチャンク埋め込みを近似表現として活用し, これを直接デコーダに供給する.これにより, デコーダの入力長が短縮され, トークン割り当て効率が向上する.第二に, 検索時に計算されたチャンク埋め込みを再利用できるため, 冗長な計算が不要になる.第三に, アテンション計算の複雑性が, コンテキスト内のトークン数ではなくチャンク数に対して二乗でスケールするようになり, 計算コストが大幅に削減される.REFRAGは, 以前の方法とは異なり, デコーダの自己回帰性を維持しながら, トークンチャンクの圧縮を任意の位置でサポートしており, マルチターンやエージェントアプリケーションにも対応可能である.この「どこでも圧縮」機能は, 軽量な強化学習[RL]ポリシーによってさらに強化され, フルチャンクトークン入力が必要な場合と低コストの近似チャンク埋め込みで十分な場合とを自動的に選択する.これにより, 計算量の多いトークン埋め込みへの依存を最小限に抑え, RAG設定におけるクエリのほとんどのチャンクを凝縮する.

REFRAGのモデルアーキテクチャは, デコーダのみの基盤モデル[例えばLLaMA]と軽量なエンコーダモデル[例えばRoBERTa]で構成されている.質問[メイン入力トークン]とコンテキスト[検索されたパッセージ]が与えられると, コンテキストは固定サイズのチャンクに分割される.エンコーダモデルがこれらのチャンクを処理してチャンク埋め込みを生成し, この埋め込みは投影層によってデコーダモデルのトークン埋め込みのサイズに合わせられる.これらの投影されたチャンク埋め込みは, 質問のトークン埋め込みと共にデコーダモデルに供給され, 応答が生成される.RAGアプリケーションでは, コンテキストが入力の大部分を占めるため, この設計により入力シーケンスが大幅に短縮され, 遅延とメモリ使用量が削減される.

REFRAGのトレーニングは, いくつかのステップで行われる.まず, エンコーダとデコーダをアライメントするために, 次パラグラフ予測タスクを用いた継続事前学習[CPT]選択的トークン圧縮を導入しており, 重要なコンテキストチャンクは非圧縮のまま拡張することで, 応答予測を改善する.このRLポリシーは, 次パラグラフ予測のパープレキシティを負の報酬としてガイドされ, どのチャンクを元の形で保持するかを決定する.ポリシーネットワークはチャンク埋め込みとマスキングを活用して順次チャンク拡張を最適化し, デコーダの自己回帰特性を維持しつつ, 圧縮の柔軟な配置を可能にする.

実験結果では, REFRAGがRAG, マルチターン会話, 長文要約など, 多様な長文コンテキストタスクで厳密に検証された.REFRAGは, パープレキシティの損失なしに, TTFTを最大30.85倍加速し[以前の手法と比較して3.75倍の改善], また, 圧縮によりLLMのコンテキストサイズを16倍に拡張できることが示されている.理論分析によれば, 短いコンテキスト長ではTTFTとスループットで最大k倍の加速を, 長いコンテキスト長では両方のメトリックで最大k²倍の加速を達成する.例えば, k=16の圧縮率で, REFRAGは16.53倍のTTFT加速を達成し, CEPE[既存の最先端手法]よりも優れている.LLaMAモデルや他の最先端のベースラインと比較しても, 精度を損なうことなく大幅な高速化を実現し, 拡張されたコンテキストウィンドウにより, REFRAGは人気のアプリケーションで精度をさらに向上させることが示されている.アブレーションスタディでは, カリキュラム学習が再構成タスクの成功に不可欠であり, 再構成タスクの事前学習も継続事前学習タスクの成功に重要であることが示されている.また, RLベースの選択的圧縮ポリシーは, 低い圧縮率を全面的に適用するよりも一貫して優れた性能を示しており, 性能を損なうことなく圧縮率を動的に調整する実用性が強調されている.マルチターン会話では, LLaMAFTの限られた4kトークンのコンテキストウィンドウが会話履歴の切り捨てによって重要な情報損失をもたらすのに対し, REFRAGの圧縮アプローチは多数のパッセージでも堅牢な性能を維持し, LLaMAFTを上回る.長文要約タスクでは, 同じデコーダトークン数[すなわち同じ遅延]においてREFRAGが最高の性能を達成し, 高い圧縮率で文書からより多くの情報を組み込むことで, 性能が向上することが示されている.

REFRAGは, アテンションメカニズムの複雑性変更やコンテキストの疎化など, LLMの推論遅延最適化に関する既存の研究とは異なり, RAGコンテキストに固有の構造と疎性を活用する特化された手法である点が特徴である.CEPEのような先行研究は, クロスアテンションを用いてKVキャッシュメモリとアテンション計算を削減するが, 因果関係を損なうためマルチターンRAGや要約タスクには不向きであり, トークン圧縮も利用していない.また, 圧縮型トランスフォーマーやプロンプト圧縮の先行研究と比較しても, REFRAGはチャンク埋め込みの事前計算とプロンプト内の任意の位置での使用を可能にする点で初めてであり, 推論時にチャンク埋め込みを再計算することなく適応的な圧縮率を学習できる点が異なる.

結論として, REFRAGはRAGアプリケーションに特化した新しい効率的なデコーディングフレームワークである.RAGコンテキストに固有の疎性とブロック対角のアテンションパターンを活用することで, コンテキスト表現を圧縮, 感知, 拡張し, メモリ使用量と推論遅延[特にTTFT]を大幅に削減する.幅広い長文コンテキストアプリケーションでの広範な実験により, REFRAGがパープレキシティやダウンストリームの精度を損なうことなく, TTFTを最大30.85倍加速すること[既存の最先端手法と比較して3.75倍]を実証した.REFRAGは, 遅延に敏感な知識集約型アプリケーションでLLMを展開するための, 実用的でスケーラブルなソリューションを提供すると考えられている.