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

arXiv: https://arxiv.org/abs/2607.24776 · HTML: https://mlsanity.com/paper/2607.24776/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Generation and context integration, one of 15 directions in Retrieval-augmented generation.
- Builds on 14 of its 25 references that are in the map.
- First posted on arXiv 15 Jun 2026.

## Placed in

- Retrieval-augmented generation / [Generation and context integration](https://mlsanity.com/rag/generation-and-context-integration/)

## Builds on (most cited first, 14 of 14)

- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://mlsanity.com/paper/1908.10084/) (2019)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Certified Adversarial Robustness via Randomized Smoothing](https://mlsanity.com/paper/1902.02918/) (2019)
- [Text Embeddings by Weakly-Supervised Contrastive Pre-training](https://mlsanity.com/paper/2212.03533/) (2022)
- [BEIR: A Heterogenous Benchmark for Zero-shot Evaluation of Information Retrieval Models](https://mlsanity.com/paper/2104.08663/) (2021)
- [Passage Re-ranking with BERT](https://mlsanity.com/paper/1901.04085/) (2019)
- [Fact or Fiction: Verifying Scientific Claims](https://mlsanity.com/paper/2004.14974/) (2020)
- [Gradient-based Hyperparameter Optimization through Reversible Learning](https://mlsanity.com/paper/1502.03492/) (2015)
- [Enhancing Retrieval-Augmented Large Language Models with Iterative Retrieval-Generation Synergy](https://mlsanity.com/paper/2305.15294/) (2023)
- [Differential Properties of Sinkhorn Approximation for Learning with Wasserstein Distance](https://mlsanity.com/paper/1805.11897/) (2018)
- [A Unified Framework for Implicit Sinkhorn Differentiation](https://mlsanity.com/paper/2205.06688/) (2022)
- [No Parameter Left Behind: How Distillation and Model Size Affect Zero-Shot Retrieval](https://mlsanity.com/paper/2206.02873/) (2022)

## Abstract (arXiv)

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 Cij=(1-coslangle zi,zjrangle)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 W2, 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 (W2--KL interpolation), samjko (2× 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× fewer leaked distractors.