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

arXiv:2607.24776v1 Announce Type: cross 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)=ext{relevance}+ext{entropy}+ext{redundancy} under Wasserstein-2 gradient flow via the Jordan--Kinderlehrer--Otto proximal scheme. The ground metric C_{ij}=(1-oslangle z_i,z_jrangle)^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: extbf{nmjko} (learned ground metric), extbf{wjko} (W^2--KL interpolation), extbf{samjko} (2imes speedup), and extbf{ualrank} (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, 2imes fewer leaked distractors.

Source: arXiv cs.LG | 2026-07-29