Schrödingerization based quantum algorithms for regularized Wasserstein proximal operators

Abstract: We develop a quantum algorithm for the regularized Wasserstein proximal operator, which is a fundamental tool in optimal transport and mean-field games. The regularization introduces a small diffusive term into the continuity equation of the Benamou-Brenier formulation, which results in a forward-backward PDE system consisting of a Fokker-Planck equation and a viscous Hamilton-Jacobi equation with a quadratic Hamiltonian. Through the Cole-Hopf transformation, both equations are converted to forward heat equations, whose coupling requires a Hadamard division to prepare the initial data for the second heat equation and a Hadamard product to recover the terminal density. We solve these heat equations via the Schrödingerization method and implement the Hadamard division and product operations using simple matrix-vector multiplication representations. The complete quantum algorithm prepares an $\varepsilon$-approximation of the terminal density state with $\mathcal{O}(d N_x T \log^2(1/\varepsilon))$ query complexity, up to constants depending on the potential and initial density, where $d$ is the spatial dimension, $N_x$ is the number of grid points per spatial dimension and $T$ is the evolution time. The complexity depends only {\it linearly} on $d N_x$, yielding an {\it exponential} speedup over classical methods, whose cost scales as $N_x^d$ per time step. Numerical experiments validate the effectiveness of the proposed algorithm.

• The paper presents a quantum algorithm that leverages Schrödingerization to linearize and efficiently compute regularized Wasserstein proximal operators, demonstrating exponential quantum advantage.
• The methodology employs the Cole-Hopf transformation to decouple a nonlinear PDE system into two forward heat equations, enabling effective quantum simulation.
• Numerical results validate the approach with 1D/2D instances, highlighting its potential for accelerating high-dimensional optimal transport, generative models, and mean-field games.

Schrödingerization-Based Quantum Algorithms for Regularized Wasserstein Proximal Operators

Overview and Motivation

The paper "Schrödingerization based quantum algorithms for regularized Wasserstein proximal operators" (2606.28752) addresses the computational challenge of evaluating regularized Wasserstein proximal operators (WPO), a primitive that arises ubiquitously in optimal transport, mean-field games, and generative modeling. The WPO, when combined with suitable functionals, serves as the backbone for gradient flows in the Wasserstein space and underpins modern algorithms in sampling, variational inference, and diffusion models. However, high-dimensional distributions and PDE-coupled formulations render classical algorithms intractable due to the curse of dimensionality. The authors systematically derive a quantum algorithm that exploits the Schrödingerization framework, dramatically reducing computational complexity from exponential (in the spatial dimension) to linear, thereby offering an exponential quantum advantage under certain conditions.

Problem Formulation and Structure

The regularized WPO is formulated by discretizing the Benamou-Brenier action with an added diffusive regularization (Fokker-Planck) term: $$\rho_T = \arg\min_{q} \left\{ \mathcal{F}(q) + \frac{1}{2T} \mathcal{W}_2^2(\rho_0, q) \right\} + \text{diffusion regularization}$$ for linear functionals $\mathcal{F}(q) = \int V(x) q(x) dx$ with $V$ a given potential. The resulting Euler-Lagrange system comprises a coupled forward Fokker-Planck equation and a backward viscous Hamilton-Jacobi equation. The authors exploit the specific structure of this system to enable exact linearization.

Cole-Hopf Reduction

Crucially, via the Cole-Hopf transformation, this coupled nonlinear PDE system is converted into two decoupled forward heat equations for auxiliary fields: $\eta$ (arising from the viscous Hamilton-Jacobi equation) and $\psi$ (from the Fokker-Planck equation), with explicit coupling between initial/terminal conditions:

• Solve for $\eta$ forwards in time from $\eta_0(x) = e^{-V(x)/2\beta}$.
• Prepare the initial condition for $\psi$ as $\psi_0(x) = \rho_0(x)/\eta(T, x)$ via a Hadamard (componentwise) division.
• Solve for $\psi$ forwards in time, then recover the terminal density via a componentwise product: $\mathcal{F}(q) = \int V(x) q(x) dx$0.

Quantum Algorithmic Framework

Heat Equation Solver—Schrödingerization Method

The Schrödingerization technique recasts non-unitary linear PDE evolution as a block-encoded, unitary process permitting efficient quantum simulation. The heat equation for a discretized field $\mathcal{F}(q) = \int V(x) q(x) dx$1 is mapped to a diagonalizable ODE system in the Fourier basis, and quantum simulation is performed using quantum signal processing and phase estimation blocks. The complexity for preparing an $\mathcal{F}(q) = \int V(x) q(x) dx$2-approximate quantum state representing $\mathcal{F}(q) = \int V(x) q(x) dx$3 is: $\mathcal{F}(q) = \int V(x) q(x) dx$4 where $\mathcal{F}(q) = \int V(x) q(x) dx$5 is dimension, $\mathcal{F}(q) = \int V(x) q(x) dx$6 grid points per coordinate, $\mathcal{F}(q) = \int V(x) q(x) dx$7 total evolution time, and $\mathcal{F}(q) = \int V(x) q(x) dx$8 target precision. This linear scaling is the origin of the exponential speedup over classical approaches, which scale as $\mathcal{F}(q) = \int V(x) q(x) dx$9 per time step.

Quantum Hadamard Product and Division

The essential non-unitary interactions—Hadamard product for density recovery, and Hadamard division for coupling the auxiliary equations—are implemented via diagonal block encoding and quantum linear systems algorithms (QLSA). Given state preparation oracles, one can construct circuits that, with query complexity polynomial in condition numbers and logarithmic in $V$0, prepare the required normalized states. Notably, the Hadamard division is implemented efficiently (to $V$1 precision, with $V$2 the condition number), a nontrivial extension beyond prior art on quantum multipliers.

Complexity and Data Dependence

An important technical contribution is the explicit norm and condition number analysis, showing that all quantities involved in normalization, division, and products grow at most exponentially with the range of the potential $V$3. With the scaling-invariance of the WPO and suitable grid choices, the resulting exponential factor is decoupled from the discretization and regularization parameter, leading to the final query complexity: $V$4 All dominant computational costs are thus reduced from exponential to linear in the dimensional grid size.

Numerical Results and Empirical Validation

The developed algorithm is validated with 1D/2D regularized WPO instances using realistic input densities and potentials. The quantum algorithm outputs are cross-verified with classical solvers (both via heat equation and kernel integral methods) and demonstrate agreement up to discretization and numerical errors. Systematic parameter sweeps with respect to regularization $V$5 and transport time $V$6 confirm expected qualitative dependence of the final density (spreading, mass concentration), matching theoretical predictions.

Implications and Future Directions

This work rigorously establishes that quantum algorithms can efficiently implement regularized Wasserstein proximal operators, a central primitive in optimization, PDEs, and AI. The practical upshot is that quantum computers may, in principle, evade the high-dimensional bottleneck that paralyzes classical simulation of optimal transport, mean-field games, or generative models based on diffusion/score-based methods. Robust quantum WPO subroutines could synthesize with quantum generative modeling, optimal sampling, PDE solvers, and control theory for solving currently intractable high-dimensional instances.

A few speculative directions with substantial potential impact include:

• Quantum-accelerated JKO time-discretization for nonlinear PDEs (e.g., nonlinear Fokker-Planck, porous media equations).
• Incorporation into variational and score-based quantum generative models, leveraging the mathematical equivalence between WPO flows and reverse-time SDEs.
• Quantum solution of equilibrium in high-dimensional mean-field games, where WPO arises in discretized Nash equilibria.
• Extensions to Schrödinger bridge and controlled-diffusion problems in quantum stochastic optimal control.

The authors note that generalization beyond linear energy functionals $V$7 (to nonlinear functionals, interaction energies) and further optimization of state preparation for inputs remain challenging open problems.

Conclusion

This paper advances the state of quantum scientific computing by delivering a comprehensive quantum algorithm for regularized Wasserstein proximal operators, tightly integrating the mathematical structure of PDEs, optimal transport theory, and advanced quantum simulation primitives. Through precise reduction to heat equations and optimal quantum implementations of matrix-vector transformations, the approach achieves exponential complexity reductions relative to all known classical algorithms. This creates a rigorous pathway to quantum speedups in optimal transport-based algorithms, paving the way for quantum-accelerated applications in high-dimensional scientific and AI domains.