Trust-Region Optimization in Wasserstein Space: A Second-Order Method for Non-Convex Energy Minimization

Background: Why Optimize in Wasserstein Space?

In both particle system modeling and probability distribution approximation, a shared mathematical core is minimizing potential-interaction energies in Wasserstein space. Whether you're searching for low-energy configurations of interacting particles or approximating a target distribution with a set of samples, the problem ultimately reduces to energy minimization over the space of probability measures with finite second moments.

The challenge is that these energy functions are often non-convex. Non-convexity means first-order methods (such as gradient flows) can easily get trapped in regions with poor local structure, making convergence speed and reliability hard to guarantee. The natural remedy is to incorporate second-order information — that is, to use the curvature (Hessian) of the energy to accelerate and stabilize iterations. But second-order models are only locally reliable, and controlling the trustworthy range of these local approximations is precisely the central concern of this paper.

The Core Method: Transplanting Trust-Region Framework to Wasserstein Space

This paper systematically transfers the classical trust-region optimization method from numerical optimization to the Wasserstein space of probability measures. The fundamental philosophy of trust-region methods is: define a "trustworthy" region around the current point, approximate the objective with a quadratic model within that region, solve for the minimizer, and then decide whether to accept the step and how to adjust the trust-region radius based on the ratio of actual to predicted decrease.

The construction of this method in Wasserstein space involves several key components:

Quadratic Models Along Pushforward Curves

The method constructs quadratic models along pushforward curves, using an explicit self-adjoint second-order variational operator to characterize the curvature of the energy in the space of measures. This allows second-order information to be embedded in the iterative process in a mathematically rigorous way.

Pushforward curves are a natural way to connect two measures in Wasserstein space: given a reference measure $\rho$ and a transport map $T$, one defines the curve $\rho_t = ((1-t)\text{id} + tT)_# \rho$ — the family of pushforward measures along linearly interpolated maps. At $t=0$, the tangent vector of this curve is exactly $T - \text{id}$, corresponding to the "displacement" direction in Wasserstein space. By expanding the energy as a Taylor series along such curves, one can extract the first-order gradient and second-order Hessian operator, providing a geometrically consistent basis for constructing quadratic approximations. Unlike in Euclidean space, the Hessian here is a self-adjoint operator acting on the function space $L^2(\rho)$, and its spectral properties directly affect the difficulty of solving the trust-region subproblem.

$L^2$ Step-Size Radius and the Steihaug-Toint Subproblem Solver

The trust-region "radius" is defined using the $L^2(\rho)$ norm to measure step size. The subproblem — minimizing the quadratic model subject to the trust-region constraint — is solved using the Steihaug-Toint solver, a well-established tool in numerical optimization for large-scale trust-region subproblems, here adapted to the inner product structure of measure space.

The Steihaug-Toint method is a classical algorithm for solving large-scale trust-region subproblems $\min_{|s| \leq \Delta} m(s) = g^\top s + \frac{1}{2} s^\top B s$. Its core idea is to monitor during conjugate gradient (CG) iterations whether the trust-region boundary is encountered: if the CG direction hits negative curvature or the step exceeds the radius, an exact line search is performed to the boundary and the algorithm terminates; otherwise, iteration continues to convergence. This approach avoids a full eigendecomposition of the Hessian — each iteration only requires one matrix-vector product — making it especially suitable for large-scale or operator-form Hessians. In Wasserstein space, "matrix-vector products" are generalized to the action of the second-order variational operator on functions in $L^2(\rho)$, preserving the computational efficiency of the algorithm.

Ratio Test and Radius Update

Each iteration is governed by a ratio test: if the actual energy decrease is sufficiently consistent with the model's predicted decrease, the step is accepted and the trust-region may be expanded; otherwise, the step is rejected and the radius is shrunk. This mechanism ensures that even when the quadratic model is inaccurate in some regions, the overall iteration remains controlled.

Wasserstein space (also known as optimal transport space) is a metric space of probability measures with finite second moments, where the distance is defined by the optimal transport cost (the 2-Wasserstein distance): $W_2(\mu,\nu)^2 = \inf_{\gamma \in \Pi(\mu,\nu)} \int |x-y|^2 , d\gamma(x,y)$, where the infimum is taken over all couplings with marginals $\mu$ and $\nu$. This space naturally carries a Riemannian manifold structure — the tangent space consists of $L^2$ gradient vector fields, and gradient flows correspond to PDEs (such as the Fokker-Planck equation). It is precisely this geometric structure that enables classical Euclidean optimization theory to be transferred to the space of measures, though concepts like curvature and constraints must be carefully redefined — which is the core technical challenge in extending trust-region methods to this setting.

Theoretical Guarantees: Convergence and Complexity Analysis

The paper provides rigorous convergence results under two sets of assumptions, which constitute the most substantial theoretical contribution of the work.

Under the conditions that the energy is bounded below and that the Hessians of both the potential and the interaction kernel are globally bounded, the authors prove three key properties: the objective energy is monotonically non-increasing, the Wasserstein gradient norm converges to zero, and an $\varepsilon$-stationary point is reached within $O(\varepsilon^{-2})$ outer iterations (including rejected ones). This complexity matches classical results for trust-region methods in non-convex optimization, showing that the extension to Wasserstein space incurs no additional convergence cost.

Furthermore, if the potential satisfies a quadratic coercivity condition, then every weak accumulation point is a stationary point. This analysis holds for any initial measure with finite second moments, giving the results a quite broad scope.

A particularly noteworthy result concerns empirical measures: in this case, the entire iteration reduces to a finite-dimensional trust-region method under the $L^2(\rho_N)$ inner product, and when the initial objective gap is uniformly bounded, the complexity constant is independent of both the number of particles and the dimension. This means the method has strong scalability potential for large-scale particle systems or high-dimensional problems — it does not degrade significantly as particle count or dimension increases.

Numerical Experiments: From Soft-Particle Energy to MMD Minimization

The experimental section covers several representative scenarios to validate the method's practical performance:

• Smooth soft-particle energy: a canonical test problem for particle interaction modeling;
• Maximum Mean Discrepancy (MMD) minimization: distribution approximation tasks targeting non-Gaussian distributions, where MMD is a widely used distributional distance in machine learning;
• Component ablation studies: individually verifying the contribution of the quadratic model, trust-region radius, and subproblem solver to overall performance;
• Scaling studies: testing the method's scalability in both the number of particles and the dimension.

The experimental design is well-structured — covering both physically motivated energy minimization problems and distribution approximation tasks in the machine learning context, while using ablation and scaling studies to corroborate the theoretical claims about complexity being independent of particle count and dimension.

Maximum Mean Discrepancy (MMD) is a probability distribution distance measure based on reproducing kernel Hilbert spaces (RKHS), defined as $\text{MMD}^2(\mu,\nu) = \mathbb{E}{x,x'\sim\mu}[k(x,x')] - 2\mathbb{E}{x\sim\mu,y\sim\nu}[k(x,y)] + \mathbb{E}_{y,y'\sim\nu}[k(y,y')]$, where $k$ is a positive definite kernel (e.g., the Gaussian kernel). MMD minimization is widely used as a distribution approximation objective in generative models (such as MMD-GAN), variational inference, and particle methods. For empirical measures, MMD has a differentiable closed-form gradient and fits naturally within the potential-interaction energy framework for particle systems, making it an ideal benchmark for validating the applicability of Wasserstein trust-region methods in machine learning settings.

Significance and Positioning

The value of this work lies in its rigorous integration of mature second-order optimization techniques with Wasserstein geometry, a central tool in modern probabilistic modeling. Compared to purely first-order Wasserstein gradient flow methods, the trust-region framework provides explicit control over the reliability of local approximations, offering greater robustness in non-convex energy landscapes. The theoretical result showing complexity independent of particle count and dimension leaves open exciting possibilities for applying this method to large-scale sampling, generative modeling, and particle system simulation.

For researchers working on optimal transport, variational inference, and particle methods, this is a paper that delivers solid contributions on both the algorithmic design and convergence theory fronts.