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

• URL: http://arxiv.org/abs/2607.24776v1
• Date: Mon, 15 Jun 2026 04:34:49 GMT
• Title: JKO-RAG: Distributional Retrieval as Wasserstein Free-Energy Gradient Flow
• Abstract summary: We propose jko, which frames reranking as minimising a free-energy functional $F(p)=textrelevance+textentropy+textredundancy$ under Wasserstein-2 gradient flow.<n>Four extensions are introduced: textbfnmjko (learned ground metric), textbfbwjko ($W2$-KL), textbfsamjko ($2times$ speedup), and textbfdualrank (OT dual potentials as confidence signals
• Score: 0.0
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Rate-Distortion Function for Encrypted Traffic Side-Channel Defense [23.361601797200475]
  We prove that the optimal defense has an exponential-tilting structure governed by KKT conditions.<n>On real-world website-fingerprinting defenses, the framework locates Front ($_mathrmgap=0.028$,bits), WTF-PAD ($0.034$,bits, and TrafficSliver ($0.124$,bits) above the theoretical curve.
  arXiv  Detail & Related papers  (2026-07-20T12:38:49Z)
• OAM-Induced Lattice Rotation Reveals a Fractional Optimum in Fault-Tolerant GKP Quantum Sensing [0.0]
  We establish that orbital-angular-momentum encoding and Gottesman-Kitaev-Preskill (GKP) lattice geometry are structurally coupled.<n>Using an end-to-end differentiable Strawberry Fields--TensorFlow circuit, we jointly optimise $ell$, the lattice aspect ratio $r$, and the finite-energy envelope $$ to maximise quantum Fisher information.
  arXiv  Detail & Related papers  (2026-05-13T09:49:16Z)
• Which Coherence Decoheres? Basis-Dependent Decoherence Rates in Symmetry-Broken Collective Spin Systems [0.0]
  Lindblad dephasing rates differ by a factor approaching $2$ as $Ntoinfty$ and reaching $2.42$ near the quantum-critical crossover.<n>The discrepancy has a single algebraic origin: parity forces $langle E_i|hatJ_z|E_irangle=0$ exactly, eliminating the cross-term that doubles the localised-state rate.
  arXiv  Detail & Related papers  (2026-05-01T11:39:51Z)
• Optimal Scalar Quantization for Matrix Multiplication: Closed-Form Density and Phase Transition [50.36362492608702]
  We study entrywise scalar quantization of two matrices prior to multiplication.<n>We obtain a closed-form optimal point density [ star(u) propto exp!left(-fracu26right)bigl( (1-2)+2u22bigr), qquad u=fracx_X,, and prove a correlation-driven phase transition.
  arXiv  Detail & Related papers  (2026-03-20T01:53:44Z)
• Implicit Bias of the JKO Scheme [9.659199680879262]
  We study emphJKO-Flow, Wasserstein gradient flow on $J$, in several simple numerical examples.<n>We show that $_k$ are approximated to order $2$ by Wasserstein gradient flow on a emphmodified energy.
  arXiv  Detail & Related papers  (2025-11-18T18:48:37Z)
• Enjoying Non-linearity in Multinomial Logistic Bandits [56.28491566735463]
  We consider the multinomial logistic bandit problem, a variant of where a learner interacts with an environment by selecting actions to maximize expected rewards.<n>We extend recent work on understanding the impact of the non-linearity of the logistic model to the multinomial setting and propose an efficient algorithm.<n>Our method yields a problem-dependent regret bound of order $ smashwidetildemathcalO( R d sqrtKT/kappa_*) $, where $R$ is the norm of the vector of rewards and $K$
  arXiv  Detail & Related papers  (2025-07-07T08:18:25Z)
• Neural Sampling from Boltzmann Densities: Fisher-Rao Curves in the Wasserstein Geometry [1.609940380983903]
  We deal with the task of sampling from an unnormalized Boltzmann density $rho_D$ by learning a Boltzmann curve given by $f_t$. Inspired by M'at'e and Fleuret, we propose an which parametrizes only $f_t$ and fixes an appropriate $v_t$. This corresponds to the Wasserstein flow of the Kullback-Leibler divergence related to Langevin dynamics.
  arXiv  Detail & Related papers  (2024-10-04T09:54:11Z)
• On the $O(\frac{\sqrt{d}}{T^{1/4}})$ Convergence Rate of RMSProp and Its Momentum Extension Measured by $\ell_1$ Norm [54.28350823319057]
  This paper considers the RMSProp and its momentum extension and establishes the convergence rate of $frac1Tsum_k=1T.<n>Our convergence rate matches the lower bound with respect to all the coefficients except the dimension $d$.<n>Our convergence rate can be considered to be analogous to the $frac1Tsum_k=1T.
  arXiv  Detail & Related papers  (2024-02-01T07:21:32Z)
• A Unified Framework for Uniform Signal Recovery in Nonlinear Generative Compressed Sensing [68.80803866919123]
  Under nonlinear measurements, most prior results are non-uniform, i.e., they hold with high probability for a fixed $mathbfx*$ rather than for all $mathbfx*$ simultaneously. Our framework accommodates GCS with 1-bit/uniformly quantized observations and single index models as canonical examples. We also develop a concentration inequality that produces tighter bounds for product processes whose index sets have low metric entropy.
  arXiv  Detail & Related papers  (2023-09-25T17:54:19Z)
• Near-optimal fitting of ellipsoids to random points [68.12685213894112]
  A basic problem of fitting an ellipsoid to random points has connections to low-rank matrix decompositions, independent component analysis, and principal component analysis. We resolve this conjecture up to logarithmic factors by constructing a fitting ellipsoid for some $n = Omega(, d2/mathrmpolylog(d),)$. Our proof demonstrates feasibility of the least squares construction of Saunderson et al. using a convenient decomposition of a certain non-standard random matrix.
  arXiv  Detail & Related papers  (2022-08-19T18:00:34Z)
• Computationally Efficient Horizon-Free Reinforcement Learning for Linear Mixture MDPs [111.75736569611159]
  We propose the first computationally efficient horizon-free algorithm for linear mixture MDPs. Our algorithm adapts a weighted least square estimator for the unknown transitional dynamic. This also improves upon the best-known algorithms in this setting when $sigma_k2$'s are known.
  arXiv  Detail & Related papers  (2022-05-23T17:59:18Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.