Eliciting In-context Retrieval and Reasoning for Long-context Large Language Models
- 長文コンテキストを扱う LLM が、RAG の検索器を置き換えられるかを調べた論文である。
- 著者らは、紛らわしい関連文書を意図的に混ぜた ICR2 ベンチマークを作り、従来の LOFT が性能を高く見積もりやすいことを示す。
- Mistral-7B に retrieve-then-generate ファインチューニングと retrieval-attention probing を組み合わせると、複数タスクで GPT-4-Turbo に近い性能を得る。
Abstract(日本語訳)
近年の長文コンテキスト言語モデル(LCLM)の進展は、パイプラインを単純化することにより、Retrieval-Augmented Generation(RAG)を変える可能性を持つ。拡張されたコンテキストウィンドウにより、LCLM は知識ベース全体を処理し、検索と推論を直接行うことができる。この能力を、本論文では In-Context Retrieval and Reasoning(ICR2)と定義する。しかし、LOFT のような既存のベンチマークは、過度に単純化されたコンテキストを与えるため、LCLM の性能をしばしば過大評価する。これに対処するため、本論文では ICR2 を導入する。ICR2 は、強力な検索器で取得した紛らわしいパッセージを含めることで、より現実的な状況における LCLM を評価するベンチマークである。さらに、LCLM の性能を高める三つの方法を提案する。第一に retrieve-then-generate ファインチューニング、第二に retrieval-attention-probing であり、これはデコード中に attention head を用いて長いコンテキストをフィルタリングし、ノイズを除く。第三に、生成ヘッドと並行して検索ヘッドを共同訓練する方法である。五つのよく知られた LCLM を LOFT と ICR2 で評価した結果、Mistral-7B に適用した最良の手法は大きな改善を示した。Exact Match で見ると、通常の RAG と教師ありファインチューニングに比べて、LOFT ではそれぞれ +17 点、+15 点、ICR2 ではそれぞれ +13 点、+2 点を得た。さらに、このモデルははるかに小さいにもかかわらず、多くのタスクで GPT-4-Turbo を上回った。
論文の面白いところ
この論文の要点は、長いコンテキストをそのまま入れれば検索問題が解ける、という素朴な期待を丁寧に疑っている点にある。RAG では通常、検索器が文書集合から候補を選び、その候補を LLM が読んで答える。長文コンテキストモデルでは、文書集合をまとめてコンテキストに置けば、検索器を省けるように見える。しかし実際の文書集合には、質問に関係はあるが答えにはつながらない文書が多い。著者らはこの「関係はあるが誤誘導する文書」を confounder と呼び、評価に明示的に組み込む。ここが、単なる長文 QA の性能比較ではない点である。モデルが長く読めることと、必要な根拠を選び損ねないことは同じではない。論文はその差を、ベンチマーク設計と手法設計の両方から扱っている。実用上も、社内文書検索やナレッジベース QA では似た名前の製品、古い仕様、別地域の規則が混ざるため、この問題は避けにくい。
問題設定
対象は、長文コンテキスト言語モデル(Long-context Language Model, LCLM)が、与えられた大きなコンテキスト内から根拠を探し、質問応答や事実検証、対話応答を行う設定である。著者らはこの能力を In-Context Retrieval and Reasoning(ICR2)と呼ぶ。従来の LOFT は、正解根拠となる文書を入れたうえで、残りをランダムな文書で埋める設計であった。この場合、不要文書は質問とあまり関係しないことが多く、モデルにとっては正解根拠が目立ちやすい。現実の検索では、BM25 や dense retriever が質問に似た文書を大量に拾うため、誤答を誘う文書のほうがむしろ自然に混ざる。たとえば「カリフォルニアの giant redwoods はどこにあるか」という質問では、Humboldt County を示す文書だけでなく、Redwood City やカリフォルニアの別施設に関する文書も候補になりうる。ICR2 はこの状況を KILT の Wikipedia ベースのタスクから構成する。正解根拠は人手注釈に従って入れ、紛らわしい文書は Contriever、DPR、DrQA、BM25、BLINK などの検索器で集める。評価は NaturalQuestions、HotpotQA、FEVER、Wizard of Wikipedia を含み、32K トークンの文脈長で行われる。
提案手法
著者らは、LCLM に文脈内検索を行わせるために三つの方法を試す。第一は retrieve-then-generate ファインチューニングである。これは、モデルにいきなり答えを出させるのではなく、まずコンテキスト中の関連パッセージを取り出し、その後で答えを生成させる訓練である。実装には、関連パッセージ本文を出力する Retrieve-Then-Answer(RTA)と、関連パッセージの ID だけを出力する Cite-Context-ID(CCI)がある。第二は Retrieval Attention Probing(RAP)である。これは再訓練を必須とせず、検索に反応しやすい attention head を検証データで選び、推論時にその head が強く注目するパッセージだけを残す方法である。長い文書全体を最後まで読ませるのではなく、モデル内部の注意の偏りを使ってコンテキストを短くする。第三は、生成用のヘッドとは別に検索ヘッドを加えて共同訓練する方法である。検索ヘッドはクエリと各パッセージの表現から関連度を計算し、上位のパッセージだけを生成側に渡す。三つの方法のうち、最も安定した結果を出したのは retrieve-then-generate ファインチューニングと RAP の組み合わせであった。
結果
まず、通常の RAG は closed-book 設定より多くの場合で良いが、Oracle RAG との間には大きな差が残る。これは、外部知識が必要である一方、長いコンテキストから正しい根拠を選ぶ部分がまだ難しいことを示す。ICR2 では、LOFT より性能が下がり、論文中では exact match が最大 51% 低下する場合があると報告されている。Mistral-7B を用いた主実験では、通常の RAG に対して SFT 系の手法が全体に改善を示す。SFT-RTA と RAP の組み合わせは、Mistral-7B の中で最良であり、LOFT と ICR2 の七タスク中五タスクで最高値を得た。要旨に示された平均では、この最良手法は通常の RAG に比べて LOFT で +17 点、ICR2 で +13 点の Exact Match 改善を示す。直接回答だけを学習する SFT と比べても、LOFT で +15 点、ICR2 で +2 点であった。RAP は追加のデコード段を持つが、選ばれた文脈が短くなるため、遅延の増加は比較的小さい。共同検索ヘッドは ICR2 で通常の RAG を上回るものの、SFT-RTA や RAP ほど安定していない。論文の制約として、主な実験が 32K トークンと Mistral-7B に寄っている点は留意すべきである。
具体例
社内文書 QA を考えると、この論文の設定は分かりやすい。入力は「東京オフィスの出張精算で、領収書の提出期限はいつか」という質問と、規程、過去のお知らせ、別拠点の経費ルール、古い改定前の文書を含む長いコンテキストである。通常の長文 LLM は、コンテキスト内に「提出期限は月末から 10 営業日以内」と「大阪オフィスでは 5 営業日以内」と「2023 年以前は翌月末まで」のような近い文書が並ぶと、質問に合う根拠を取り違えることがある。ICR2 は、このような紛らわしい候補を評価時に意図的に入れる。retrieve-then-generate では、モデルはまず東京オフィスの現行規程にあたるパッセージを抜き出し、その根拠に基づいて期限を答えるように訓練される。RAP を使う場合は、検索に効いている attention head が強く見ている文書だけを残し、不要な拠点別規程や旧版規程を減らしてから答えを生成する。期待される出力は、単に「10 営業日以内」と答えるだけでなく、その根拠が東京オフィスの現行規程であることに合った回答である。間違えやすい点は、語彙が近いだけの文書を根拠にしてしまうことである。この論文の貢献は、その誤りを評価に含め、さらにモデル側にも根拠選択の手続きを明示的に学ばせた点にある。