Papers › JKO-RAG: Distributional Retrieval as Wasserstein Free-Energy Gradient Flow
JKO-RAG: Distributional Retrieval as Wasserstein Free-Energy Gradient Flow
Levi Segal, Murari Ambati
Title, abstract, authors and date from arXiv's metadata (CC0); this paper is not in the Papers with Code archive (frozen 2025-07-28).
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)=relevance+entropy+redundancy under Wasserstein-2 gradient flow via the Jordan--Kinderlehrer--Otto proximal scheme. The ground metric Cᵢⱼ=(1-cos⟨zᵢ,zⱼ⟩)² 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², 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²--KL interpolation), \textbf{\samjko} (2× 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× fewer leaked distractors.
Code
Reachability, where shown, is from one Syntology probe window (2026-09-16 to 2026-09-18); repositories not probed show nothing. GitHub stars are not tracked.
Code Syntology ran Syntology
Syntology holds the repository link but has not harvested or run code from it.
Results from the paper
The Papers with Code archive ends with its 2025-07-28 snapshot. This paper's arXiv identifier, 2607.24776, was issued in July 2026, after that date, so the archive has no leaderboard rows for it.
Report a problem or propose a change · a person checks every report against the paper or source before anything changes; decisions are listed on /corrections