💻 computer science

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

The paper proposes JKO-RAG, a novel retrieval reranking framework that models passage selection as a Wasserstein-2 free-energy gradient flow to leverage semantic geometry, thereby achieving superior robustness against query paraphrasing and distractors compared to traditional cross-encoders across multiple benchmarks.

Original authors: Levi Segal, Murari Ambati

Published 2026-07-29

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Levi Segal, Murari Ambati

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to find the perfect playlist for a party. You have a massive library of songs, and you want to pick the best ones. In the world of artificial intelligence, this is called "retrieval." Usually, computers act like a strict DJ who picks songs one by one based on how well they match your request, creating a simple ranked list. But here's the catch: the computer that actually plays the music (the "language model") doesn't listen to songs one by one; it hears the whole playlist at once. If you pick two songs that are almost identical, it's a waste of space, and if you pick songs that clash, the vibe is ruined. Scientists have long known that the group of songs matters more than the individual rankings, and that the "distance" between songs (how similar they sound) is a geometric shape, not just a list. This paper steps into that gap, asking: what if we didn't just pick a list, but instead treated the selection process like a fluid flowing over a landscape to find the most stable, diverse, and relevant mix?

The authors of this paper, Levi Segal and Murari Ambati, propose a new way to build these AI playlists called JKO-RAG. They argue that the old way of picking documents is like trying to arrange a crowd of people by asking each person individually if they want to be there. It ignores how the people interact with each other. Instead, they suggest treating the selection process like a free-energy gradient flow. Think of this as a magical, invisible landscape where "relevance" is a valley (you want to go down there) and "redundancy" (picking the same thing twice) is a hill you want to avoid. The computer doesn't just jump to the bottom; it flows there like water, smoothing out the bumps and avoiding the hills.

The secret sauce in their method is a mathematical tool called Wasserstein geometry. Imagine you have a pile of sand (your documents) and you want to move it to a new spot. A standard method might just grab a handful and move it, ignoring the shape of the ground. The Wasserstein method, however, understands the terrain. It knows that moving a grain of sand from a smooth patch to a bumpy patch costs more energy than moving it across a flat field. By using this "terrain-aware" math, the system becomes incredibly good at keeping the playlist stable. Even if you change the request slightly—like asking for "a song about summer" instead of "a summer song"—the playlist doesn't fall apart. It stays anchored.

The paper finds that this new method is a game-changer for stability, even if it doesn't always win on raw speed or simple ranking scores. When the researchers tested their system, they found that it was 22% to 38% more stable than the previous best methods when the questions were rephrased. It also leaked 2 times fewer "hard distractors" (tricky, fake documents that look like the real thing but aren't). The authors proved mathematically that this happens because their method acts like a shock absorber for the AI's decision-making. When the input changes, the Wasserstein method dampens the movement, keeping the core group of selected documents steady, whereas older methods let the whole group jitter and shift wildly.

They also introduced four cool upgrades to this system:

NM-JKO: A version that learns its own map of the terrain instead of using a generic one.
BW-JKO: A slider that lets you tune how much "geometry" you want, bridging the gap between old and new methods.
SAM-JKO: A speed-boosting trick that makes the process twice as fast without losing quality.
DUAL-RANK: A way to give the AI a "confidence score," so it knows when to say, "I'm not sure, maybe don't pick this one," which helps avoid bad choices.

In short, the paper shows that by treating document selection as a geometric flow rather than a simple list, we can build AI systems that are much more robust and reliable. They didn't just guess this; they ran simulations and mathematical proofs to show exactly why the "geometry-aware" approach works, predicting that the more you tune the system to be "stronger" in its geometric flow, the more stable it becomes. It's a shift from asking "Is this document good?" to "Does this document fit well with the others in this specific landscape?"

Technical Summary: JKO-RAG

Problem Statement

Retrieval-Augmented Generation (RAG) pipelines typically treat the reranking step as a collection of independent binary relevance judgments, producing a ranked list of candidate passages. The authors argue this is a fundamental mismatch: the downstream consumer (a language model) conditions on a set of passages, not a list, and the selection problem is inherently geometric and distributional. Existing diversity-aware methods like Maximal Marginal Relevance (MMR) and Determinantal Point Processes (DPP) address these issues but lack iterative refinement or convergence guarantees. Furthermore, standard reranking ignores the semantic geometry of the embedding space, potentially retrieving near-duplicate passages that score well individually but provide redundant context.

Methodology: JKO-RAG

The paper proposes JKO-RAG, which reframes retrieval as a distributional gradient flow on the probability simplex $\Delta^{M-1}$ over a candidate pool of $M$ documents.

The Framework

The method iterates the Jordan–Kinderlehrer–Otto (JKO) proximal scheme to minimize a free-energy functional $F(p)$ :
$p_{t+1} = \arg \min_{p \in \Delta^{M-1}} \left( \frac{1}{2h} W_{C, \varepsilon}^2(p, p_t) + F_q(p) \right)$

Free Energy Functional ( $F_q$ ): Composed of three terms:
- Data Fidelity: Weighted sum of cross-encoder and dense retrieval scores ( $E_i$ ).
- Entropy: A regularization term ( $\lambda \sum p_i \log p_i$ ) to prevent probability collapse.
- Redundancy: A quadratic penalty ( $\frac{\rho}{2} p^\top K p$ ) where $K_{ij} = \max(0, \cos \langle z_i, z_j \rangle)$ penalizes the selection of semantically similar (near-duplicate) documents.
Proximal Term (Wasserstein-2): The distance metric is the entropic regularized Optimal Transport (OT) cost $W_{C, \varepsilon}^2$ .
- Ground Metric: $C_{ij} = (1 - \cos \langle z_i, z_j \rangle)^2$ , encoding the squared spherical distance on the embedding manifold.
- Solver: The proximal step is solved via Adam optimization on logits, utilizing a Sinkhorn solver with an envelope-theorem trick for efficient backpropagation.

Extensions

The paper introduces four extensions to the core framework:

NM-JKO: Learns a low-rank ground metric via InfoNCE to better align the OT cost with task-specific retrieval geometry.
BW-JKO: Interpolates between KL-divergence (mirror descent) and Wasserstein-2 using a Bregman–Wasserstein parameter $\alpha$ .
SAM-JKO: A score-aware multi-resolution approach that clusters candidates by relevance before running fine-grained JKO, achieving a $2\times$ speedup.
DUAL-RANK: Uses Sinkhorn dual potentials as calibrated confidence signals for selective coverage (abstention).

Key Contributions & Theory

The paper's central contribution is a linear-response theory that explains why the Wasserstein geometry improves robustness, moving beyond empirical observation to a mechanistic proof.

Theoretical Findings

Linear-Response Decomposition (Theorem 1): The authors analyze the sensitivity of a single JKO step to energy perturbations (e.g., query paraphrasing). They show that the response operator $\delta p = -R_D \delta E$ depends on the inverse of a Hessian matrix $A_D$ .
The Geometric Source of Stability: The difference between the Wasserstein ( $W^2$ $W^{2}$ ) and KL-proximal methods lies solely in the proximal Hessian:
- KL: Uses a diagonal, geometry-blind operator ( $\text{diag}(1/p)$ ).
- Wasserstein: Uses a dense, geometry-aware entropic-OT operator ( $H_W$ ).
Geometric Contraction (Proposition 1): The dense Hessian $H_W$ acts as a "stiffness" term that penalizes mass transport across the semantic metric. Consequently, the Wasserstein response contracts the geometric movement of the distribution under perturbation, whereas the KL response allows mass to move freely across clusters.
Step Size Dependence (Corollary 1): The stability advantage is monotonically decreasing with the step size $h$ . As $h \to \infty$ (weak proximal), the operators converge, explaining why tuned configurations (large $h$ ) yield similar nDCG scores. As $h \to 0$ (strong proximal), the stability gap widens.

Experimental Results

The authors evaluate JKO-RAG on five BEIR datasets (SciFact, NFCorpus, TREC-COVID, FiQA, SCIDOCS).

Retrieval Quality (nDCG)

JKO-RAG significantly outperforms the cross-encoder reranker on all five datasets.
When hyperparameters are tuned (weak proximal, $h=2.0$ ), JKO-RAG achieves nDCG scores statistically tied with a strong KL-proximal baseline. This confirms the theory that the geometric advantage does not primarily drive raw ranking accuracy in the tuned regime.

Robustness and Stability

The decisive advantage of JKO-RAG appears in robustness metrics, consistent with the theoretical predictions:

Query Paraphrase: Under lexical perturbations (dropping stopwords, adding hedges), JKO-RAG is 22–38% more stable (measured by Wasserstein distance between distributions) than the KL-proximal baseline.
Distractor Injection: JKO-RAG leaks 2 $\times$ fewer injected hard distractors (near-duplicates) into the top-10 results compared to KL.
Frequency Analysis: The geometric contraction is uniform across frequency bands of the Gibbs-affinity graph, damping both intra-cluster and cross-cluster perturbations.
Certified Radius: A new analysis decomposes robustness into "distributional stability" (where JKO-RAG excels) and "exact-set rigidity" (where the sharper KL map is more rigid). The authors argue that for downstream LMs conditioning on a weighted set, the distributional notion is the operative one.

Significance and Claims

The paper claims that JKO-RAG provides a theoretically grounded framework for distributional retrieval. Its significance lies in:

Mechanistic Explanation: It transforms the intuition that "Wasserstein is geometry-aware" into a verifiable theorem (Theorem 1) and a falsifiable prediction regarding step-size dependence (Corollary 1), both of which are empirically verified.
Orthogonal Contribution: The framework is backbone-agnostic and composable on top of any first-stage retriever. The stability gains are orthogonal to model scale; substituting a larger encoder would improve absolute scores but leave the $W^2$ -vs-KL stability analysis unchanged.
New Metrics: It introduces "certified-radius" analysis for retrieval and demonstrates that geometric OT duals can serve as informative confidence signals for selective abstention.

The authors acknowledge limitations, noting that the theory is a first-order (linearized) analysis of a single step and that the method operates on a pre-selected finite pool. However, they position the work as a productive research platform for understanding the geometric dynamics of neural retrieval.