0

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

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.

Preview
Year
2026
Hosting
Full text hostedCC-BY-4.0

Cite

Notes

Only stored in your browser.

Attribution

Abstract & full text
arxiv.org/abs/2607.24776CC-BY-4.0
TL;DR
Semantic Scholar
Attribution policy →

Abstract

RAG pipelines return a ranked list of passages. We argue this is a mismatch: the downstream language model conditions on a set, and the selection problem is fundamentally geometric. We propose \jko, which frames reranking as minimising a free-energy functional F(p)=relevance+entropy+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 linear-response theory explaining 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: \nmjko (learned ground metric), \bwjko (W^2--KL interpolation), \samjko (2\times speedup), and \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.