技術要約: JKO-RAG

問題提起

Retrieval-Augmented Generation (RAG) パイプラインは、通常、リランキング（再順位付け）ステップを独立したバイナリの関連性判定の集合として扱い、候補となるパッセージのランク付きリストを生成する。著者らは、これが根本的なミスマッチであると主張している。なぜなら、ダウンストリームの消費者（言語モデル）は、リストではなく集合としてのパッセージを条件として扱うためであり、選択問題の本質は幾何学的かつ分布的であるからである。MMR（Maximal Marginal Relevance）やDPP（Determinantal Point Processes）のような既存の多様性重視の手法はこれらの問題に対処しているが、反復的な洗練や収束の保証を欠いている。さらに、標準的なリランキングは埋め込み空間のセマンティックな幾何学を無視しており、個々のスコアは高いものの、冗長なコンテキストを提供する近接した重複パッセージを回収してしまう可能性がある。

手法: JKO-RAG

本論文は、JKO-RAGを提案し、リトリーバルを、$M$ 個の候補ドキュメント上の確率単体 $\Delta^{M-1}$ 上の**分布的勾配流（distributional gradient flow）**として再定義する。

フレームワーク

本手法は、自由エネルギー汎関数 $F(p)$ を最小化するために、Jordan–Kinderlehrer–Otto (JKO) 近似スキームを反復する：
$$p_{t+1} = \arg \min_{p \in \Delta^{M-1}} \left( \frac{1}{2h} W_{C, \varepsilon}^2(p, p_t) + F_q(p) \right)$$

1. 自由エネルギー汎関数 ($F_q$): 以下の3つの項で構成される：

   • データ忠実度 (Data Fidelity): クロスエンコーダと高密度リトリーバのスコア ($E_i$) の加重和。
   • エントロピー (Entropy): 確率の崩壊を防ぐための正則化項 ($\lambda \sum p_i \log p_i$)。
   • 冗長性 (Redundancy): セマンティックに類似した（近接重複した）ドキュメントの選択にペナルティを与える二次形式の罰則項 ($\frac{\rho}{2} p^\top K p$)。ここで $K_{ij} = \max(0, \cos \langle z_i, z_j \rangle)$ である。

2. 近接項 (Proximal Term - Wasserstein-2): 距離指標は、エントロピー正則化された最適輸送 (OT) コスト $W_{C, \varepsilon}^2$ である。

   • 基底計量 (Ground Metric): $C_{ij} = (1 - \cos \langle z_i, z_j \rangle)^2$。これは埋め込み多様体上の球面距離の二乗を符号化している。
   • ソルバー: 近接ステップは、ロジット上での Adam 最適化を通じて解かれ、効率的なバックプロパゲーションのためのエンベロープ定理のトリックを用いた Sinkhorn ソルバーを利用する。

拡張機能

本論文では、コアフレームワークに対する4つの拡張を紹介している：

• NM-JKO: InfoNCE を通じて低ランクの基底計量を学習し、OT コストをタスク固有のリトリーバル幾何学により良く適合させる。
• BW-JKO: KL ダイバージェンス（ミラー降下法）と Wasserstein-2 の間を、Bregman–Wasserstein パラメータ $\alpha$ を用いて補間する。
• SAM-JKO: 候補を関連性によってクラスター化してから細粒度の JKO を実行するスコア認識型マルチレゾリューション・アプローチであり、2倍の高速化を実現する。
• DUAL-RANK: Sinkhorn 双対ポテンシャルを、選択的カバレッジ（棄却）のための較正された信頼信号として使用する。

主要な貢献と理論

本論文の中心的な貢献は、なぜ Wasserstein 幾何学が堅牢性を向上させるのかを説明する**線形応答理論（linear-response theory）**であり、経験的な観察を超えてメカニズム的な証明を行っている点にある。

理論的知見

1. 線形応答分解 (定理 1): 著者らは、エネルギー摂動（例：クエリのパラフレーズ）に対する単一の JKO ステップの感度を分析している。彼らは、応答演算子 $\delta p = -R_D \delta E$ がヘッセ行列 $A_D$ の逆行列に依存することを示す。
2. 幾何学的な安定性の源泉: Wasserstein ($W^2$) と KL 近似法の違いは、純粋に近接ヘッセ行列に起因する：
   • KL: 対角的で幾何学に無関心な演算子 ($\text{diag}(1/p)$) を使用する。
   • Wasserstein: 密な、幾何学を考慮したエントロピック-OT 演算子 ($H_W$) を使用する。
3. 幾何学的収縮 (命題 1): 密なヘッセ行列 $H_W$ は、セマンティックな計量にわたる質量の輸送にペナルティを与える「剛性（stiffness）」項として機能する。その結果、Wasserstein 応答は、摂動下での分布の幾何学的移動を収縮させるが、KL 応答は質量がクラスター間を自由に移動することを許容してしまう。
4. ステップサイズへの依存性 (系 1): 安定性の優位性は、ステップサイズ $h$ とともに単調に減少する。$h \to \infty$ (弱い近接) のとき、演算子は収束し、チューニングされた構成（大きな $h$）が同様の nDCG スコアを得る理由を説明する。$h \to 0$ (強い近接) のとき、安定性のギャップは広がる。

実験結果

著者らは、5 つの BEIR データセット (SciFact, NFCorpus, TREC-COVID, FiQA, SCIDOCS) で JKO-RAG を評価している。

リトリーバル品質 (nDCG)

• JKO-RAG は、すべての 5 つのデータセットにおいてクロスエンコーダ・リランカーを大幅に上回る。
• ハイパーパラメータが調整されている場合（弱い近接、$h=2.0$）、JKO-RAG は強力な KL 近似ベースラインと同等の nDCG スコアを達成する。これは、幾何学的優位性が主に生のランキング精度を駆動するのではなく、理論通りであることを裏付けている。

堅牢性と安定性

JKO-RAG の決定的な優位性は、理論的予測と一致する堅牢性の指標に現れている：

• クエリのパラフレーズ: 語彙的な摂動（ストップワードの削除、ヘッジ語の追加）の下で、JKO-RAG は KL 近似ベースラインよりも22–38% 高い安定性（分布間の Wasserstein 距離で測定）を示す。
• ディストラクター（妨害要素）の注入: JKO-RAG は、KL と比較して、注入されたハードなディストラクター（近接重複）をトップ10の結果に混入させる割合が 2倍少ない。
• 頻度分析: 幾何学的収縮は Gibbs-affinity グラフの周波数帯域全体で一様であり、クラスター内およびクラスター間の摂動の両方を減衰させる。
• 認定半径 (Certified Radius): 新しい分析により、堅牢性を「分布的安定性」（JKO-RAG が優れる）と「厳密な集合の剛性」（より鋭い KL マップの方が剛性が高い）に分解している。著者らは、重み付き集合を条件とするダウンストリーム LM に対しては、分布的な概念が作用するものであると主張している。

意義と主張

本論文は、JKO-RAG が理論的に裏付けられた分布的リトリーバルのフレームワークを提供すると主張している。その意義は以下の点にある：

1. メカニズムの説明: 「Wasserstein は幾何学を考慮している」という直感を、検証可能な定理（定理 1）と、ステップサイズ依存性に関する反証可能な予測（系 1）へと昇華させ、これらはいずれも経験的に検証されている。
2. 直交する貢献: このフレームワークはバックボーンに依存せず、あらゆる第一段階のリトリーバの上に構成可能である。安定性の向上はモデルの規模とは直交しており、より大きなエンコーダに置き換えても絶対的なスコアは向上するが、$W^2$ 対 KL の安定性分析自体は変わらない。
3. 新しい指標: リトリーバルのための「認定半径」分析を導入し、幾理学的な OT 双対が、選択的棄却のための情報豊富な信頼信号として機能できることを示している。

著者らは限界についても認めており、理論は単一ステップの一次（線形化）分析であること、および本手法は事前に選択された有限のプールに対して動作することを指摘している。しかし、彼らはこの研究を、ニューラル・リトリーバルの幾何学的ダイナミクスを理解するための生産的な研究プラットフォームとして位置づけている。