PISCO: Pretty Simple Compression for Retrieval-Augmented Generation

生成日:

PISCO: Pretty Simple Compression for Retrieval-Augmented Generation

論文の面白いところ

この論文の中心は、RAG の文書圧縮を、文書の復元問題ではなく、答えを変えないための表現学習として扱う点にある。従来のソフト圧縮では、文書を短いベクトル列に押し込むために自己符号化や言語モデル事前学習を行い、その後に QA で調整することが多かった。PISCO はその段階を大きく削り、文書と質問をそのまま読んだ教師モデルの出力を、圧縮表現だけを見た学生モデルが再現するように学習する。発想は素朴で、圧縮しても質問への答えが同じならよい、というものだが、実験ではこの素朴さがよく働いている。著者らは 10B トークン規模の事前学習も試しているが、QA 性能への寄与は小さいと述べる。これは、圧縮表現が原文をきれいに再構成できても、質問に必要な情報を使えるとは限らないことを示している。さらに、教師モデルは必ずしも巨大なものがよいわけではなく、文脈に基づく説明を含む出力を返すモデルが有利だったという観察も興味深い。実用上は、文書集合をあらかじめ圧縮しておける場合に、推論時の遅延と費用を下げる手段として読める。

問題設定

RAG は、質問に関係する文書を検索し、その文書を大規模言語モデル(Large Language Model; LLM)の入力に加えて答えを生成する方法である。検索文書を増やせば必要な根拠に当たる可能性は高くなるが、Transformer 系モデルでは入力長に応じて計算量が重くなる。コンテキスト長の上限もあり、多数の文書をそのまま入れることは難しい。文書圧縮はこの制約を緩める方法で、不要な文を削るハード圧縮と、文書を埋め込みや key-value cache に変換するソフト圧縮に大別できる。ハード圧縮は人が読めるが、圧縮率には限界がある。ソフト圧縮は高い圧縮率を狙えるが、既存手法では非圧縮入力に比べて精度低下が大きく、事前学習や注釈付きデータを必要とすることが多い。この論文は、RAG の QA 性能をできるだけ落とさず、文書を高倍率で圧縮することを目標にする。とくに、一般的な文書質問応答、未知ドメイン、多言語データ、文書数の増加に対して、圧縮表現がどこまで使えるかを調べている。

提案手法

PISCO は、各文書に少数のメモリトークンを付けて LLM に入力し、そのメモリトークン位置の隠れ状態を文書の圧縮表現として取り出す。16 倍圧縮の設定では、128 トークン程度の文書を 8 個の埋め込みベクトルに写す。圧縮側と生成側には別々の LoRA アダプタを入れ、メモリトークンの埋め込みも同時に最適化する。生成側のデコーダを固定せずに微調整する点も重要で、付録の実験では、デコーダを凍結すると多くのデータセットで性能が落ちる。学習では、まず検索された複数文書と質問を教師 LLM に与え、答えを生成させる。次に学生モデルは、原文ではなく圧縮された文書表現と質問を受け取り、教師の答えを交差エントロピー損失で再現する。これは系列レベル知識蒸留(Sequence-level Knowledge Distillation; SKD)であり、正解ラベルに対するトークンレベル蒸留とは異なる。著者らは、教師が作ったラベルを事前に保存しておけば、複数の学習実行で再利用できるとも述べている。

結果

主実験では、Natural Questions、TriviaQA、HotpotQA、ASQA、PopQA を用い、5 件の検索文書を与える設定で評価した。Mistral-7B、Llama-3.1-8B、SOLAR-10.7B を基盤にした PISCO を試し、16 倍圧縮で既存の圧縮手法を上回った。Mistral 版 PISCO の平均精度は 0.68 で、同条件の COCOM の 0.61 を上回り、非圧縮 Mistral の 0.71 に近い。Solar 版では平均 0.74 となり、非圧縮 Solar の 0.73 をわずかに上回った。著者らは、圧縮が不要情報を落とすことで軽いノイズ除去として働いた可能性を述べている。Mistral 版の 16 倍圧縮では、非圧縮 Mistral に比べて推論時の浮動小数点演算量が 17% になり、5.7 倍の速度向上が示された。GPT-4o によるペア比較でも、PISCO-Mistral は COCOM よりよい回答を多く出し、非圧縮 Mistral とはおおむね同程度と判定された。未知ドメインや多言語の評価では一部低下があるものの、基盤モデルに近い挙動を保ち、文書数を 50 件まで増やした評価でも性能低下は緩やかだった。制約として、文書をその場で圧縮する用途には向かず、文書集合を事前に圧縮して多数の問い合わせに使い回す場合に適している。

具体例

たとえば、利用者が「ペルーとチリのワイン産地で作られる無色の蒸留酒は何か」と尋ねたとする。通常の RAG では、検索器がピスコに関する文書を数件取り出し、それらの本文をそのまま LLM のプロンプトに入れる。PISCO では、各文書はあらかじめ短い埋め込み列に変換されており、推論時には本文の長いトークン列ではなく、その圧縮表現が質問とともにデコーダへ渡される。モデルは圧縮表現の中から「ペルー」「チリ」「ワイン産地」「蒸留したブドウ果汁」といった手掛かりに相当する情報を使い、「それは Pisco である」と短く答えることが期待される。このとき、圧縮表現が単に文書を丸暗記している必要はない。質問に関係する事実を、生成時に取り出せる形で保持していればよい。間違えやすい点は、文書中にブランデー、ワイン、オルホなど近い酒名が出る場合である。圧縮が粗すぎると、産地や製法の対応が混ざり、一般的な蒸留酒名や別の酒名を答えるおそれがある。PISCO の学習は、教師モデルが文書を読んで出した答えを再現させるため、このような質問応答に必要な区別を圧縮表現へ残す方向に働く。