{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/arxiv-2607-24776","title":"JKO-RAG: Distributional Retrieval as Wasserstein Free-Energy Gradient Flow","arxiv_id":"2607.24776","date":"2026-06-15","proceeding":null,"authors":["Levi Segal","Murari Ambati"],"abstract":"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.","url_abs":"https://arxiv.org/abs/2607.24776","url_pdf":"https://arxiv.org/pdf/2607.24776","source":{"archive":null,"snapshot":"2025-07-28","note":"not in the Papers with Code archive (frozen at the snapshot)","row_kind":"graph","title_abstract_authors_date":"arXiv metadata, CC0 1.0 (https://info.arxiv.org/help/license)"},"code_links":[],"tasks":[],"methods":[],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"atlas_url":null,"mcp":null,"developers":"https://syntology.ai/developers","read_at":"2026-09-24T18:15:14+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"mentioned_in_github":null,"is_official":null,"provenance":"deterministic:regex_extraction","mentioned_in_paper":null,"url":"https://github.com/MurariAmbati/jko-rag","reach":{"status":"gone","observed_at":"2026-09-16","how":"tree_404+repo_404"}}],"summary":{},"by_repo_kind":{},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":0,"samples":[]},"arxiv_metadata":{"licence":"arXiv metadata, CC0 1.0 (https://info.arxiv.org/help/license)","fields":["title","abstract","authors","date"],"primary_category":"cs.LG","source":"arxiv_2026.jsonl"},"syntology_extracted_results":null}