Generalized Local Squared Wasserstein-2 Loss

• Generalized Local Squared Wasserstein-2 Loss is a family of localized squared W2 objectives that customizes transport penalties based on context-dependent criteria.
• It constructs the loss by summing or averaging squared W2 discrepancies computed over local patches, time slices, or parameter neighborhoods to capture nuanced variations.
• Applications span spline interpolation, diffusion regularization in images, conditional matching in inverse problems, and Bayesian asymptotics, demonstrating its versatility in improving model fidelity.

Generalized Local Squared Wasserstein-2 Loss denotes a family of objectives that replace a single global transport discrepancy by a sum, average, or asymptotic expansion of squared $W_2$ terms computed in localized contexts. In the literature, the localization mechanism is not uniform: it may be temporal and barycentric along a discrete curve of measures, infinitesimal in the Riemannian geometry induced by $W_2$, conditional on neighborhoods in input or initial-condition space, or local in parameter space through second-order expansions of $W_2^2$-induced risks. The shared principle is the same: penalize deviations from a local Wasserstein reference rather than only a global mismatch (Justiniano et al., 2023, Lin et al., 2019, Xia et al., 2024, Requadt et al., 7 Oct 2025).

1. Multiple meanings of locality

The expression is used for several related constructions rather than a single standardized formula. This suggests that it is best understood as a family name for squared $W_2$-based local objectives whose precise form depends on the ambient problem: spline interpolation in Wasserstein space, smoothness regularization on image manifolds, conditional distribution matching in uncertainty quantification, or local quadratic approximation in Bayesian asymptotics (Justiniano et al., 2023, Lin et al., 2019, Xia et al., 7 Mar 2025, Xia et al., 7 Jul 2025).

Setting	Locality mechanism	Representative form
Wasserstein splines	interior time index $k$, anchored barycentric average of $\mu_{k-1},\mu_{k+1}$ around $\mu_k$	$$4K^3\sum_{k=1}^{K-1}W_2^2\!\big(\mu_k,\operatorname{Bar}_{\mu_k}(\mu_{k-1},\mu_{k+1})\big)+\delta K\sum_{k=0}^{K-1}W_2^2(\mu_k,\mu_{k+1})$$
W2 diffusion regularization	small $W_2$-geodesic displacements in input space	$$\mathcal L_{\mathrm{GLSW2}}(f)\approx \frac{\eta^2}{2}\,\mathbb E_x\,\operatorname{Tr}\!\big(J_f(x)^\top L(x)J_f(x)\big)$$
Conditional UQ / inverse problems	neighborhoods $W_2$0 in input space	$W_2$1
Time-decoupled dynamics	neighborhoods in initial conditions and separate time slices	$W_2$2
Parametric Bayes asymptotics	local expansion around $W_2$3	$W_2$4

A recurrent misconception is that “local” must mean spatial locality in the underlying state space. In these works it can instead mean temporal locality, tangent-space locality, conditional locality with respect to covariates, or local asymptotic structure in parameter space. Likewise, “generalized” may refer either to anchored barycenters, generalized costs for mixed continuous–categorical outputs, time-decoupled formulations, or unbalanced transport with source terms.

2. Optimal-transport foundations and generalized variants

For probability measures $W_2$5 with finite second moments, the squared Wasserstein-2 distance is defined by

$W_2$6

and in the Benamou–Brenier dynamic formulation,

$W_2$7

In the spline framework, this Riemannian viewpoint is combined with local barycentric constructions. For $W_2$8 and $W_2$9, the $W_2^2$0-barycenter is

$W_2^2$1

while the anchored generalized barycenter with base $W_2^2$2 is

$W_2^2$3

for a three-marginal optimal plan $W_2^2$4 whose $W_2^2$5- and $W_2^2$6-marginals are optimal couplings (Justiniano et al., 2023).

A distinct generalization allows mass removal and creation. For $W_2^2$7 and $W_2^2$8,

$W_2^2$9

and

$W_2$0

For $W_2$1, the generalized Benamou–Brenier formula states

$W_2$2

with

$W_2$3

On this basis, a “Generalized Local Squared Wasserstein-2 Loss” can be formed patchwise,

$W_2$4

thereby localizing both transport and source penalties (Piccoli et al., 2013).

3. Anchored barycentric spline loss in Wasserstein space

The most explicit practical identification of the generalized local squared $W_2$5 loss appears in the discrete spline model for measure-valued interpolation. The continuous spline energy is defined as

$W_2$6

with action regularization

$W_2$7

and regularized energy $W_2$8. After time discretization with $W_2$9, the discrete action is

$k$0

The non-anchored squared-acceleration term is

$k$1

whereas the anchored or generalized local version is

$k$2

The corresponding regularized energies are $k$3 and $k$4. The practical loss is then

$k$5

subject to interpolation constraints $k$6 and one discrete boundary condition among natural, Hermite, or periodic (Justiniano et al., 2023).

For $k$7, existence of discrete regularized spline interpolations follows from tightness and lower semicontinuity. On the Gaussian subspace $k$8, $k$9 is explicit,

$\mu_{k-1},\mu_{k+1}$0

the anchored barycenter covariance is explicit, and the discrete energies satisfy the consistency estimates

$\mu_{k-1},\mu_{k+1}$1

Numerically, the fully discrete functional is optimized with entropy-regularized OT via Sinkhorn and a variant of Nesterov’s accelerated gradient descent on simplex-constrained weights. The update uses extrapolation, automatic differentiation through Sinkhorn, multiplicative updates

$\mu_{k-1},\mu_{k+1}$2

followed by renormalization. The paper explicitly states that it does not use iPALM; rather, it uses an accelerated first-order scheme tailored to the simplex constraints. The method is applied to numerical spline interpolation and to synthesized textures, where local feature distributions are spline-interpolated and then matched by a generator in feature space.

4. Infinitesimal geometry, diffusion regularization, and local quadratic asymptotics

In discriminative learning, a generalized local squared $\mu_{k-1},\mu_{k+1}$3 loss arises from the Riemannian geometry induced by Wasserstein-2 on images or histograms. If $\mu_{k-1},\mu_{k+1}$4 denotes an image and $\mu_{k-1},\mu_{k+1}$5 is the inverse of a weighted graph Laplacian $\mu_{k-1},\mu_{k+1}$6, then locally

$\mu_{k-1},\mu_{k+1}$7

Using an input-dependent additive noise model

$\mu_{k-1},\mu_{k+1}$8

the second-order expansion of the augmented risk yields

$\mu_{k-1},\mu_{k+1}$9

The stand-alone smoothness version is

$\mu_k$0

with scalar specialization $\mu_k$1 in the small-displacement expansion (Lin et al., 2019).

This formulation makes locality metric-dependent rather than patch-dependent. The operator $\mu_k$2 privileges mass-transport directions on the pixel graph, so the resulting regularizer is anisotropic and input-dependent, unlike isotropic Jacobian penalties. The same paper reports empirical robustness improvements on CIFAR-10, including natural test error $\mu_k$3 and FGSM $\mu_k$4 robust test error $\mu_k$5 for W2 grad regularization, with convolutional implementations whose overhead is negligible relative to standard backpropagation (Lin et al., 2019).

A different local interpretation appears in Bayesian asymptotics. For a parametric family $\mu_k$6, the squared Wasserstein loss

$\mu_k$7

is shown, near $\mu_k$8, to admit a local quadratic expansion

$\mu_k$9

with $$4K^3\sum_{k=1}^{K-1}W_2^2\!\big(\mu_k,\operatorname{Bar}_{\mu_k}(\mu_{k-1},\mu_{k+1})\big)+\delta K\sum_{k=0}^{K-1}W_2^2(\mu_k,\mu_{k+1})$$0 positive definite and $$4K^3\sum_{k=1}^{K-1}W_2^2\!\big(\mu_k,\operatorname{Bar}_{\mu_k}(\mu_{k-1},\mu_{k+1})\big)+\delta K\sum_{k=0}^{K-1}W_2^2(\mu_k,\mu_{k+1})$$1 locally. In this setting, “generalized local squared $$4K^3\sum_{k=1}^{K-1}W_2^2\!\big(\mu_k,\operatorname{Bar}_{\mu_k}(\mu_{k-1},\mu_{k+1})\big)+\delta K\sum_{k=0}^{K-1}W_2^2(\mu_k,\mu_{k+1})$$2 loss” refers to the locally quadratic structure of $$4K^3\sum_{k=1}^{K-1}W_2^2\!\big(\mu_k,\operatorname{Bar}_{\mu_k}(\mu_{k-1},\mu_{k+1})\big)+\delta K\sum_{k=0}^{K-1}W_2^2(\mu_k,\mu_{k+1})$$3 as a function of parameters, which, combined with Bernstein–von Mises theory, yields asymptotic normality and efficiency of Bayes estimators under broad regularity conditions (Requadt et al., 7 Oct 2025). In one-dimensional models the paper also provides an explicit second-derivative formula involving the optimal map $$4K^3\sum_{k=1}^{K-1}W_2^2\!\big(\mu_k,\operatorname{Bar}_{\mu_k}(\mu_{k-1},\mu_{k+1})\big)+\delta K\sum_{k=0}^{K-1}W_2^2(\mu_k,\mu_{k+1})$$4.

A closely related exact reduction occurs for predictive density estimation in location and location-scale families. There, plug-in predictive densities form a complete class under $$4K^3\sum_{k=1}^{K-1}W_2^2\!\big(\mu_k,\operatorname{Bar}_{\mu_k}(\mu_{k-1},\mu_{k+1})\big)+\delta K\sum_{k=0}^{K-1}W_2^2(\mu_k,\mu_{k+1})$$5, and the Bayes predictive density is the plug-in density evaluated at posterior means of location and scale parameters. In location-scale families with spherical scale, predictive risk reduces to

$$4K^3\sum_{k=1}^{K-1}W_2^2\!\big(\mu_k,\operatorname{Bar}_{\mu_k}(\mu_{k-1},\mu_{k+1})\big)+\delta K\sum_{k=0}^{K-1}W_2^2(\mu_k,\mu_{k+1})$$6

so the Wasserstein loss becomes a quadratic parameter loss (Matsuda et al., 2019).

5. Empirical neighborhood losses, time decoupling, and mixed-type outputs

In uncertainty quantification and inverse problems with latent randomness, the local squared Wasserstein-2 method matches conditional output laws in neighborhoods of the input. For $$4K^3\sum_{k=1}^{K-1}W_2^2\!\big(\mu_k,\operatorname{Bar}_{\mu_k}(\mu_{k-1},\mu_{k+1})\big)+\delta K\sum_{k=0}^{K-1}W_2^2(\mu_k,\mu_{k+1})$$7 and an approximate model $$4K^3\sum_{k=1}^{K-1}W_2^2\!\big(\mu_k,\operatorname{Bar}_{\mu_k}(\mu_{k-1},\mu_{k+1})\big)+\delta K\sum_{k=0}^{K-1}W_2^2(\mu_k,\mu_{k+1})$$8, the ideal objective is

$$4K^3\sum_{k=1}^{K-1}W_2^2\!\big(\mu_k,\operatorname{Bar}_{\mu_k}(\mu_{k-1},\mu_{k+1})\big)+\delta K\sum_{k=0}^{K-1}W_2^2(\mu_k,\mu_{k+1})$$9

where $W_2$0 and $W_2$1. Because $W_2$2 is unavailable, the empirical local objective is

$W_2$3

where neighborhoods are defined by $W_2$4. The paper solves the discrete OT exactly with Python POT using ot.emd2, without entropic regularization in the main experiments, and proves

$W_2$5

exhibiting the bias–variance trade-off in $W_2$6 (Xia et al., 2024).

For stochastic dynamical systems, locality is combined with time decoupling. Given local empirical conditional distributions $W_2$7 and $W_2$8 at time $W_2$9, the local discrepancy is

$$\mathcal L_{\mathrm{GLSW2}}(f)\approx \frac{\eta^2}{2}\,\mathbb E_x\,\operatorname{Tr}\!\big(J_f(x)^\top L(x)J_f(x)\big)$$0

and the time-decoupled loss is

$$\mathcal L_{\mathrm{GLSW2}}(f)\approx \frac{\eta^2}{2}\,\mathbb E_x\,\operatorname{Tr}\!\big(J_f(x)^\top L(x)J_f(x)\big)$$1

This avoids OT on path space and is proved to be necessary for recovering the parameter distribution: under the stated Lipschitz and moment assumptions,

$$\mathcal L_{\mathrm{GLSW2}}(f)\approx \frac{\eta^2}{2}\,\mathbb E_x\,\operatorname{Tr}\!\big(J_f(x)^\top L(x)J_f(x)\big)$$2

The formulation is developed for ODEs, jump–diffusions, and SPDEs, with POT-based EMD in experiments and Adam or AdamW for training stochastic neural networks (Xia et al., 7 Mar 2025).

A further generalization addresses mixed outputs with both continuous and categorical components. On $$\mathcal L_{\mathrm{GLSW2}}(f)\approx \frac{\eta^2}{2}\,\mathbb E_x\,\operatorname{Tr}\!\big(J_f(x)^\top L(x)J_f(x)\big)$$3, the paper defines the squared mixed-type cost

$$\mathcal L_{\mathrm{GLSW2}}(f)\approx \frac{\eta^2}{2}\,\mathbb E_x\,\operatorname{Tr}\!\big(J_f(x)^\top L(x)J_f(x)\big)$$4

and the generalized Wasserstein-2 distance $$\mathcal L_{\mathrm{GLSW2}}(f)\approx \frac{\eta^2}{2}\,\mathbb E_x\,\operatorname{Tr}\!\big(J_f(x)^\top L(x)J_f(x)\big)$$5 by minimizing the expected value of this cost over couplings. The generalized local squared loss is

$$\mathcal L_{\mathrm{GLSW2}}(f)\approx \frac{\eta^2}{2}\,\mathbb E_x\,\operatorname{Tr}\!\big(J_f(x)^\top L(x)J_f(x)\big)$$6

with a time-averaged extension for spatiotemporal systems. For differentiable training, the categorical part is replaced by a surrogate $$\mathcal L_{\mathrm{GLSW2}}(f)\approx \frac{\eta^2}{2}\,\mathbb E_x\,\operatorname{Tr}\!\big(J_f(x)^\top L(x)J_f(x)\big)$$7 together with

$$\mathcal L_{\mathrm{GLSW2}}(f)\approx \frac{\eta^2}{2}\,\mathbb E_x\,\operatorname{Tr}\!\big(J_f(x)^\top L(x)J_f(x)\big)$$8

so gradients do not flow through the rounding gap. The paper proves a universal approximation theorem for stochastic neural networks under this generalized $$\mathcal L_{\mathrm{GLSW2}}(f)\approx \frac{\eta^2}{2}\,\mathbb E_x\,\operatorname{Tr}\!\big(J_f(x)^\top L(x)J_f(x)\big)$$9 metric and a finite-sample bound analogous in structure to the purely continuous case (Xia et al., 7 Jul 2025).

6. Related design patterns, applications, and limitations

Several adjacent research lines instantiate local squared $W_2$00 ideas in more specialized forms. In generative modeling via restricted convex potentials, the paper on restricted $W_2$01 approximation does not explicitly propose a local $W_2$02 objective, but it states that a local version can be built as

$W_2$03

with local feature maps $W_2$04 and local convex classes $W_2$05. The same work emphasizes restricted moment matching on tangent spaces and projected SGD with ICNN-based convex potentials (Taghvaei et al., 2019). In end-to-end $W_2$06 generative networks, the global quadratic-cost OT objective is trained with ICNN potentials and cycle-consistency regularization, and the provided adaptation to a local squared $W_2$07 loss weights correlations and cycle penalties over neighborhoods or kernels while preserving convexity in the transported variable (Korotin et al., 2019).

In vision, locality can be imposed at the class, patch, or region level. For ordered single-label classification, the squared EMD loss

$W_2$08

penalizes errors according to a ground distance matrix $W_2$09, and the paper also proposes learning $W_2$10 from CNN features during training (Hou et al., 2016). In superpixel-based segmentation, adjacent region histograms $W_2$11 are compared by

$W_2$12

and greedy adjacency-constrained merging produces a local graph-based squared $W_2$13 criterion over superpixels (Huang et al., 22 Jan 2026). In texture interpolation, the spline formulation uses local feature maps $W_2$14 on image patches, constructs feature-space empirical measures $W_2$15, and reports visually superior interpolation relative to metamorphosis splines that blend intensities rather than transporting mass in feature space (Justiniano et al., 2023).

The main limitations are likewise context-dependent. Entropic regularization introduces bias and smoothing, although it improves differentiability and numerical stability; this is emphasized both in spline computation and in OT-based learning more broadly (Justiniano et al., 2023, Korotin et al., 2019). Higher-order spline terms can overshoot near key frames, with the anchored generalized barycenter and $W_2$16 recommended as stabilizers (Justiniano et al., 2023). Neighborhood-based losses depend sensitively on the radius $W_2$17: too small a neighborhood leads to poor empirical estimation, while too large a neighborhood induces bias by mixing distinct conditional laws (Xia et al., 2024, Xia et al., 7 Mar 2025, Xia et al., 7 Jul 2025). Time decoupling reduces computation but may discard temporal-correlation information when parameter effects are expressed primarily through trajectory correlations rather than time marginals (Xia et al., 7 Mar 2025). In mixed-output settings, the choice of categorical ground penalty matters, and the paper notes sensitivity to the constant $W_2$18 in $W_2$19 (Xia et al., 7 Jul 2025).

Taken together, these formulations establish a coherent but plural notion of generalized local squared Wasserstein-2 loss. The core object is always a squared $W_2$20-type comparison, but the operative notion of locality changes with the modeling task: barycentric averaging in time, infinitesimal diffusion in input geometry, empirical conditioning in covariate neighborhoods, time-slice aggregation for stochastic dynamics, patchwise unbalanced transport, or local quadratic structure in parameter space. This diversity is not a defect of the term; it is a direct reflection of the fact that $W_2$21 supplies geometry, while “local” specifies where that geometry is sampled, linearized, or constrained.