LLaMA モデルによるテキスト検索パイプラインの最適化

リンク一覧

概要と序論

方法

実験

アブレーション研究と分析

関連作業

結論、謝辞および参考文献

2 方法

2.1 準備

タスク定義 クエリ Q と、n 個のドキュメントからなるコーパス C = {D1, D2, …, Dn} が与えられた場合、テキスト検索の目標は、クエリ Q に最も関連性の高い k 個のドキュメント (k ≪ n) を見つけることです。リトリーバーとリランカーで構成された多段階検索パイプラインでは、リトリーバーのタスクは、類似度メトリック Sim(Q, D) ∈ R に基づいて、クエリに関連する上位 k 個の候補を効率的に生成することです。リランカーのタスクは、より効果的であるが通常はより計算コストの高いリランキング モデルを使用して、これらの k 個の候補ドキュメントをさらに並べ替えて、関連性の順序を改善することです。このコンテキストでの「ドキュメント」は、文、節、または完全なドキュメントを含む任意の情報スニペットを指すことに注意してください。多段階パイプラインには複数のリランカーを含めることができますが、このホワイト ペーパーでは 1 つのリランカーに焦点を当てます。

現代のリトリーバーは通常、テキストをベクトル表現にエンコードするバイエンコーダーアーキテクチャに従っており、Sim(Q, D) はクエリ Q とドキュメント D のベクトル表現のドット積として計算されます (Karpukhin ら、2020)。対照的に、(ポイントワイズ) リランカーは通常、クエリと候補ドキュメントの両方を入力として受け取り、関連性スコアを直接生成します。これらのスコアは、候補を並べ替えるために使用されます (Nogueira ら、2019; Gao ら、2021)。

ラマ LLaMA (Touvron et al., 2023a) は、Transformer アーキテクチャに基づく、自己回帰型のデコーダーのみの大規模言語モデルです。このモデルは、膨大な量の Web データで事前トレーニングされた数十億のパラメーターが特徴です。単方向であるということは、モデルの注意メカニズムが予測を行う際に入力シーケンス内の先行する要素のみを考慮することを意味します。具体的には、入力シーケンス x = (t1, t2, …, tn−1) が与えられた場合、モデルは先行するトークンのみに基づいて次のトークン tn の確率を計算します。予測プロセスは数学的に P(tn|t1, …, tn−1) と表すことができます。ここで、P は確率を表し、tn はシーケンス内の次の要素を表します。

2.2 レトリーバー

RepLLaMA と呼ばれる私たちのリトリーバー モデルは、DPR (Karpukhin et al.、2020) で提案されたバイエンコーダー デンス リトリーバー アーキテクチャに従いますが、バックボーン モデルは LLaMA で初期化されます。

これまでの高密度リトリーバー モデルの研究は、BERT のような双方向エンコーダーのみのモデルを使用することが多く、先頭に付加された (CLS) トークンの表現をテキスト入力の高密度表現として採用しています。ただし、LLaMA は単方向であるため、入力クエリまたはドキュメントにシーケンス終了トークンを追加して、LLaMA への入力シーケンスを形成します。したがって、クエリまたはドキュメントのベクトル埋め込みは次のように計算されます。

ここで、Decoder(·) は LLaMA モデルを表し、各入力トークンの最後のレイヤー トークン表現を返します。シーケンス終了トークンの表現を入力シーケンス t1 . . . tk の表現として取ります。これはクエリ Q またはドキュメント D のいずれかです。D と Q の関連性は、対応する密な表現 VQ と VD のドット積として Sim(Q, D) = として計算されます。

次に、モデルは InfoNCE 損失に従ってエンドツーエンドで最適化されます。

ここで、D+ はクエリ Q に関連するドキュメント (人間のラベルに基づく) を表し、{DN } はクエリに関連しないドキュメント セットを表します。ネガティブ ドキュメント セットには、既存の検索システムのトップ ランクの結果からサンプリングされたハード ネガティブと、同じトレーニング バッチ内の他のクエリに関連付けられたポジティブ ドキュメントとハード ネガティブ ドキュメントから派生したバッチ内ネガティブの両方が含まれます。実際には、高密度検索トレーニングでは、ハード ネガティブとバッチ内ネガティブのより大きなセットが効果的です。

推論フェーズでは、クエリは通常リアルタイムでエンコードされ、HNSW (Malkov and Yashunin、2020) などの効率的な近似最近傍検索ライブラリを使用して、事前にエンコードされたコーパス内で上位 k 件の類似ドキュメントが検索されます。ただし、この研究では、モデルの有効性を評価するために、フラット インデックスを使用して正確な最近傍検索を実行することを選択しました。

2.3 リランカー

RankLLaMA と呼ばれる当社の再ランク付けモデルは、ポイントワイズ再ランク付けモデルとしてトレーニングされています。このアプローチでは、クエリと候補ドキュメントを一緒にモデル入力として渡し、モデルはクエリに対するドキュメントの関連性を示すスコアを生成します (Nogueira 他、2019)。

より詳細には、RankLLaMA はクエリとドキュメントのペアを次のように再ランク付けします。

ここで、Linear(·) は、シーケンス終了トークンの最後のレイヤー表現をスカラーに投影する線形投影レイヤーです。リトリーバーと同様に、モデルは対照損失によって最適化されます。ただし、この場合、ネガティブ ドキュメントにはバッチ内ネガティブは含まれません。

マルチステージ パイプラインで特定のリトリーバーからの候補を再ランク付けするように最適化されたリランカーをトレーニングするには、そのリトリーバーからのトップランクの結果からハードネガティブをサンプリングする必要があります。具体的には、この場合、RankLLaMA のハードネガティブ トレーニング データは、RepLLaMA のトップランクの結果から選択されます。

推論段階では、RepLLaMA によって取得された上位候補ドキュメントが並べ替えられます。この並べ替えは、RankLLaMA が各クエリとドキュメントのペアに割り当てる関連性スコアに基づいて行われ、ドキュメントは関連性の高い順に並べられます。