arXiv:2607.24776cs.IRcs.LG2026-07

用几何方法优化检索,让结果更稳定且少冗余。

JKO-RAG: Distributional Retrieval as Wasserstein Free-Energy Gradient Flow

  • 将检索重排看作在语义空间中沿最优传输路径最小化能量函数。
  • 在五个数据集上超越交叉编码器,对改写查询稳定性提升22%-38%。
  • 理论可验证,适合追求鲁棒检索的系统开发者。

RAG 管道返回的是经过排序的段落列表,但我们认为这存在根本性错配:下游语言模型依赖的是一个集合,而选择问题本质上是几何问题。我们提出 JKO-RAG,将重排建模为在 Wasserstein-2 梯度流下最小化自由能泛函 $F(p) = ext{相关性} + ext{熵} + ext{冗余}$,采用 Jordan--Kinderlehrer--Otto 的近似变分方案。基础度量 $C_{ij} = (1 - an heta_{ij})^2$ 编码嵌入流形的语义几何结构。核心贡献是提出线性响应理论解释为何 Wasserstein 几何更优:与 KL 相比,W² 的近似海森矩阵是稠密且几何感知的,而 KL 是对角且几何无感的,这一差异抑制了查询同义句引起的质量传递。该理论给出可证伪预测:稳定性优势随步长 $h$ 单调递减。我们通过自由能下降、频域扰动响应、$h$-依赖性验证和认证半径分析加以实证。引入四项扩展: extbf{ mjko}(学习型基础度量)、 extbf{wjko}($W^2$–KL 插值)、 extbf{ amjko}(速度提升2倍)、 extbf{ dualrank}(OT 对偶势作为置信度信号)。在五个 BEIR 基准测试中,JKO 在所有任务上均优于交叉编码器;其决定性优势在于鲁棒性——在改写查询下稳定性提高22%~38%,泄露干扰项减少约2倍。

原文摘要 · Abstract (English)

RAG pipelines return a \emph{ranked list} of passages. We argue this is a mismatch: the downstream language model conditions on a \emph{set}, and the selection problem is fundamentally geometric. We propose \jko, which frames reranking as minimising a free-energy functional $F(p)=\text{relevance}+\text{entropy}+\text{redundancy}$ under Wasserstein-2 gradient flow via the Jordan--Kinderlehrer--Otto proximal scheme. The ground metric $C_{ij}=(1-\cos\langle z_i,z_j\rangle)^2$ encodes the semantic geometry of the embedding manifold. Our central contribution is a \emph{linear-response theory} explaining \emph{why} the Wasserstein geometry helps: the Wasserstein and KL retrieval maps differ only in their proximal Hessian -- dense and geometry-aware for $W^2$, diagonal and geometry-blind for KL -- and this difference damps the mass transport that query paraphrase induces. The theory yields a falsifiable prediction: the stability advantage is monotonically decreasing in step size $h$. We verify this empirically via free-energy descent, frequency-resolved perturbation response, the predicted $h$-dependence, and a certified-radius analysis. Four extensions are introduced: \textbf{\nmjko} (learned ground metric), \textbf{\bwjko} ($W^2$--KL interpolation), \textbf{\samjko} ($2\times$ speedup), and \textbf{\dualrank} (OT dual potentials as confidence signals). Across five BEIR benchmarks, \jko\ outperforms the cross-encoder on all five; the decisive advantage is robustness -- 22--38\% more stable under paraphrase, $2\times$ fewer leaked distractors.

检索增强最优传输几何建模鲁棒性

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。