どんなもの？

近年、大規模言語モデル (LLM) の台頭により、自然言語処理 (NLP) 分野は目覚ましい発展を遂げています。特に、外部知識源を活用してLLMの能力を強化する手法として、検索拡張生成 (RAG: Retrieval-Augmented Generation) が注目されています。しかし、RAGにはリアルタイム検索による遅延や、適切な文書選択の難しさ、システムの複雑化といった課題が存在します。
本論文では、これらの課題を解決するために、キャッシュ拡張生成 (CAG: Cache-Augmented Generation) という新しいパラダイムを提案しています。

先行研究と比べてどこがすごい？

従来のRAGでは、質問が与えられるたびに外部知識源から関連情報を検索する必要がありました。CAGでは、予め関連するすべての文書をLLMの拡張コンテキストにロードし、その実行時パラメータをキャッシュしておくことで、リアルタイム検索を不要にしています。これにより、検索の遅延をなくし、検索エラーを最小限に抑えながら、コンテキストの関連性を維持することができます。

技術や手法のキモはどこ？

CAGの重要なポイントは、LLMの拡張コンテキストとキーバリュー (KV) キャッシュの活用です。

拡張コンテキスト:
- 近年のLLMは、従来のモデルよりもはるかに長いテキストシーケンスを処理できる拡張コンテキストを備えています。CAGでは、この拡張コンテキストに、タスクに必要なすべての知識を事前にロードします。
KVキャッシュ:
- LLMは、入力テキストを処理する際に、キーバリュー (KV) キャッシュと呼ばれる内部表現を生成します。このキャッシュには、入力テキストに関する情報が圧縮された形式で保存されます。CAGでは、事前にロードした知識に対するKVキャッシュを計算し、保存しておきます。
- 推論時には、このキャッシュを利用することで、リアルタイム検索を行うことなく、事前にロードされた知識に基づいて回答を生成できます。

追加調査：KVキャッシュはテキスト？　ベクトル？

KVキャッシュの具体的な計算方法については、論文では言及されていません。しかし、KVキャッシュはLLM内部で生成されるものであり、一般的にはテキストを処理して作成されたベクトルであると考えられます。
ユーザーの質問と一緒にLLMに渡されるKVキャッシュは、直接テキストとして渡されるのではなく、処理されたベクトル表現として渡されます。
詳しく説明すると、

推論前に、関連文書群をLLMに入力として与え、KVキャッシュを計算します。この際、LLMは入力テキストをベクトル表現に変換し、それをKVキャッシュとして保存します。
推論時には、質問も同様にベクトル表現に変換されます。
そして、質問のベクトル表現と、事前に計算されたKVキャッシュが組み合わされ、LLMに入力として渡されます。

このように、LLM内部では、テキストはベクトル表現として処理され、KVキャッシュもベクトルとして保存・利用されます。
ただし、具体的なKVキャッシュの計算方法やベクトル表現の種類は、使用するLLMのアーキテクチャや実装に依存するため、ソース資料だけでは断定できません。

追加調査2：拡張コンテキストとKVキャッシュの具体的な処理は？

例として、論文中のSQuAD実験におけるCAGの処理の流れを以下に示します。

外部知識の事前ロード: SQuADデータセットのすべての文書 (DS) をLlama 3.1モデルの拡張コンテキストにロードします。
KVキャッシュの事前計算: ロードした文書群 (DS) に対して、Llama 3.1モデルのKVキャッシュ (CS_KV) を計算し、保存します。
推論: ユーザーからの質問 (@8) が与えられると、事前計算されたKVキャッシュ (CS_KV) と質問 (@8) を組み合わせたものをLlama 3.1モデルに入力として与え、回答を生成します。

このように、CAGは、事前にKVキャッシュを計算しておくことで、推論時に必要な計算量を大幅に削減し、高速な回答生成を実現しています。

どうやって有効だと検証した？

本論文では、質問応答ベンチマークとして広く知られるSQuAD と HotPotQA を用いて、CAGの有効性を検証しています。これらのデータセットは、それぞれ単一文書内での正確な回答と、複数文書を横断する複雑な推論を必要とするため、包括的な評価を行うことができます。
実験では、参照テキストの長さを変えた複数のテストセットを作成し、CAGと従来のRAGシステム (BM25を用いた疎検索システムとOpenAI Indexesを用いた密検索システム) の性能を比較しました。評価指標には、BERTScore を使用しています。
その結果、CAGはほとんどの場合において最高のBERTScoreを達成し、両方のRAGシステムを凌駕しました。これは、CAGが検索エラーを排除し、関連するすべての情報に基づいて全体的な推論を行うことができるためです。また、CAGは従来のRAGシステムと比較して、生成時間が大幅に短縮されることも示されました。

ただし、論文で取り上げられているのは、比較的小規模な知識ベースを用いた実験であることに注意が必要です。より大規模な知識ベースを用いた場合や、リアルタイムでの情報更新が重要なタスクにおいては、RAGがCAGよりも優れている可能性も考えられます。

議論はある？

本論文では、RAGを完全に置き換えることを提案していますが、すべてのユースケースでCAGが最適解であるとは限りません。
例えば、非常に大規模な知識ベースや、動的に変化する情報を含む知識ベースに対しては、CAGは適さない可能性があります。そのような場合は、事前にロードされたコンテキストと選択的な検索を組み合わせたハイブリッドアプローチが有効と考えられます。

メリット・デメリットまとめ

CAG (Cache-Augmented Generation)
- メリット
  - 推論速度が速い: 事前にKVキャッシュを計算することで、推論時にリアルタイムで検索を行う必要がなくなり、高速な回答生成が可能となります。
  - 検索エラーが発生しない: 関連文書全体が事前に拡張コンテキストにロードされているため、検索エラーが発生するリスクがありません。
  - システム構成がシンプル: 検索コンポーネントが不要となるため、システム構成がシンプルになり、メンテナンスや開発のオーバーヘッドが削減されます
- デメリット
  - 大規模な知識ベースには適用困難: 拡張コンテキストの容量制限により、扱える知識ベースの規模が制限されます。
  - 知識ベースの更新が困難: 知識ベースに変更が生じた場合、KVキャッシュを再計算する必要があり、更新に手間がかかります。
RAG (Retrieval-Augmented Generation)
- メリット
  - 大規模な知識ベースに対応可能: 外部のデータベースや検索エンジンを利用するため、大規模な知識ベースにも対応できます。
  - 知識ベースの更新が容易: 知識ベースが更新された場合でも、検索インデックスを更新するだけで対応できます。
- デメリット
  - 推論速度が遅い: リアルタイムで検索を行う必要があるため、CAGと比較して推論速度が遅くなります。
  - 検索エラーが発生する可能性: 検索エンジンの精度やクエリとの適合性によって、検索エラーが発生する可能性があります。
  - システム構成が複雑: 検索コンポーネントと生成コンポーネントを統合する必要があるため、システム構成が複雑になります

次に読むべき論文は？

この論文で紹介されているCAGは、比較的小規模な知識ベースを用いたタスクにおいて、RAGよりも優れた性能を発揮することが示されています。しかし、より大規模な知識ベースやリアルタイムな情報更新が必要なタスクにおいて、CAGがどのように機能するかはまだ明らかではありません。
今後の研究において、これらの課題に取り組むためには、大規模な知識ベースに対応できるRAGの利点と、高速な推論と検索エラーの回避というCAGの利点を組み合わせた、ハイブリッドな手法を検討する必要があるでしょう。
具体的には、以下のような論文を読むことをお勧めします。

Retrieval Augmented Generation or Long-Context LLMs? A Comprehensive Study and Hybrid Approach (Zhuowan Li, Cheng Li, Mingyang Zhang, Qiaozhu Mei, and Michael Bendersky, 2024)
- この論文では、RAGと長文コンテキストLLMの包括的な比較を行い、ハイブリッドアプローチを提案しています。
Long Context RAG Performance of Large Language Models (Quinn Leng, Jacob Portes, Sam Havens, Matei Zaharia, and Michael Carbin, 2024)
- この論文では、大規模言語モデルの長文コンテキストRAGの性能について調査しています。

所感

目的次第ではありそうだけど、ちょっとしたFAQやキャラクターAIを作るならCAGで良さそう。
AIと会話するサービスだって、裏側でRAGを使うことによってレスポンスが遅くなるなら、CAGを試す価値はあるなと思いました。
ソースコードもあるので、機会があれば開発に進める。上司にプッシュしよう。

参考文献

"Retrieval Augmented Generation or Long-Context LLMs? A Comprehensive Study and Hybrid Approach" (Zhuowan Li, Cheng Li, Mingyang Zhang, Qiaozhu Mei, and Michael Bendersky, 2024)
"Don’t Do RAG: When Cache-Augmented Generation is All You Need for Knowledge Tasks" (Brian J Chan, Chao-Ting Chen, Jui-Hung Cheng, and Hen-Hsen Huang, 2024)
https://github.com/hhhuang/CAG