A Local Sinkhorn Framework for Conditional Distribution Reconstruction of Multidimensional Random Fields

# A Local Sinkhorn Framework for Conditional Distribution Reconstruction of Multidimensional Random Fields
Source: [https://arxiv.org/html/2608.11613](https://arxiv.org/html/2608.11613)
\[type=editor, auid=000,bioid=1, orcid=0000\-0002\-2116\-4712\]

\[type=editor, auid=000,bioid=1, orcid=0009\-0009\-6685\-0861\]

††highlights:We propose a scalable local entropically regularized optimal transport framework to train stochastic neural networks for learning random fields\. We derive generalization error bounds that reveal how our proposed framework could efficiently learn multidimensional random fields\. Our proposed approach accurately reconstructs stochastic random fields and dynamical systems and outperforms several machine\-learning\-based benchmark approaches in uncertainty quantification\.Qijing Shenqijing\.shen@ndm\.ox\.ac\.ukorganization=School of Mathematics, University of Birmingham, addressline=Watson Building, Edgbaston, city=Birmingham, postcode=B15 2TT, country=United Kingdomorganization=Department of Mathematics, University of Houston, addressline=3551 Cullen Blvd, city=Houston, postcode=77204, state=Texas, country=United Statesorganization=Nuffield Department of Medicine, University of Oxford, United Kingdom, addressline=Old Road Campus, city=Oxford, postcode=OX3 7BN, country=United Kingdom

April 2026

###### Abstract

In this paper, we propose a local Sinkhorn divergence framework for conditional distribution reconstruction of multidimensional random fields\. By utilizing the debiased Sinkhorn divergence, our proposed approach develops a differentiable and computationally efficient local distribution matching objective to train stochastic neural networks \(SNNs\)\. Furthermore, we establish theoretical generalization error estimates for our local Sinkhorn divergence framework, which explicitly characterizes the trade\-off between approximation bias and statistical efficiency controlled by the regularization parameter and reveals how our proposed local Sinkhorn divergence loss function can be efficiently applied to learning multidimensional random field models\. The proposed framework provides a scalable alternative to exact local optimal transport for conditional distribution reconstruction, offering a practical compromise between geometric fidelity, statistical efficiency, and computational scalability for uncertainty quantification and probabilistic scientific machine learning\. Through various numerical examples, we compare our proposed local Sinkhorn divergence framework with other loss functions to train SNNs and with other machine\-learning\-based uncertainty quantification frameworks, demonstrating that the proposed local Sinkhorn divergence framework achieves an effective balance between reconstruction accuracy and computational efficiency while maintaining good scalability for multidimensional stochastic systems\.

###### keywords

Random Field Reconstruction ,Sinkhorn Divergence ,Uncertainty Quantification ,Stochastic Neural Network ,Scientific Machine Learning

††credit:Conceptualization, Methodology, Formal analysis, Investigation, Software, Writing\-original draft, Writing\-review and editing††credit:Methodology, Formal analysis, Software, Writing\-original draft, Writing\-review and editing††corresponding:Corresponding author## 1Introduction

The identification of random fields from observational data has become an increasingly important topic in uncertainty quantification, stochastic modeling, and scientific machine learning\. Many physical and engineering systems, including porous media flow, turbulent transport, groundwater dynamics, and biological processes, are inherently affected by uncertain parameters or random forcing, leading to stochastic responses that cannot be adequately characterized by deterministic models alone[ghanem1991stochastic](https://arxiv.org/html/2608.11613#bib.bib13);[xiu2010numerical](https://arxiv.org/html/2608.11613#bib.bib33);[sullivan2015introduction](https://arxiv.org/html/2608.11613#bib.bib25)\. Consequently, instead of predicting only the conditional expectation, it is often desirable to reconstruct the entire conditional probability distribution of the stochastic solution, thereby providing a more comprehensive description of uncertainty propagation\. However, the primary challenge is usually not learning the conditional mean, but accurately reconstructing the multidimensional conditional distributions in a computationally scalable manner\.

Recent advances in machine learning have enabled data\-driven approximations of multidimensional stochastic systems without requiring explicit knowledge of the underlying governing equations\. Various deep generative models, including conditional variational autoencoders \(CVAEs\), generative adversarial networks \(GANs\), conditional normalizing flows \(CNFs\), and diffusion models, have demonstrated remarkable capability in learning complex probability distributions from data[kingma2013auto](https://arxiv.org/html/2608.11613#bib.bib17);[goodfellow2014gan](https://arxiv.org/html/2608.11613#bib.bib14);[rezende2015variational](https://arxiv.org/html/2608.11613#bib.bib22);[ho2020denoising](https://arxiv.org/html/2608.11613#bib.bib16)\. Nevertheless, for stochastic random field reconstruction, the learned distributions should preserve the geometric structure of probability measures associated with nearby physical states\. This requirement naturally motivates the use of optimal transport \(OT\), whose Wasserstein metric provides a physically meaningful distance between probability distributions and has recently attracted increasing attention in scientific machine learning[villani2009optimal](https://arxiv.org/html/2608.11613#bib.bib29);[peyre2019computational](https://arxiv.org/html/2608.11613#bib.bib20)\. Compared with divergence\-based discrepancies, such as the Kullback–Leibler divergence or maximum mean discrepancy[gretton2012kernel](https://arxiv.org/html/2608.11613#bib.bib15), Wasserstein distances remain informative even when two probability measures have disjoint supports and provide meaningful gradients during optimization[arjovsky2017wasserstein](https://arxiv.org/html/2608.11613#bib.bib2);[peyre2019computational](https://arxiv.org/html/2608.11613#bib.bib20)\. These advantages have motivated numerous OT\-based machine learning methods, including Wasserstein GANs, gradient\-flow formulations, and OT regularization for scientific computing[frogner2015learning](https://arxiv.org/html/2608.11613#bib.bib10);[arjovsky2017wasserstein](https://arxiv.org/html/2608.11613#bib.bib2)\. However, the exact computation of Wasserstein distances requires solving a large\-scale linear programming problem whose complexity grows rapidly with sample size, making repeated evaluations during neural network training prohibitively expensive\. Our recent work introduced a local squared Wasserstein\-22\(W2W\_\{2\}\) framework for random field reconstruction, in which neighboring samples associated with nearby input states were compared through local optimal transport distances\. By exploiting the locality of conditional distributions, the method was shown to accurately reconstruct stochastic fields from sparse observations while preserving important geometric structures of the underlying probability measures[xia2026local](https://arxiv.org/html/2608.11613#bib.bib31)\. However, the local squaredW2W\_\{2\}distance relies on repeatedly solving exact optimal transport problems during neural network training\. As the number of neighborhood evaluations and samples increases, the computational cost becomes the dominant bottleneck, significantly limiting scalability for multidimensional stochastic systems and large datasets\. Additionally, the generalization error bound may deteriorate when the squaredW2W\_\{2\}distance between two probability measures is estimated using empirical probability measures, particularly in high\-dimensional settings\.

Recently, Cuturi introduced entropic regularization of optimal transport, leading to the Sinkhorn algorithm that dramatically accelerates OT computations while remaining fully differentiable[cuturi2013sinkhorn](https://arxiv.org/html/2608.11613#bib.bib6)\. Subsequently, Feydyet al\.proposed the Sinkhorn divergence, which removes the entropic bias and interpolates smoothly between the exact Wasserstein distance and kernel\-based maximum mean discrepancy[feydy2019interpolating](https://arxiv.org/html/2608.11613#bib.bib8)\. The Sinkhorn divergence has since been successfully applied to generative modeling, Bayesian inference, and scientific machine learning because of its computational efficiency and stable gradient properties[genevay2018learning](https://arxiv.org/html/2608.11613#bib.bib12);[feydy2019interpolating](https://arxiv.org/html/2608.11613#bib.bib8)\. Furthermore, through introducing an extra entropic regularization term, the generalization error when estimating the Sinkhorn divergence using the empirical measures may not deteriorate with the dimensionality of the associated random variables[genevay2018learning](https://arxiv.org/html/2608.11613#bib.bib12)\.

Motivated by these developments, this work develops a local Sinkhorn\-divergence framework for conditional distribution reconstruction of stochastic fields\. Building upon our previous localW2W\_\{2\}formulation, the proposed approach replaces the expensive exact OT computation with Sinkhorn divergence, resulting in substantially improved computational efficiency while preserving the local geometric information of conditional probability distributions\. The proposed local Sinkhorn divergence bridges exact optimal transport and kernel\-based distribution matching by introducing an entropic regularization that provides a controllable bias while significantly improving computational efficiency and statistical scalability\. Owing to the differentiability and favorable numerical stability of Sinkhorn divergence, the resulting learning framework is considerably more efficient in reconstructing conditional distributions of the to\-be\-learned random fields\. Furthermore, we shall show from both theoretical and empirical aspects that incorporating the regularization term could lead to improved accuracy when reconstructing multidimensional random field models using empirical distributions\. Numerical experiments on stochastic Darcy flow and multidimensional stochastic FitzHugh–Nagumo \(FHN\) systems demonstrate that the proposed approach achieves reconstruction accuracy comparable to that of the exact Wasserstein formulation while significantly reducing the overall training cost\.

The main contributions of this work are summarized as follows\.

- •We extend our recent local optimal transport framework for random field reconstruction from the exact Wasserstein distance to the debiased Sinkhorn divergence, leading to a computationally efficient, scalable, and fully differentiable local distribution matching method for training SNNs to reconstruct random fields\.
- •We provide a theoretical analysis of the proposed local Sinkhorn formulation by combining the neighborhood approximation error with the regularization and statistical convergence properties of empirical Sinkhorn divergence\. The resulting generalization bounds reveal how the local Sinkhorn divergence may partially alleviate the curse of dimensionality through introducing an entropic regularization term\.
- •Through various numerical examples, we demonstrate that the proposed local Sinkhorn divergence provides an effective compromise between approximation accuracy and computational efficiency and outperforms several machine\-learning\-based uncertainty quantification benchmarks, making it particularly suitable for multidimensional uncertainty quantification problems\.

The remainder of this paper is organized as follows\. Section[2](https://arxiv.org/html/2608.11613#S2)introduces the proposed local Sinkhorn divergence framework together with its theoretical analysis and generalization error bounds\. Section[3](https://arxiv.org/html/2608.11613#S3)presents three numerical examples, including one\-dimensional conditional distribution reconstruction, stochastic Darcy flow, and stochastic FHN systems, where the proposed method is compared with several local distribution matching losses and existing machine\-learning benchmark approaches\. Finally, Section[4](https://arxiv.org/html/2608.11613#S4)concludes the paper with a summary of the main findings and discusses several directions for future research\. Frequently used notations throughout this paper are summarized in Table[1](https://arxiv.org/html/2608.11613#S1.T1)\.

Table 1:Frequently used notations throughout this paper\.NotationDescription𝒙∈𝒟\\bm\{x\}\\in\\mathcal\{D\}Input variable\.𝒟⊆ℝn\\mathcal\{D\}\\subseteq\\mathbb\{R\}^\{n\}Input domain\.NNTotal number of training samples\.ω\\omega\(ω^\\hat\{\\omega\}\)Random variable representing uncertainty\.𝒚⁡\(𝒙,ω\)∈ℝd\\bm\{y\}\(\\bm\{x\},\\omega\)\\in\\mathbb\{R\}^\{d\}Target random field\.𝒚^​\(𝒙,ω^\)∈ℝd\\hat\{\\bm\{y\}\}\(\\bm\{x\},\\hat\{\\omega\}\)\\in\\mathbb\{R\}^\{d\}Reconstructed random field\.μ𝒙\\mu\_\{\\bm\{x\}\}Conditional probability distribution of𝒚⁡\(𝒙,ω\)\\bm\{y\}\(\\bm\{x\},\\omega\)\.μ^𝒙\\hat\{\\mu\}\_\{\\bm\{x\}\}Conditional probability distribution of𝒚^​\(𝒙,ω^\)\\hat\{\\bm\{y\}\}\(\\bm\{x\},\\hat\{\\omega\}\)\.μ𝒙,δe\\mu\_\{\\bm\{x\},\\delta\}^\{\\text\{e\}\}Empiricalδ\\delta\-neighborhood measure associated with𝒚⁡\(𝒙,ω\)\\bm\{y\}\(\\bm\{x\},\\omega\)\(details given in Section[2](https://arxiv.org/html/2608.11613#S2)\)\.μ^𝒙,δe\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}^\{\\text\{e\}\}Empiricalδ\\delta\-neighborhood measure associated with𝒚^​\(𝒙,ω^\)\\hat\{\\bm\{y\}\}\(\\bm\{x\},\\hat\{\\omega\}\)\(details given in Section[2](https://arxiv.org/html/2608.11613#S2)\)\.δ\\deltaNeighborhood radius\.ν\\nuProbability measure on the input domain\.νe\\nu^\{\\text\{e\}\}Empirical measure induced by the observed inputs\.W22​\(μ,μ^\)W\_\{2\}^\{2\}\(\\mu,\\hat\{\\mu\}\)SquaredW2W\_\{2\}distance between the two probability measuresμ,μ^\\mu,\\hat\{\\mu\}\.Sε​\(μ,μ^\)S\_\{\\varepsilon\}\(\\mu,\\hat\{\\mu\}\)Debiased Sinkhorn divergence between the two probability measuresμ,μ^\\mu,\\hat\{\\mu\}\.ε\\varepsilonSinkhorn regularization parameter\.𝒲¯2\\overline\{\\mathcal\{W\}\}\_\{2\}Averaged squared Wasserstein discrepancy\.𝒮¯ε,δe\\overline\{\\mathcal\{S\}\}\_\{\\varepsilon,\\delta\}^\{\\text\{e\}\}Averaged empirical local Sinkhorn divergence\.
## 2A local Sinkhorn divergence method for conditional distribution reconstruction of random fields

In this section, we formally introduce the random field model we reconstruct and the local Sinkhorn divergence framework we propose\. Consider the reconstruction of conditional probability distributions from scattered observations\. Let

𝒚𝒙=f⁡\(𝒙,ω\),x∈𝒟⊂ℝn,\\bm\{y\}\_\{\\bm\{x\}\}=f\(\\bm\{x\},\\omega\),\\qquad x\\in\\mathcal\{D\}\\subset\\mathbb\{R\}^\{n\},\(2\.1\)be the unknown random field to be reconstructed, whereω∈Ω\\omega\\in\\Omegadenotes the underlying random variable andf:𝒟×Ω→ℝdf:\\mathcal\{D\}\\times\\Omega\\rightarrow\\mathbb\{R\}^\{d\}is an unknown stochastic mapping\. For each input location𝒙\\bm\{x\}, only realizations of the random responsey⁡\(𝒙,ω\)y\(\\bm\{x\},\\omega\)are available, while the conditional probability measure of𝒚⁡\(𝒙,ω\)\\bm\{y\}\(\\bm\{x\};\\omega\), denoted asμ𝒙\\mu\_\{\\bm\{x\}\}, is unknown\.

Our objective is to construct an approximation

𝒚^𝒙=f^​\(𝒙,ω^\),𝒙∈𝒟,\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}=\\hat\{f\}\(\\bm\{x\},\\hat\{\\omega\}\),\\qquad\\bm\{x\}\\in\\mathcal\{D\},\(2\.2\)wheref^\\hat\{f\}is another stochastic mapping, andω^∈Ω^\\hat\{\\omega\}\\in\\hat\{\\Omega\}is another random variable\. The goal is to match the distribution of𝒚⁡\(𝒙,ω\)\\bm\{y\}\(\\bm\{x\};\\omega\)using the distribution of𝒚^​\(𝒙,ω^\)\\hat\{\\bm\{y\}\}\(\\bm\{x\};\\hat\{\\omega\}\), denoted byμ^𝒙\\hat\{\\mu\}\_\{\\bm\{x\}\}for every𝒙∈𝒟\\bm\{x\}\\in\\mathcal\{D\}\.

We employ a stochastic neural network \(SNN\) similar to our previous work[xia2026local](https://arxiv.org/html/2608.11613#bib.bib31)as the approximate random field modelf^\\hat\{f\}in Eq\. \([2\.2](https://arxiv.org/html/2608.11613#S2.E2)\)\. The structure of the SNN is described in Fig\.[1](https://arxiv.org/html/2608.11613#S2.F1)\. Unlike deterministic neural networks, the SNN contains random parameters \(as the random variableω^\\hat\{\\omega\}in Eq\. \([2\.2](https://arxiv.org/html/2608.11613#S2.E2)\)\) sampled during forward propagation, enabling the model to generate multiple realizations for the same input location\. Furthermore, it is proved in\([xia2025new](https://arxiv.org/html/2608.11613#bib.bib32), Appendix H\)that this SNN can approximate any random field model under an averaged squaredW2W\_\{2\}metric, subject to nonrestrictive assumptions\.

![Refer to caption](https://arxiv.org/html/2608.11613v1/concept_new.png)Figure 1:Illustration of the SNN used for conditional distribution reconstruction\. Each linear layer samples its weights from Gaussian distributions whose means and variances are trainable during forward propagation, enabling stochastic realizations to be generated for the same input\. Residual connections \(ResNet\) could be employed to improve optimization stability\. The generated samples are used to construct empirical conditional distributions of the reconstructed𝒚^​\(𝒙,ω^\)\\hat\{\\bm\{y\}\}\(\\bm\{x\};\\hat\{\\omega\}\)in Eq\. \([2\.2](https://arxiv.org/html/2608.11613#S2.E2)\)\. The stochastic weights in the SNN serve as the random variableω^\\hat\{\\omega\}in Eq\. \([2\.2](https://arxiv.org/html/2608.11613#S2.E2)\)\. TheReLUactivation function may be replaced with other activation functions\.First, we briefly review the Sinkhorn divergence between two probability measures\. Let𝒚,𝒚^∈ℝn\\bm\{y\},\\hat\{\\bm\{y\}\}\\in\\mathbb\{R\}^\{n\}be two random variables satisfying:

𝔼⁡\[‖𝒚‖2\]<∞,𝔼⁡\[‖𝒚^‖2\]<∞\.\\mathbb\{E\}\\\!\\left\[\\\|\\bm\{y\}\\\|^\{2\}\\right\]<\\infty,\\qquad\\mathbb\{E\}\\\!\\left\[\\\|\\hat\{\\bm\{y\}\}\\\|^\{2\}\\right\]<\\infty\.\(2\.3\)Throughout this paper,∥⋅∥\\\|\\cdot\\\|denotes theℓ2\\ell^\{2\}norm of a vector\. The probability distributions of𝒚\\bm\{y\}and𝒚^\\hat\{\\bm\{y\}\}are denoted byμ\\muandμ^\\hat\{\\mu\}, respectively\. Let

c⁡\(y,y^\)=‖𝒚−𝒚^‖2c\(y,\\hat\{y\}\)=\\\|\\bm\{y\}\-\\hat\{\\bm\{y\}\}\\\|^\{2\}\(2\.4\)be the quadratic transportation cost\. The entropically regularized optimal transport cost is defined as:

Wε2\(μ,μ^\)=minπ∈Π⁡\(μ,μ^\)\{∫ℝn×ℝnc\(𝒚,𝒚^\)dπ\(𝒚,𝒚^\)\+εKL\(π∥μ⊗μ^\)\},W\_\{\\varepsilon\}^\{2\}\(\\mu,\\hat\{\\mu\}\)=\\min\_\{\\pi\\in\\Pi\(\\mu,\\hat\{\\mu\}\)\}\\left\\\{\\int\_\{\\mathbb\{R\}^\{n\}\\times\\mathbb\{R\}^\{n\}\}c\(\\bm\{y\},\\hat\{\\bm\{y\}\}\)\\,\\mbox\{d\}\\pi\(\\bm\{y\},\\hat\{\\bm\{y\}\}\)\+\\varepsilon\\,\\mathrm\{KL\}\(\\pi\\,\\\|\\,\\mu\\otimes\\hat\{\\mu\}\)\\right\\\},\(2\.5\)whereΠ⁡\(μ,μ^\)\\Pi\(\\mu,\\hat\{\\mu\}\)denotes the set of probability measures onℝn×ℝn\\mathbb\{R\}^\{n\}\\times\\mathbb\{R\}^\{n\}whose marginals areμ\\muandμ^\\hat\{\\mu\}, respectively\. In Eq\. \([2\.5](https://arxiv.org/html/2608.11613#S2.E5)\),

KL\(π∥μ⊗μ^\)=∫ℝn×ℝnlog\(d​πd​\(μ⊗μ^\)\)dπ\\mathrm\{KL\}\(\\pi\\,\\\|\\,\\mu\\otimes\\hat\{\\mu\}\)=\\int\_\{\\mathbb\{R\}^\{n\}\\times\\mathbb\{R\}^\{n\}\}\\log\\\!\\left\(\\frac\{\\mbox\{d\}\\pi\}\{\\mbox\{d\}\(\\mu\\otimes\\hat\{\\mu\}\)\}\\right\)\\,\\mbox\{d\}\\pidenotes the Kullback–Leibler divergence ofπ\\piwith respect to the product measureμ⊗μ^\\mu\\otimes\\hat\{\\mu\}, provided thatπ≪μ⊗μ^\\pi\\ll\\mu\\otimes\\hat\{\\mu\}; otherwise,KL\(π∥μ⊗μ^\)=\+∞\\mathrm\{KL\}\(\\pi\\,\\\|\\,\\mu\\otimes\\hat\{\\mu\}\)=\+\\infty, andε\>0\\varepsilon\>0is the entropic regularization parameter\.

Unlike the classical Wasserstein distance, the regularized transport cost does not vanish whenμ=μ^\\mu=\\hat\{\\mu\}\. To remove this entropic bias, the Sinkhorn divergence is defined as

Sε​\(μ,μ^\)=Wε2​\(μ,μ^\)−12​Wε2​\(μ,μ\)−12​Wε2​\(μ^,μ^\)\.S\_\{\\varepsilon\}\(\\mu,\\hat\{\\mu\}\)=W\_\{\\varepsilon\}^\{2\}\(\\mu,\\hat\{\\mu\}\)\-\\frac\{1\}\{2\}W\_\{\\varepsilon\}^\{2\}\(\\mu,\\mu\)\-\\frac\{1\}\{2\}W\_\{\\varepsilon\}^\{2\}\(\\hat\{\\mu\},\\hat\{\\mu\}\)\.\(2\.6\)The Sinkhorn divergence in Eq\. \([2\.6](https://arxiv.org/html/2608.11613#S2.E6)\) is symmetric, non\-negative, differentiable with respect to the empirical samples, and satisfiesSε​\(μ,μ\)=0S\_\{\\varepsilon\}\(\\mu,\\mu\)=0; furthermore, the Sinkhorn divergence converges to the Wasserstein distance asε→0\\varepsilon\\rightarrow 0while remaining differentiable and computationally tractable, whereas choosingε→∞\\varepsilon\\rightarrow\\inftyrecovers the maximum mean discrepancy[feydy2019interpolating](https://arxiv.org/html/2608.11613#bib.bib8)\. Yet, the Sinkhorn divergence is more computationally efficient through introducing an additional regularizing term in Eq\. \([2\.5](https://arxiv.org/html/2608.11613#S2.E5)\) while preserving the advantage of the squaredW2W\_\{2\}distance for comparing two probability measures\.

To simplify the subsequent analysis, we introduce the following assumptions\.

###### Assumption 2\.1\.

Assume that the uncertainty model𝒚𝒙\\bm\{y\}\_\{\\bm\{x\}\}and the approximation model𝒚^𝒙\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}in Eqs\. \([2\.1](https://arxiv.org/html/2608.11613#S2.E1)\) and \([2\.2](https://arxiv.org/html/2608.11613#S2.E2)\) satisfy the following conditions\.

1. 1\.Both random fields are uniformly bounded,i\.e\., there existsM0\>0M\_\{0\}\>0such that ‖𝒚⁡\(𝒙,ω\)‖≤M0,‖𝒚^​\(𝒙,ω^\)‖≤M0,∀𝒙∈𝒟\.\\\|\\bm\{y\}\(\\bm\{x\},\\omega\)\\\|\\leq M\_\{0\},\\qquad\\\|\\hat\{\\bm\{y\}\}\(\\bm\{x\},\\hat\{\\omega\}\)\\\|\\leq M\_\{0\},\\qquad\\forall\\bm\{x\}\\in\\mathcal\{D\}\.\(2\.7\)
2. 2\.The mappingsffandf^\\hat\{f\}are uniformly Lipschitz continuous with respect to the input variable\. That is, there exists a constantL\>0L\>0such that ∥f\(𝒙,ω\)−f\(𝒙′,ω\)∥≤L∥𝒙−𝒙′∥,∥f^\(𝒙,ω^\)−f^\(𝒙′,ω^\)∥≤L∥𝒙−𝒙′∥,∀𝒙,𝒙′∈𝒟\.\\\|f\(\\bm\{x\},\\omega\)\-f\(\\bm\{x\}^\{\\prime\},\\omega\)\\\|\\leq L\\\|\\bm\{x\}\-\\bm\{x\}^\{\\prime\}\\\|,\\,\\,\\,\\,\\\|\\hat\{f\}\(\\bm\{x\},\\hat\{\\omega\}\)\-\\hat\{f\}\(\\bm\{x\}^\{\\prime\},\\hat\{\\omega\}\)\\\|\\leq L\\\|\\bm\{x\}\-\\bm\{x\}^\{\\prime\}\\\|,\\,\\,\\forall\\bm\{x\},\\bm\{x\}^\{\\prime\}\\in\\mathcal\{D\}\.\(2\.8\)
3. 3\.The random variablesω\\omegaandω^\\hat\{\\omega\}are independent of the input variable𝒙\\bm\{x\}\.

Following our previous work on local optimal transport[xia2026local](https://arxiv.org/html/2608.11613#bib.bib31), we define the weighted average squaredW2W\_\{2\}metric between the ground\-truth and approximate random fields𝒚𝒙\\bm\{y\}\_\{\\bm\{x\}\}and𝒚^𝒙\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}in Eqs\. \([2\.1](https://arxiv.org/html/2608.11613#S2.E1)\) and \([2\.2](https://arxiv.org/html/2608.11613#S2.E2)\):

𝒲¯2​\(𝒚𝒙,𝒚^𝒙\)≔∫𝒟W22​\(μ𝒙,μ^𝒙\)​ν​\(d​𝒙\)\.\\overline\{\\mathcal\{W\}\}\_\{2\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\\coloneqq\\int\_\{\\mathcal\{D\}\}W\_\{2\}^\{2\}\(\\mu\_\{\\bm\{x\}\},\\hat\{\\mu\}\_\{\\bm\{x\}\}\)\\,\\nu\(\\mbox\{d\}\\bm\{x\}\)\.\(2\.9\)
𝒲¯2​\(𝒚𝒙,𝒚^𝒙\)\\overline\{\\mathcal\{W\}\}\_\{2\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)defined in Eq\. \([2\.9](https://arxiv.org/html/2608.11613#S2.E9)\) is nonnegative, and if the weighting measure of the input variableν\\nuis strictly positive over𝒟\\mathcal\{D\}, then𝒲¯2​\(𝒚𝒙,𝒚^𝒙\)=0\\overline\{\\mathcal\{W\}\}\_\{2\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)=0holds if and only if

W22​\(μ𝒙,μ^𝒙\)=0,for almost every​𝒙∈𝒟,W\_\{2\}^\{2\}\(\\mu\_\{\\bm\{x\}\},\\hat\{\\mu\}\_\{\\bm\{x\}\}\)=0,\\qquad\\text\{for almost every \}\\bm\{x\}\\in\\mathcal\{D\},which implies that the probability measure of𝒚𝒙\\bm\{y\}\_\{\\bm\{x\}\}can be matched by the probability measure of𝒚^𝒙\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}for almost every𝒙∈𝒟\\bm\{x\}\\in\\mathcal\{D\}\. Consequently, minimizing \([2\.9](https://arxiv.org/html/2608.11613#S2.E9)\) is equivalent to matching the conditional distributionsμ𝒙\\mu\_\{\\bm\{x\}\}usingμ^𝒙\\hat\{\\mu\}\_\{\\bm\{x\}\}throughout the computational domain\. However, in practical applications only a finite set of scattered observations\{\(𝒙i,𝒚𝒙i\)\}i=1N\\\{\(\\bm\{x\}\_\{i\},\\bm\{y\}\_\{\\bm\{x\}\_\{i\}\}\)\\\}\_\{i=1\}^\{N\}is available\. Since the conditional probability measuresμ𝒙\\mu\_\{\\bm\{x\}\}andμ^𝒙\\hat\{\\mu\}\_\{\\bm\{x\}\}cannot be evaluated directly from finite observations, we approximate them by local empirical measures constructed from neighboring samples\.

Given a training sample𝒙i∈\{𝒙1,…​𝒙N\}\\bm\{x\}\_\{i\}\\in\\\{\\bm\{x\}\_\{1\},\.\.\.\\bm\{x\}\_\{N\}\\\}, we define its neighborhood

B\(𝒙i,δ\)=\{𝒙j:∥𝒙j−𝒙i∥<δ,j=1,…,N\},B\(\\bm\{x\}\_\{i\},\\delta\)=\\\{\\bm\{x\}\_\{j\}:\\\|\\bm\{x\}\_\{j\}\-\\bm\{x\}\_\{i\}\\\|<\\delta,\\,\\,j=1,\.\.\.,N\\\},\(2\.10\)whereδ\>0\\delta\>0is the neighborhood radius\. For implementing the neighborhood technique, we randomly choose a subset

ℬ⁡\(𝒙i,δ\)⊆B⁡\(𝒙i,δ\)\.\\mathcal\{B\}\(\\bm\{x\}\_\{i\},\\delta\)\\subseteq B\(\\bm\{x\}\_\{i\},\\delta\)\.\(2\.11\)Letμxi,δe\\mu\_\{x\_\{i\},\\delta\}^\{\\text\{e\}\}andμ^𝒙i,δe\\hat\{\\mu\}\_\{\\bm\{x\}\_\{i\},\\delta\}^\{\\text\{e\}\}denote the corresponding empirical conditional probability measures constructed from the neighboring observations,i\.e\.,μ𝒙i,δe\\mu\_\{\\bm\{x\}\_\{i\},\\delta\}^\{\\text\{e\}\}is the empirical distribution of𝒚⁡\(𝒙,ω\),𝒙∈ℬ⁡\(𝒙i,δ\)\\bm\{y\}\(\\bm\{x\};\\omega\),\\bm\{x\}\\in\\mathcal\{B\}\(\\bm\{x\}\_\{i\},\\delta\)andμ^𝒙i,δe\\hat\{\\mu\}\_\{\\bm\{x\}\_\{i\},\\delta\}^\{\\text\{e\}\}is the empirical distribution of𝒚^​\(𝒙,ω\),𝒙∈ℬ⁡\(𝒙i,δ\)\\hat\{\\bm\{y\}\}\(\\bm\{x\};\\omega\),\\bm\{x\}\\in\\mathcal\{B\}\(\\bm\{x\}\_\{i\},\\delta\)\.μ𝒙i,δe\\mu\_\{\\bm\{x\}\_\{i\},\\delta\}^\{\\text\{e\}\}andμ^𝒙i,δe\\hat\{\\mu\}\_\{\\bm\{x\}\_\{i\},\\delta\}^\{\\text\{e\}\}are used to approximateμ𝒙\\mu\_\{\\bm\{x\}\}andμ^𝒙\\hat\{\\mu\}\_\{\\bm\{x\}\}, respectively\. We then introduce the proposed local Sinkhorn divergence loss:

Sε,δe​\(𝒚𝒙,𝒚^𝒙\)≔∫𝒟Sε​\(μ𝒙,δe,μ^𝒙,δe\)​νe​\(d​𝒙\),S\_\{\\varepsilon,\\delta\}^\{\\text\{e\}\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\\coloneqq\\int\_\{\\mathcal\{D\}\}S\_\{\\varepsilon\}\(\\mu\_\{\\bm\{x\},\\delta\}^\{\\text\{e\}\},\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}^\{\\text\{e\}\}\)\\,\\nu^\{\\text\{e\}\}\(\\mbox\{d\}\\bm\{x\}\),\(2\.12\)whereνe\\nu^\{\\text\{e\}\}denotes the empirical measure associated with the observed input variable𝒙\\bm\{x\}\. Eq\. \([2\.12](https://arxiv.org/html/2608.11613#S2.E12)\) is not a single Sinkhorn divergence between two probability measures\. Instead, it computes the average Sinkhorn divergence between the local empirical conditional distributions over the entire input domain\. The superscripteindicates empirical probability measures estimated from neighboring observations, while the subscriptδ\\deltadenotes the neighborhood radius\.

Compared with theW2W\_\{2\}metrics, the Sinkhorn divergence introduces an entropic regularization, allowing the optimal transport problem to be solved efficiently by iterative matrix scaling rather than network simplex optimization[cuturi2013sinkhorn](https://arxiv.org/html/2608.11613#bib.bib6);[feydy2019interpolating](https://arxiv.org/html/2608.11613#bib.bib8)\. Consequently, the resulting loss is fully differentiable and naturally amenable to GPU\-parallel computation, while recent theoretical results also demonstrate improved statistical convergence properties for empirical Sinkhorn divergences compared with empirical Wasserstein distances[genevay2019sample](https://arxiv.org/html/2608.11613#bib.bib11)\. Specifically, we establish two different forms of the generalization error bounds on how the local Sinkhorn divergence Eq\. \([2\.12](https://arxiv.org/html/2608.11613#S2.E12)\) approximates the averaged squared Wasserstein discrepancy in Eq\. \([2\.9](https://arxiv.org/html/2608.11613#S2.E9)\)\.

###### Theorem 2\.1\(Generalization bound based on the generalization error of the local squaredW2W\_\{2\}loss\)\.

Under Assumptions[2\.1](https://arxiv.org/html/2608.11613#S2.Thmassumption1), letμ𝒙\\mu\_\{\\bm\{x\}\}andμ^𝒙\\hat\{\\mu\}\_\{\\bm\{x\}\}denote the true conditional probability measures associated with𝒚𝒙\\bm\{y\}\_\{\\bm\{x\}\}and𝒚^𝒙\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}at location𝒙∈𝒟\\bm\{x\}\\in\\mathcal\{D\}, respectively\. Letμ𝒙,δe\\mu\_\{\\bm\{x\},\\delta\}^\{\\text\{e\}\}andμ^𝒙,δe\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}^\{\\text\{e\}\}be the corresponding empirical neighborhood measures built from the sample set\{\(𝒙i,𝒚𝒙i\)\}i=1N\\\{\(\\bm\{x\}\_\{i\},\\bm\{y\}\_\{\\bm\{x\}\_\{i\}\}\)\\\}\_\{i=1\}^\{N\}with neighborhood radiusδ\>0\\delta\>0\. We assumeε<2​M02d\\varepsilon<\\frac\{2M\_\{0\}^\{2\}\}\{\\sqrt\{d\}\}\. For𝒮¯ε,δe​\(𝒚𝒙,𝒚^𝒙\)\\overline\{\\mathcal\{S\}\}\_\{\\varepsilon,\\delta\}^\{\\text\{e\}\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)and𝒲¯2​\(𝒚𝒙,𝒚^𝒙\)\\overline\{\\mathcal\{W\}\}\_\{2\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)defined in Eqs\. \([2\.9](https://arxiv.org/html/2608.11613#S2.E9)\) and \([2\.12](https://arxiv.org/html/2608.11613#S2.E12)\), there exists a constantC\>0C\>0such that:

𝔼⁡\[\|𝒮¯ε,δe​\(𝒚𝒙,𝒚^𝒙\)−𝒲¯2​\(yx,y^x\)\|\]≤2​ε​log⁡\(2​e2​M02d​ε\)\+4​M0N\+8​C​M0​𝔼​\[h⁡\(N⁡\(x,δ\),d\)\]\+16​M02​δ,\\displaystyle\\mathbb\{E\}\\\!\\left\[\\left\|\\overline\{\\mathcal\{S\}\}\_\{\\varepsilon,\\delta\}^\{\\text\{e\}\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\-\\overline\{\\mathcal\{W\}\}\_\{2\}\(y\_\{x\},\\hat\{y\}\_\{x\}\)\\right\|\\right\]\\leq 2\\varepsilon\\log\\\!\\left\(\\frac\{2e^\{2\}M\_\{0\}^\{2\}\}\{\\sqrt\{d\}\\,\\varepsilon\}\\right\)\+\\frac\{4M\_\{0\}\}\{\\sqrt\{N\}\}\+8CM\_\{0\}\\,\\mathbb\{E\}\\\!\\bigl\[h\(N\(x,\\delta\),d\)\\bigr\]\+16M\_\{0\}^\{2\}\\,\\delta,\(2\.13\)whereNNis the number of training samples, and

N⁡\(𝒙,δ\)≔\|ℬ⁡\(𝒙,δ\)\|N\(\\bm\{x\},\\delta\)\\coloneqq\|\\mathcal\{B\}\(\\bm\{x\},\\delta\)\|\(2\.14\)denotes the effective number of samples chosen from the neighborhoodB⁡\(𝒙,δ\)B\(\\bm\{x\},\\delta\)centered at𝒙\\bm\{x\}defined in Eq\. \([2\.10](https://arxiv.org/html/2608.11613#S2.E10)\), and

h⁡\(N,d\):=\{2N−1/4log\(1\+N\)1/2,d≤4,2​N−14\+\(∏i=1d\(σi,𝒙σ1,𝒙\)​N\)−1d\+\(∏i=1d\(σ^i,𝒙σ^1,𝒙\)​N\)−1d,d\>4\.h\(N,d\):=\\begin\{cases\}2N^\{\-1/4\}\\log\(1\+N\)^\{1/2\},&d\\leq 4,\\\\\[5\.69054pt\] 2N^\{\-\\frac\{1\}\{4\}\}\+\\big\(\\prod\_\{i=1\}^\{d\}\(\\frac\{\\sigma\_\{i,\\bm\{x\}\}\}\{\\sigma\_\{1,\\bm\{x\}\}\}\)N\\big\)^\{\-\\frac\{1\}\{d\}\}\+\\big\(\\prod\_\{i=1\}^\{d\}\(\\frac\{\\hat\{\\sigma\}\_\{i,\\bm\{x\}\}\}\{\\hat\{\\sigma\}\_\{1,\\bm\{x\}\}\}\)N\\big\)^\{\-\\frac\{1\}\{d\}\},&d\>4\.\\end\{cases\}\(2\.15\)In Eq\. \([2\.15](https://arxiv.org/html/2608.11613#S2.E15)\),

σi,𝒙≔\(∫ℝdyi​μ𝒙​\(d​𝒚\)\)16,σ^i,𝒙≔\(∫ℝdyi​μ^𝒙​\(d​𝒚\)\)16\.\\sigma\_\{i,\\bm\{x\}\}\\coloneqq\\big\(\\int\_\{\\mathbb\{R\}^\{d\}\}y\_\{i\}\\mu\_\{\\bm\{x\}\}\(\\mbox\{d\}\\bm\{y\}\)\\big\)^\{\\frac\{1\}\{6\}\},\\,\\,\\hat\{\\sigma\}\_\{i,\\bm\{x\}\}\\coloneqq\\big\(\\int\_\{\\mathbb\{R\}^\{d\}\}y\_\{i\}\\hat\{\\mu\}\_\{\\bm\{x\}\}\(\\mbox\{d\}\\bm\{y\}\)\\big\)^\{\\frac\{1\}\{6\}\}\.\(2\.16\)

The proof of Theorem[2\.1](https://arxiv.org/html/2608.11613#S2.Thmtheorem1)is mainly based on the generalization error bound of the local squaredW2W\_\{2\}distance \(\([xia2025efficient](https://arxiv.org/html/2608.11613#bib.bib30), Theorem 2\.3\)\) and is given in Appendix[B](https://arxiv.org/html/2608.11613#A2)\. Theorem[2\.1](https://arxiv.org/html/2608.11613#S2.Thmtheorem1)provides a generalization error bound on estimating how the empirical Sinkhorn divergence loss functionS¯ε,δe\\bar\{S\}\_\{\\varepsilon,\\delta\}^\{\\text\{e\}\}, when utilizing the neighborhood technique, approximates the averaged squaredW2W\_\{2\}distance between the ground truth random field𝒚𝒙\\bm\{y\}\_\{\\bm\{x\}\}and the approximate model𝒚^𝒙\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\.

On the other hand, if we use the generalization error bound for estimating the Sinkhorn divergence using the empirical distributions, we can obtain another generalization error bound for our local Sinkhorn divergence loss function\.

###### Theorem 2\.2\(Generalization bound based on the generalization error of the Sinkhorn divergence\)\.

For the averaged local Sinkhorn loss defined in Eq\. \([2\.12](https://arxiv.org/html/2608.11613#S2.E12)\) and the corresponding population averaged squaredW2W\_\{2\}distance defined in Eq\. \([2\.9](https://arxiv.org/html/2608.11613#S2.E9)\), assume that the hypotheses of\([genevay2019sample](https://arxiv.org/html/2608.11613#bib.bib11), Theorem 3\)hold, andε<2​M02d\\varepsilon<\\frac\{2M\_\{0\}^\{2\}\}\{\\sqrt\{d\}\}\. Then the following generalization error bound holds:

𝔼⁡\[\|𝒮¯ε,δe​\(𝒚𝒙,𝒚^𝒙\)−𝒲¯2​\(yx,y^x\)\|\]≤2​ε​log⁡\(2​e2​M02d​ε\)\+𝒪⁡\(e5​M02εOPENN⁡\(𝒙,δCLOSE\)​\(1\+ε−⌊d/2⌋\)\)\+16​M02​δ,\\displaystyle\\mathbb\{E\}\\\!\\left\[\\left\|\\overline\{\\mathcal\{S\}\}\_\{\\varepsilon,\\delta\}^\{\\text\{e\}\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\-\\overline\{\\mathcal\{W\}\}\_\{2\}\(y\_\{x\},\\hat\{y\}\_\{x\}\)\\right\|\\right\]\\leq 2\\varepsilon\\log\\\!\\Big\(\\frac\{2e^\{2\}M\_\{0\}^\{2\}\}\{\\sqrt\{d\}\\,\\varepsilon\}\\Big\)\+\\mathcal\{O\}\\bigg\(\\frac\{e^\{\\frac\{5M\_\{0\}^\{2\}\}\{\\varepsilon\}\}\}\{\\sqrt\{N\(\\bm\{x\},\\delta\}\)\}\\Big\(1\+\\varepsilon^\{\-\\lfloor d/2\\rfloor\}\\Big\)\\bigg\)\+16M\_\{0\}^\{2\}\\delta,\(2\.17\)whereN⁡\(𝒙,δ\)N\(\\bm\{x\},\\delta\)denotes the effective number of samples in each neighborhood \(defined in Eq\. \([2\.14](https://arxiv.org/html/2608.11613#S2.E14)\)\)\.

The proof of Theorem[2\.2](https://arxiv.org/html/2608.11613#S2.Thmtheorem2)is mainly based on the generalization error bound of using the empirical measures to evaluate the Sinkhorn divergence \(\([genevay2019sample](https://arxiv.org/html/2608.11613#bib.bib11), Theorem 3\)\) and is given in Appendix[B](https://arxiv.org/html/2608.11613#A2)\.

To exemplify how the size of the neighborhoodδ\\deltaand the number of samples in each neighborhoodN⁡\(𝒙,δ\)N\(\\bm\{x\},\\delta\)influences the generalization error bound, consider the case when the empirical probability measure of the input variable is close to the uniform distribution and the neighborhood size satisfiesN⁡\(𝒙,δ\)≍C​N​δnN\(\\bm\{x\},\\delta\)\\asymp CN\\delta^\{n\}\(NNis the total number of training samples\) for the intrinsic dimensionnnof the input variable𝒙\\bm\{x\}, then using Theorem[2\.1](https://arxiv.org/html/2608.11613#S2.Thmtheorem1), we have:

𝔼⁡\[\|𝒮¯ε,δe​\(𝒚𝒙,𝒚^𝒙\)−𝒲¯2​\(𝒚𝒙,𝒚^𝒙\)\|\]≲2​ε​log⁡\(2​e2​M02d​ε\)\+4​M0N\+8​C​M0​h​\(C​N​δn,d\)\+16​M02​δ\.\\displaystyle\\mathbb\{E\}\\\!\\left\[\\left\|\\overline\{\\mathcal\{S\}\}\_\{\\varepsilon,\\delta\}^\{\\text\{e\}\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\-\\overline\{\\mathcal\{W\}\}\_\{2\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\\right\|\\right\]\\lesssim 2\\varepsilon\\log\\\!\\left\(\\frac\{2e^\{2\}M\_\{0\}^\{2\}\}\{\\sqrt\{d\}\\,\\varepsilon\}\\right\)\+\\frac\{4M\_\{0\}\}\{\\sqrt\{N\}\}\+8CM\_\{0\}\\,h\(CN\\delta^\{n\},d\)\+16M\_\{0\}^\{2\}\\,\\delta\.\(2\.18\)
On the other hand, using Theorem[2\.2](https://arxiv.org/html/2608.11613#S2.Thmtheorem2)yields:

𝔼⁡\[\|𝒮¯ε,δe​\(𝒚𝒙,𝒚^𝒙\)−𝒲¯2​\(𝒚𝒙,𝒚^𝒙\)\|\]≲2​ε​log⁡\(2​e2​M02d​ε\)\+𝒪⁡\(e5​M02εC​N​δn​\(1\+ε−⌊d/2⌋\)\)\+16​M02​δ\.\\displaystyle\\mathbb\{E\}\\\!\\left\[\\left\|\\overline\{\\mathcal\{S\}\}\_\{\\varepsilon,\\delta\}^\{\\text\{e\}\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\-\\overline\{\\mathcal\{W\}\}\_\{2\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\\right\|\\right\]\\lesssim 2\\varepsilon\\log\\\!\\left\(\\frac\{2e^\{2\}M\_\{0\}^\{2\}\}\{\\sqrt\{d\}\\,\\varepsilon\}\\right\)\+\\mathcal\{O\}\\bigg\(\\frac\{e^\{\\frac\{5M\_\{0\}^\{2\}\}\{\\varepsilon\}\}\}\{\\sqrt\{CN\\delta^\{n\}\}\}\\Bigl\(1\+\\varepsilon^\{\-\\lfloor d/2\\rfloor\}\\Bigr\)\\bigg\)\+16M\_\{0\}^\{2\}\\delta\.\(2\.19\)
The two bounds Eqs\. \([2\.18](https://arxiv.org/html/2608.11613#S2.E18)\) and \([2\.19](https://arxiv.org/html/2608.11613#S2.E19)\) suggest an important practical implication\. When the data are highly heterogeneous, the conditional distributions may vary sharply across the domain\. When the approximate random field model can capture such heterogeneity such that∏i=1d\(σi,𝒙σ1,𝒙\)1d\+∏i=1d\(σ^i,𝒙σ^1,𝒙\)1d\\prod\_\{i=1\}^\{d\}\(\\frac\{\\sigma\_\{i,\\bm\{x\}\}\}\{\\sigma\_\{1,\\bm\{x\}\}\}\)^\{\\frac\{1\}\{d\}\}\+\\prod\_\{i=1\}^\{d\}\(\\frac\{\\hat\{\\sigma\}\_\{i,\\bm\{x\}\}\}\{\\hat\{\\sigma\}\_\{1,\\bm\{x\}\}\}\)^\{\\frac\{1\}\{d\}\}is small, the approximation error induced by overly large entropic regularization whenε\\varepsilonis large can dominate the generalization error\. In this regime, it is preferable to choose a relatively small regularization parameterε\\varepsilon, so that the local Sinkhorn loss remains close to the local squaredW2W\_\{2\}loss\. Then, compared with the local squaredW2W\_\{2\}loss, the additional generalization error induced is essentially only the Sinkhorn regularization bias, while the neighborhood\-based approximation and sampling errors remain unchanged\.

On the other hand, when the noise in the target variable𝒚𝒙\\bm\{y\}\_\{\\bm\{x\}\}is homogeneous, the conditional distributions vary more smoothly and a moderately large regularization parameter, e\.g\.ε=O⁡\(1\)\\varepsilon=O\(1\)at the scale of the cost function, can be advantageous, so that the convergence rate as the number of training samplesN→∞N\\rightarrow\\inftyis𝒪⁡\(\(N​δn\)−12\)\\mathcal\{O\}\\big\(\(N\\delta^\{n\}\)^\{\-\\frac\{1\}\{2\}\}\\big\)\. In this case, entropic smoothing improves numerical stability and can reduce the effective sample complexity penalty associated with the Sinkhorn approximation\. As a result, by tuningε\\varepsilon, the local Sinkhorn loss provides a flexible trade\-off between statistical accuracy and computational robustness, and may partially alleviate the curse of dimensionality compared with the exact local squaredW2W\_\{2\}loss\. Furthermore, asε\\varepsilondecreases, the proposed local Sinkhorn loss approaches the local squared Wasserstein distance, leading to higher geometric fidelity but slower statistical convergence\. Conversely, increasingε\\varepsilonimproves the statistical convergence and computational efficiency of the empirical Sinkhorn divergence while introducing additional regularization bias\. Therefore, selecting a moderate regularization parameter also provides a favorable compromise between approximation accuracy and computational efficiency\.

Finally, the neighborhood radiusδ\\deltamust be chosen to balance two competing sources of error\. As shown in Eqs\. \([2\.18](https://arxiv.org/html/2608.11613#S2.E18)\) and \([2\.19](https://arxiv.org/html/2608.11613#S2.E19)\), increasingδ\\deltaenlarges the number of samples within each neighborhood and therefore reduces the empirical estimation error, namely8​C​M0​h​\(C​N​δn,d\)8CM\_\{0\}\\,h\(CN\\delta^\{n\},d\)and𝒪⁡\(exp⁡\(5​M02/ε\)C​N​δn​\(1\+ε−⌊d/2⌋\)\)\\mathcal\{O\}\\\!\\left\(\\frac\{\\exp\\\!\\left\(5M\_\{0\}^\{2\}/\\varepsilon\\right\)\}\{\\sqrt\{CN\\delta^\{n\}\}\}\\left\(1\+\\varepsilon^\{\-\\lfloor d/2\\rfloor\}\\right\)\\right\)\. However, a larger neighborhood inevitably increases the systematic approximation error, quantified by the term16​M02​δ16M\_\{0\}^\{2\}\\delta\. Consequently, an appropriate neighborhood radius should strike a balance between statistical accuracy and locality: it should not be too small to provide adequate samples for reliable estimation of the local empirical distribution, while remaining not too large to preserve the local structure of the underlying conditional distribution and avoid excessive approximation bias\.

## 3Numerical experiments

In this section, we carry out numerical experiments to test our proposed local Sinkhorn approach for training SNNs\. Test settings and hyperparameters are listed in Table[4](https://arxiv.org/html/2608.11613#A3.T4)\. Pseudocode for minimizing our proposed local Sinkhorn to train the SNN in Fig\.[1](https://arxiv.org/html/2608.11613#S2.F1)is given in Algorithm[1](https://arxiv.org/html/2608.11613#alg1)\.

Algorithm 1The pseudocode of the local Sinkhorn approach to train an SNN1:Given

NNobserved data

\{\(𝒙i,𝒚𝒙i\)\}i=1N\\\{\(\\bm\{x\}\_\{i\},\\bm\{y\}\_\{\\bm\{x\}\_\{i\}\}\)\\\}\_\{i=1\}^\{N\}, the neighborhood radius

δ\\delta, the mini\-batch size

nbn\_\{b\}, the number of local samples

nn, the minimum number of local samples required for each center

nminn\_\{\\min\}, the maximal number of local samples recorded for each center

nmaxn\_\{\\max\}, the number of epochs

EE, and the Sinkhorn regularization parameter

ε\\varepsilon
2:Initialize the SNN in Fig\.[1](https://arxiv.org/html/2608.11613#S2.F1)

3:For each

xix\_\{i\}, find its neighborhood

Bi≔\{𝒙j:‖𝒙j−𝒙i‖≤δ\}B\_\{i\}\\coloneqq\\\{\\bm\{x\}\_\{j\}:\\\|\\bm\{x\}\_\{j\}\-\\bm\{x\}\_\{i\}\\\|\\leq\\delta\\\}\. Do not use

𝒙i\\bm\{x\}\_\{i\}as an eligible training neighborhood center if its neighborhood contains too few samples \(

\|Bi\|<nmin\|B\_\{i\}\|<n\_\{\\min\}\)\. If there are too many neighbors such that

\|Bi\|\>nmax\|B\_\{i\}\|\>n\_\{\\max\}, randomly keep

nmaxn\_\{\\max\}neighbors to form the updated

BiB\_\{i\}
4:Input

\{𝒙i\}i=1N\\\{\\bm\{x\}\_\{i\}\\\}\_\{i=1\}^\{N\}into the SNN model to obtain predictions

\{𝒚^i\}i=1N\\\{\\hat\{\\bm\{y\}\}\_\{i\}\\\}\_\{i=1\}^\{N\}
5:for

epoch=1,…,E\\mathrm\{epoch\}=1,\\ldots,Edo

6:Randomly select

nbn\_\{b\}neighborhood centers

7:foreach selected center

𝒙i\\bm\{x\}\_\{i\}do

8:Randomly choose

min⁡\(n,\|Bi\|\)\\min\(n,\|B\_\{i\}\|\)samples from

BiB\_\{i\}to construct a local mini\-batch \(

ℬ⁡\(𝒙i,δ\)\\mathcal\{B\}\(\\bm\{x\}\_\{i\},\\delta\)defined in Eq\. \([2\.11](https://arxiv.org/html/2608.11613#S2.E11)\)\)

9:Input the sampled local inputs into the SNN to obtain predictions

10:Compute the local Sinkhorn loss

ℒi=Sε​\(μ^𝒙i,δe,μ𝒙i,δe\)\\mathcal\{L\}\_\{i\}=S\_\{\\varepsilon\}\\\!\\left\(\\hat\{\\mu\}\_\{\\bm\{x\}\_\{i\},\\delta\}^\{\\text\{e\}\},\\mu\_\{\\bm\{x\}\_\{i\},\\delta\}^\{\\text\{e\}\}\\right\)
11:endfor

12:Calculate the total loss

ℒ=1nb​∑i=1nbℒi\\mathcal\{L\}=\\frac\{1\}\{n\_\{b\}\}\\sum\_\{i=1\}^\{n\_\{b\}\}\\mathcal\{L\}\_\{i\}
13:Perform gradient descent to minimize

ℒ\\mathcal\{L\}and update the parameters in the SNN

14:Resample the stochastic weights in the SNN using the updated parameter distributions

15:endfor

16:returnThe trained SNN

In Examples[1](https://arxiv.org/html/2608.11613#Thmexample1)and 2, the reconstruction accuracy on the testing set is evaluated by comparing the empirical conditional means and variances of the reconstructed and reference distributions\. Let\{\(𝒙itest,𝒚itest\)\}\\\{\(\\bm\{x\}\_\{i\}^\{\\text\{test\}\},\\bm\{y\}\_\{i\}^\{\\text\{test\}\}\)\\\}denote samples in the testing set, andBtest​\(𝒙itest,0\)B^\{\\text\{test\}\}\(\\bm\{x\}\_\{i\}^\{\\text\{test\}\},0\)denote the set of observations with the identical input as𝒙i\\bm\{x\}\_\{i\}:

\{𝒙jtest:𝒙jtest=𝒙itest\}\.\\\{\\bm\{x\}\_\{j\}^\{\\text\{test\}\}:\\bm\{x\}\_\{j\}^\{\\text\{test\}\}=\\bm\{x\}\_\{i\}^\{\\text\{test\}\}\\\}\.\(3\.1\)The ground\-truth and predicted means are:

𝒎i=1R​∑𝒙jtest∈B⁡\(𝒙itest,0\)𝒚jtest,𝒎^i=1R​∑𝒙jtest∈B⁡\(𝒙itest,0\)𝒚^jtest\\bm\{m\}\_\{i\}=\\frac\{1\}\{R\}\\sum\_\{\\bm\{x\}\_\{j\}^\{\\text\{test\}\}\\in B\(\\bm\{x\}\_\{i\}^\{\\text\{test\}\},0\)\}\\bm\{y\}\_\{j\}^\{\\text\{test\}\},\\qquad\\hat\{\\bm\{m\}\}\_\{i\}=\\frac\{1\}\{R\}\\sum\_\{\\bm\{x\}\_\{j\}^\{\\text\{test\}\}\\in B\(\\bm\{x\}\_\{i\}^\{\\text\{test\}\},0\)\}\\hat\{\\bm\{y\}\}\_\{j\}^\{\\text\{test\}\}\(3\.2\)be the empirical conditional mean vectors at theithi^\{\\text\{th\}\}testing center, respectively\. Similarly, let

𝒗i=1R​∑𝒙jtest∈B⁡\(𝒙itest,0\)\(𝒚jtest−𝒎i\)∘2,𝒗^i=1R​∑𝒙jtest∈B⁡\(𝒙itest,0\)\(𝒚^jtest−𝒎^i\)∘2,\\bm\{v\}\_\{i\}=\\frac\{1\}\{R\}\\sum\_\{\\bm\{x\}\_\{j\}^\{\\text\{test\}\}\\in B\(\\bm\{x\}\_\{i\}^\{\\text\{test\}\},0\)\}\\left\(\\bm\{y\}\_\{j\}^\{\\text\{test\}\}\-\\bm\{m\}\_\{i\}\\right\)^\{\\circ 2\},\\qquad\\hat\{\\bm\{v\}\}\_\{i\}=\\frac\{1\}\{R\}\\sum\_\{\\bm\{x\}\_\{j\}^\{\\text\{test\}\}\\in B\(\\bm\{x\}\_\{i\}^\{\\text\{test\}\},0\)\}\\left\(\\hat\{\\bm\{y\}\}\_\{j\}^\{\\text\{test\}\}\-\\hat\{\\bm\{m\}\}\_\{i\}\\right\)^\{\\circ 2\},\(3\.3\)denote the empirical variance vectors at theithi^\{\\text\{th\}\}testing center, where\(⋅\)∘2\(\\cdot\)^\{\\circ 2\}denotes the element\-wise square\. In Eqs\. \([3\.2](https://arxiv.org/html/2608.11613#S3.E2)\) and \([3\.3](https://arxiv.org/html/2608.11613#S3.E3)\),R≔\|Btest​\(𝒙itest,0\)\|R\\coloneqq\|B^\{\\text\{test\}\}\(\\bm\{x\}\_\{i\}^\{\\text\{test\}\},0\)\|is the number of testing samples at each testing center\. The average relative mean and variance reconstruction error is defined as follows:

ℰmean=1Nc​∑i=1Nc‖𝒎^i−𝒎i‖2‖𝒎i‖2\+ϵ0,ℰvar=1Nc​∑i=1Nc‖𝒗^i−𝒗i‖2‖𝒗i‖2\+ϵ0,\\mathcal\{E\}\_\{\\mathrm\{mean\}\}=\\frac\{1\}\{N\_\{c\}\}\\sum\_\{i=1\}^\{N\_\{c\}\}\\frac\{\\\|\\hat\{\\bm\{m\}\}\_\{i\}\-\\bm\{m\}\_\{i\}\\\|\_\{2\}\}\{\\\|\\bm\{m\}\_\{i\}\\\|\_\{2\}\+\\epsilon\_\{0\}\},\\,\\,\\mathcal\{E\}\_\{\\mathrm\{var\}\}=\\frac\{1\}\{N\_\{c\}\}\\sum\_\{i=1\}^\{N\_\{c\}\}\\frac\{\\\|\\hat\{\\bm\{v\}\}\_\{i\}\-\\bm\{v\}\_\{i\}\\\|\_\{2\}\}\{\\\|\\bm\{v\}\_\{i\}\\\|\_\{2\}\+\\epsilon\_\{0\}\},\(3\.4\)whereNcN\_\{c\}is the number of testing centers, andϵ0\\epsilon\_\{0\}is set as10−810^\{\-8\}to avoid division by zero\.

First, we consider learning a unidimensional random field model to compare with different loss functions\.

###### Example 1\.

We consider a one\-dimensional conditional distribution reconstruction problem\. The objective is to learn the conditional distribution

p⁡\(y\|x\),x∈\[0,1\]\.p\(y\\,\|\\,x\),\\,\\,x\\in\[0,1\]\.\(3\.5\)Unlike conventional regression problems where only the conditional mean is estimated, our objective is to reconstruct the entire conditional distribution from scattered observations\. The conditional response is generated from a bimodal Gaussian mixture model,

y=m⁡\(x\)\+s​d​\(x\)\+ε0,y=m\(x\)\+s\\,d\(x\)\+\\varepsilon\_\{0\},\(3\.6\)where the random variables∈\{−1,\+1\}s\\in\\\{\-1,\+1\\\}is sampled with equal probability,

ℙ⁡\(s=−1\)=ℙ⁡\(s=1\)=12,\\mathbb\{P\}\(s=\-1\)=\\mathbb\{P\}\(s=1\)=\\frac\{1\}\{2\},andε0∼𝒩⁡\(0,σ2\)\\varepsilon\_\{0\}\\sim\\mathcal\{N\}\(0,\\sigma^\{2\}\)withσ=0\.04\\sigma=0\.04\. In Eq\. \([3\.6](https://arxiv.org/html/2608.11613#S3.E6)\), the conditional mean function is chosen as

m⁡\(x\)=0\.50\+0\.20​x\+exp⁡\[−5​\(x−0\.60\)2\]\+0\.40​sin⁡\(x2\),m\(x\)=0\.50\+0\.20x\+\\exp\\\!\\left\[\-5\(x\-0\.60\)^\{2\}\\right\]\+0\.40\\sin\\\!\\left\(\\frac\{x\}\{2\}\\right\),\(3\.7\)which produces a smooth nonlinear trend over the computational domain\. The separation between the two Gaussian modes is controlled by

d⁡\(x\)=0\.38\+0\.10​exp⁡\(−\(x−0\.70\)22⋅\(0\.14\)2\),d\(x\)=0\.38\+0\.10\\exp\\\!\\left\(\-\\frac\{\(x\-0\.70\)^\{2\}\}\{2\\cdot\(0\.14\)^\{2\}\}\\right\),\(3\.8\)which allows the distance between the two modes to vary smoothly with the input variable\. Consequently, the conditional distribution is given by

p⁡\(y\|x\)=12​𝒩​\(m⁡\(x\)−d⁡\(x\),σ2\)\+12​𝒩​\(m⁡\(x\)\+d⁡\(x\),σ2\)\.p\(y\|x\)=\\frac\{1\}\{2\}\\mathcal\{N\}\\left\(m\(x\)\-d\(x\),\\sigma^\{2\}\\right\)\+\\frac\{1\}\{2\}\\mathcal\{N\}\\left\(m\(x\)\+d\(x\),\\sigma^\{2\}\\right\)\.\(3\.9\)
For the training set,Ntrain=2000N\_\{\\rm train\}=2000samples are independently generated by first samplingxi∼𝒰⁡\(0,1\)x\_\{i\}\\sim\\mathcal\{U\}\(0,1\), followed by drawingyi∼p⁡\(y\|xi\)y\_\{i\}\\sim p\(y\|x\_\{i\}\)\. Therefore, each training sample is independent\. To generate the test set,100100uniformly distributed input locations are selected as neighborhood centers\. At each testing location,2020independent realizations are generated from the conditional distribution \(Eq\.[3\.9](https://arxiv.org/html/2608.11613#S3.E9)\), providing empirical estimates of the local conditional distributions used for quantitative evaluation\.

Figure 2:Comparison of different local distribution matching losses for Example[1](https://arxiv.org/html/2608.11613#Thmexample1)\(a\) Total training time\. \(b\) Average relative errors of the conditional mean and conditional variance on the testing set\. \(c\) Comparison between the true and predicted conditional distributions obtained using the proposed Local Sinkhorn loss\. The solid curves denote the conditional means, while the shaded regions represent one standard deviation\. All loss functions are minimized to train the same SNN architecture and identical training hyperparameters\.To investigate the influence of different discrepancy measures on conditional distribution learning, six different local loss functions are considered, namely local MSE, local MAE, local Energy Distance, local MMD, local squaredW2W\_\{2\}, and the proposed local Sinkhorn divergence\. The mathematical definitions of these losses are summarized in Appendix[D](https://arxiv.org/html/2608.11613#A4)\. Fig\.[2](https://arxiv.org/html/2608.11613#S3.F2)summarizes the overall performance of the six loss functions\. Fig\.[2](https://arxiv.org/html/2608.11613#S3.F2)\(a\) compares the total training time\. Owing to the entropic regularization, the proposed Local Sinkhorn loss requires significantly less computational effort than the previous local squaredW2W\_\{2\}loss in[xia2026local](https://arxiv.org/html/2608.11613#bib.bib31)\. Fig\.[2](https://arxiv.org/html/2608.11613#S3.F2)\(b\) reports the average relative errors of the conditional mean and conditional variance on the testing set\. Distribution\-based losses consistently outperform pointwise regression losses \(Local MSE and Local MAE\), indicating that matching local empirical distributions is substantially more effective than matching individual samples\. Among all competing losses, the proposed local Sinkhorn loss achieves the smallest overall errors for both the conditional mean and conditional variance\. Fig\.[2](https://arxiv.org/html/2608.11613#S3.F2)\(c\) compares the reconstructed conditional distributions obtained by the proposed Local Sinkhorn model with the reference solution\. The predicted mean accurately follows the true nonlinear trend while the predicted standard\-deviation band agrees closely with the reference distribution over the entire computational domain\. These results demonstrate that the proposed local Sinkhorn formulation is capable of accurately reconstructing both the location and spread of the conditional distribution\.

Next, we carry out an experiment in which the target random variable𝒚𝒙\\bm\{y\}\_\{\\bm\{x\}\}is multidimensional and compare our local Sinkhorn approach with other machine\-learning uncertainty quantification benchmarks\.

###### Example 2\.

We consider a stochastic Darcy flow problem on the spatial domain𝒙∈𝒟=\[0,1\]2\\bm\{x\}\\in\\mathcal\{D\}=\[0,1\]^\{2\}\. For each realization of the uncertain model parametersω∈Ω\\omega\\in\\Omega, we solve:

−∇⋅\(a\(𝒙,ω\)∇u\(𝒙,ω\)\)=f\(𝒙\),𝒙∈D,\-\\nabla\\cdot\\left\(a\(\\bm\{x\},\\omega\)\\nabla u\(\\bm\{x\},\\omega\)\\right\)=f\(\\bm\{x\}\),\\qquad\\bm\{x\}\\in D,\(3\.10\)subject to homogeneous Dirichlet boundary conditions

u⁡\(𝒙,ω\)=0,𝒙∈∂D\.u\(\\bm\{x\},\\omega\)=0,\\qquad\\bm\{x\}\\in\\partial D\.\(3\.11\)
The permeability field is modeled as:

a⁡\(𝒙,ω\)=exp⁡\(g⁡\(𝒙,ω\)\),a\(\\bm\{x\},\\omega\)=\\exp\\\!\\left\(g\(\\bm\{x\},\\omega\)\\right\),\(3\.12\)whereg⁡\(𝒙,ω\)g\(\\bm\{x\},\\omega\)is a zero\-mean Gaussian random field generated by FFT\-based spectral synthesis\. Specifically, the random field is constructed in the Fourier domain as:

g^​\(k,ω\)=c⁡\(k\)​ξ^​\(k,ω\),\\widehat\{g\}\(k,\\omega\)=\\sqrt\{c\(k\)\}\\,\\widehat\{\\xi\}\(k,\\omega\),\(3\.13\)whereξ^​\(k,ω\)\\widehat\{\\xi\}\(k,\\omega\)are independent complex Gaussian random variables satisfyingξ^​\(k,ω\)∼𝒩​\(0,1\)\\widehat\{\\xi\}\(k,\\omega\)\\sim\\mathcal\{N\}\(0,1\), and the covariance spectrum is prescribed by

c\(k\)=\(1\+\(2πℓ\)2\|k\|2\)−α/2,c\(k\)=\\left\(1\+\(2\\pi\\ell\)^\{2\}\|k\|^\{2\}\\right\)^\{\-\\alpha/2\},\(3\.14\)with correlation lengthℓ\\elland smoothness parameterα\\alpha\. The Gaussian random field is then obtained by applying the inverse Fourier transform,

g⁡\(𝒙,ω\)=ℱ−1​\(g^​\(k,ω\)\),g\(\\bm\{x\},\\omega\)=\\mathcal\{F\}^\{\-1\}\\\!\\left\(\\widehat\{g\}\(k,\\omega\)\\right\),\(3\.15\)followed by normalization to zero mean and unit variance\.

To generate the training and testing set, we establish a uniform64×6464\\times 64grid on𝒟\\mathcal\{D\}denoted as𝒟64=\{𝒙i,j=\(i63,j63\)\}i,j=064\\mathcal\{D\}\_\{64\}=\\\{\\bm\{x\}\_\{i,j\}=\(\\frac\{i\}\{63\},\\frac\{j\}\{63\}\)\\\}\_\{i,j=0\}^\{64\}\. Then, we randomly sample𝒙i,j\\bm\{x\}\_\{i,j\}from𝒟64\\mathcal\{D\}\_\{64\}satisfying6≤i,j≤576\\leq i,j\\leq 57\. For the chosen𝒙i,j\\bm\{x\}\_\{i,j\}, a square window of size7×77\\times 7is first constructed on the underlying computational grid\. A total of1616grid points𝒙i′,j′\\bm\{x\}\_\{i^\{\\prime\},j^\{\\prime\}\}are then randomly sampled without replacement from this window\. This sampling strategy is adopted to ensure that most training samples have enough neighbors around them\. For each selected grid point𝒙i′,j′\\bm\{x\}\_\{i^\{\\prime\},j^\{\\prime\}\}, a5×55\\times 5solution patch on the grid is extracted with the selected grid point being the center\. The corresponding solution patch is defined as:

𝒚⁡\(𝒙i′,j′,ωi′,j′\)=\(u⁡\(𝒙i′−2,j′−2,ωi′−2,j′−2\),u⁡\(𝒙i′−2,j′−1,ωi′−2,j′−1\),…,u⁡\(𝒙i′\+2,j′\+2,ωi′\+2,j′\+2\)\)⊤∈ℝ25\.\\bm\{y\}\(\\bm\{x\}\_\{i^\{\\prime\},j^\{\\prime\}\},\\omega\_\{i^\{\\prime\},j^\{\\prime\}\}\)=\\Big\(u\(\\bm\{x\}\_\{i^\{\\prime\}\-2,j^\{\\prime\}\-2\},\\omega\_\{i^\{\\prime\}\-2,j^\{\\prime\}\-2\}\),u\(\\bm\{x\}\_\{i^\{\\prime\}\-2,j^\{\\prime\}\-1\},\\omega\_\{i^\{\\prime\}\-2,j^\{\\prime\}\-1\}\),\\ldots,u\(\\bm\{x\}\_\{i^\{\\prime\}\+2,j^\{\\prime\}\+2\},\\omega\_\{i^\{\\prime\}\+2,j^\{\\prime\}\+2\}\)\\Big\)^\{\\top\}\\in\\mathbb\{R\}^\{25\}\.\(3\.16\)An illustration of the training and testing set sample generation strategy is given in Fig\.[3](https://arxiv.org/html/2608.11613#S3.F3)\(a\)\. The vector𝒚⁡\(𝒙,ω\)\\bm\{y\}\(\\bm\{x\},\\omega\)serves as the target random field whose conditional distribution is to be reconstructed when𝒙\\bm\{x\}is given\. To generate the testing set, 20 testing centers are generated\. For each testing center𝒙ℓtest\\bm\{x\}\_\{\\ell\}^\{\\mathrm\{test\}\}, a large number of independent permeability realizationsω1,…,ωM\\omega\_\{1\},\\ldots,\\omega\_\{M\}are generated\. The resulting collection

\{𝒚⁡\(𝒙ℓtest,ωℓ\)\}j=1M\\left\\\{\\bm\{y\}\(\\bm\{x\}\_\{\\ell\}^\{\\mathrm\{test\}\},\\omega\_\{\\ell\}\)\\right\\\}\_\{j=1\}^\{M\}\(3\.17\)provides an empirical approximation of the conditional distributionp⁡\(𝒚∣𝒙ℓtest\)p\\\!\\left\(\\bm\{y\}\\mid\\bm\{x\}\_\{\\ell\}^\{\\mathrm\{test\}\}\\right\)\. The learned model is then evaluated by comparing the predicted conditional distribution with this empirical reference distribution at the test locations\. This benchmark naturally produces multidimensional output variables with strong spatial correlations, making it a suitable experiment to test if our local Sinkhorn approach could successfully reconstruct a multidimensional random field in which the noise is strongly heterogeneous or lies in a low\-dimensional manifold\.

To evaluate the effectiveness of the proposed local Sinkhorn divergence, we compare it with several representative approaches for conditional distribution reconstruction\. The compared loss functions include local Maximum Mean Discrepancy \(MMD\) and the local squaredW2W\_\{2\}loss proposed in our previous work[xia2026local](https://arxiv.org/html/2608.11613#bib.bib31), whose definitions are detailed in Appendix[D](https://arxiv.org/html/2608.11613#A4)\. To further compare with representative probabilistic learning approaches, we also consider Heteroscedastic Gaussian Regression, Mixture Density Networks \(MDN\)[bishop1994mixture](https://arxiv.org/html/2608.11613#bib.bib4), Conditional Variational Autoencoders \(CVAE\)[sohn2015learning](https://arxiv.org/html/2608.11613#bib.bib24), and Conditional Normalizing Flows \(CNF\)[rezende2015variational](https://arxiv.org/html/2608.11613#bib.bib22);[dinh2017density](https://arxiv.org/html/2608.11613#bib.bib7);[papamakarios2021normalizing](https://arxiv.org/html/2608.11613#bib.bib19)\. These methods represent commonly used likelihood\-based and latent\-variable models for conditional distribution estimation\. All benchmark methods are trained and evaluated using the same datasets and testing protocols to ensure a fair comparison\.

Table 2:Comparison of different conditional distribution learning methods and loss function used to train the SNN on the Darcy flow benchmark\.MethodMean ErrorVariance ErrorTraining Time \(s\)Memory \(MB\)Heteroscedastic Gaussian Regression0\.11890\.510423\.48687\.38Mixture Density Network0\.09705\.310152\.06686\.72Conditional VAE0\.29315\.868731\.10695\.09Conditional Normalizing Flow0\.32563\.5312140\.36689\.73SNN \+ LocalW2W\_\{2\}0\.05190\.3213500\.84905\.14SNN \+ Sinkhorn0\.04730\.2318308\.34791\.08SNN \+ MMD0\.07790\.7878649\.60733\.90Table[2](https://arxiv.org/html/2608.11613#S3.T2)compares the proposed neighborhood\-based distribution learning methods with several representative conditional generative baselines\. Among the likelihood\-based models, the heteroscedastic Gaussian regression is computationally the most efficient, requiring only23\.4823\.48seconds for training, but its Gaussian assumption limits the reconstruction accuracy, especially for the variance\. MDN improves the mean prediction accuracy but still suffers from poor variance reconstruction\. CVAE and CNF provide more expressive conditional generative models, although their reconstruction errors remain significantly larger than those of the neighborhood\-based approaches\.

The neighborhood\-based optimal transport methods consistently outperform the global conditional generative models\. In particular, training the SNN with the local Sinkhorn divergence achieves the smallest errors in the predicted mean and variance\. Furthermore, compared to the previous local squaredW2W\_\{2\}approach, its computational efficiency is improved, and the runtime is shortened\. The proposed neighborhood\-based Sinkhorn approach to train SNN provides the best overall trade\-off between reconstruction accuracy and computational efficiency\. These results suggest that once local conditional distributions are exploited with the neighborhood technique, the simple SNN can be efficiently trained using our local Sinkhorn loss, which is sufficient to outperform considerably more sophisticated machine\-learning conditional generative models\. Overall, the experimental results demonstrate that exploiting local neighborhoods when utilizing the OT\-based loss functions is considerably more important than increasing the complexity of the conditional generator\.

To investigate the robustness of the proposed local Sinkhorn framework with respect to different algorithmic parameters, we further perform sensitivity studies on the noise level of the training data, the neighborhood radius, and the Sinkhorn regularization parameter\. The corresponding results are presented in Fig\.[3](https://arxiv.org/html/2608.11613#S3.F3)\. For the noise sensitivity study, we vary the intensity of the underlying Gaussian random field\. The permeability field is constructed from

gσ​\(𝒙,ω\)=σ​g​\(𝒙,ω\),g\_\{\\sigma\}\(\\bm\{x\},\\omega\)=\\sigma\\,g\(\\bm\{x\},\\omega\),\(3\.18\)whereσ\>0\\sigma\>0is referred to as the noise level andggdenotes the normalized Gaussian random field generated by the FFT\-based spectral synthesis in Eq\. \([3\.15](https://arxiv.org/html/2608.11613#S3.E15)\)\. Consequently, the permeability field is given by

a⁡\(𝒙,ω\)=exp⁡\(gσ​\(𝒙,ω\)\)\.a\(\\bm\{x\},\\omega\)=\\exp\\\!\\left\(g\_\{\\sigma\}\(\\bm\{x\},\\omega\)\\right\)\.Increasingσ\\sigmaenlarges the variance of the permeability field and therefore increases the intrinsic uncertainty of the corresponding Darcy solutions\.

![Refer to caption](https://arxiv.org/html/2608.11613v1/example2_new.png)Figure 3:\(a\) An illustration of the training and testing sample selection strategy\. \(b\) Influence of the random\-field intensity parameterσ\\sigma, where the Gaussian random field is generated asgσ​\(𝒙,ω\)=σ​g0​\(𝒙,ω\)g\_\{\\sigma\}\(\\bm\{x\},\\omega\)=\\sigma g\_\{0\}\(\\bm\{x\},\\omega\)\. \(c\) Influence of the neighborhood radiusδ\\deltaused to construct local empirical conditional distributions\. \(d\) Influence of the Sinkhorn regularization parameterε\\varepsilon\. For all panels, the left axis reports the average relative reconstruction errors of the conditional mean and variance, while the right axis reports the corresponding training time\.Fig\.[3](https://arxiv.org/html/2608.11613#S3.F3)\(b\) illustrates the influence of the observation noise level\. As expected, increasing the noise level deteriorates the reconstruction accuracy of both the conditional mean and variance, while the training time increases only moderately\. This indicates that the proposed method remains computationally stable even under relatively noisy observations\. Fig\.[3](https://arxiv.org/html/2608.11613#S3.F3)\(c\) studies the effect of the neighborhood radiusδ\\delta\. A sufficiently small neighborhood does not contain enough local samples to accurately approximate the conditional distribution, whereas an excessively large neighborhood violates the locality assumption and mixes samples associated with different conditional distributions\. Consequently, an intermediate neighborhood radius provides the best trade\-off between approximation accuracy and computational cost\. Fig\.[3](https://arxiv.org/html/2608.11613#S3.F3)\(d\) shows the sensitivity of the proposed method to the Sinkhorn regularization parameterε\\varepsilon\. Whenε\\varepsilonis too small, the Sinkhorn divergence approaches the exact Wasserstein distance\. In this regime, the reconstructed variance error is large because the empirical squaredW2W\_\{2\}distance converges slowly to the ground truth squaredW2W\_\{2\}distance in multidimensional settings\. Furthermore, choosing a too smallε\\varepsilonleads to longer runtime\. Conversely, a large regularization parameter introduces excessive entropic smoothing and leads to larger reconstruction errors\. These observations agree well with the theoretical error analysis in Section[2](https://arxiv.org/html/2608.11613#S2), where the regularization parameter controls the trade\-off between approximation bias and computational efficiency, and a moderate regularization parameterε\\varepsilonis preferable to reconstruct multidimensional random field models\. Overall, the proposed local Sinkhorn framework exhibits good robustness over a reasonably wide range of neighborhood radii and regularization parameters\.

Finally, we consider reconstructing a stochastic dynamical system using our proposed local Sinkhorn approach\.

###### Example 3\.

To further demonstrate the applicability of the proposed method to multidimensional stochastic dynamical systems, we consider a network of coupled nonlinear stochastic FHN oscillators in computational neuroscience[Tuckwell1998](https://arxiv.org/html/2608.11613#bib.bib28);[Acebron2004](https://arxiv.org/html/2608.11613#bib.bib1)\.

We consider a network consisting ofNn=5N\_\{n\}=5coupled neurons, resulting in a1010\-dimensional stochastic differential equation\. Let

𝒙⁡\(t,𝒙0,ω\)=\(v1​\(t\),…,v5​\(t\),w1​\(t\),…,w5​\(t\)\)⊤∈ℝ10,\\bm\{x\}\(t;\\bm\{x\}\_\{0\};\\omega\)=\\big\(v\_\{1\}\(t\),\\ldots,v\_\{5\}\(t\),w\_\{1\}\(t\),\\ldots,w\_\{5\}\(t\)\\big\)^\{\\top\}\\in\\mathbb\{R\}^\{10\},whereviv\_\{i\}denotes the membrane potential andwiw\_\{i\}is the corresponding recovery variable\. The governing stochastic system is

d​vi​\(t\)\\displaystyle\\mbox\{d\}v\_\{i\}\(t\)=\(vi−vi33−wi\+∑j=15Ci​j​\(ω\)​vj\)​d​t\+σi​\(ω\)​d​Bi​\(t\),\\displaystyle=\\left\(v\_\{i\}\-\\frac\{v\_\{i\}^\{3\}\}\{3\}\-w\_\{i\}\+\\sum\_\{j=1\}^\{5\}C\_\{ij\}\(\\omega\)v\_\{j\}\\right\)\\mbox\{d\}t\+\\sigma\_\{i\}\(\\omega\)\\,\\mbox\{d\}B\_\{i\}\(t\),\(3\.19\)d​wi​\(t\)\\displaystyle\\mbox\{d\}w\_\{i\}\(t\)=εi\(ω\)\(vi\+c−dwi\)dt,i=1,…,5,\(v1\(0\),…,v5\(0\),w1\(0\),…,w5\(0\)\)=𝒙0,\\displaystyle=\\varepsilon\_\{i\}\(\\omega\)\\left\(v\_\{i\}\+c\-dw\_\{i\}\\right\)\\mbox\{d\}t,\\qquad i=1,\\ldots,5,\\,\\,\(v\_\{1\}\(0\),\\ldots,v\_\{5\}\(0\),w\_\{1\}\(0\),\\ldots,w\_\{5\}\(0\)\)=\\bm\{x\}\_\{0\},whereBi​\(t\)B\_\{i\}\(t\)are independent standard Brownian motions, andc=0\.7c=0\.7andd=0\.8d=0\.8are fixed deterministic parameters\.

Unlike the classical FHN model, the coupling strengthsCi​j​\(ω\)C\_\{ij\}\(\\omega\), the diffusion coefficientsσi​\(ω\)\\sigma\_\{i\}\(\\omega\), and the recovery time scalesεi​\(ω\)\\varepsilon\_\{i\}\(\\omega\)are all regarded as random variables whose values may vary across different realizations, andω\\omegadenotes the uncertainty in those model parameters\. More precisely,\(Ci​j​\(ω\),σi​\(ω\),εi​\(ω\)\)\\left\(C\_\{ij\}\(\\omega\),\\sigma\_\{i\}\(\\omega\),\\varepsilon\_\{i\}\(\\omega\)\\right\)is independently sampled for every realization, yielding a stochastic drift fielda⁡\(x,ω\)a\(x,\\omega\)and a stochastic diffusion fieldb⁡\(x,ω\)b\(x,\\omega\)\. Consequently, the dynamics can be written compactly as

d​𝒙​\(t,𝒙0,ω\)=a⁡\(𝒙⁡\(t,𝒙0,ω\),ω\)​d​t\+b⁡\(𝒙⁡\(t,𝒙0,ω\),ω\)​d​𝑩t,𝒙⁡\(0,𝒙0,ω\)=𝒙0,t∈\[0,T\],\\mbox\{d\}\\bm\{x\}\(t;\\bm\{x\}\_\{0\};\\omega\)=a\(\\bm\{x\}\(t;\\bm\{x\}\_\{0\};\\omega\),\\omega\)\\,\\mbox\{d\}t\+b\(\\bm\{x\}\(t;\\bm\{x\}\_\{0\};\\omega\),\\omega\)\\,\\mbox\{d\}\\bm\{B\}\_\{t\},\\,\\,\\bm\{x\}\(0;\\bm\{x\}\_\{0\};\\omega\)=\\bm\{x\}\_\{0\},\\,\\,t\\in\[0,T\],\(3\.20\)where both the drift and diffusion are random functions induced by the uncertain parameters,𝑩t\\bm\{B\}\_\{t\}is a five\-dimensional standard Brownian motion, and𝒙0\\bm\{x\}\_\{0\}denotes the initial condition, with𝒙⁡\(t,𝒙0,ω\)\\bm\{x\}\(t;\\bm\{x\}\_\{0\};\\omega\)explicitly indicating the dependence of the solution on𝒙0\\bm\{x\}\_\{0\}\.

The reference trajectories and predicted trajectories are generated using the Euler–Maruyama solver implemented in thetorchsdepackage\. For the training set, each trajectory is associated with an independently sampled initial condition and an independent realization of the random parameter vectorω\\omega\. The resulting dataset consists of\{\(𝒙0\(i\),𝒙\(i\)​\(t1,𝒙0\(i\),ω\(i\)\),…,𝒙\(i\)​\(tNt,𝒙0\(i\),ω\(i\)\)\)\}i=1Ntrain\\Big\\\{\\big\(\\bm\{x\}\_\{0\}^\{\(i\)\},\\bm\{x\}^\{\(i\)\}\(t\_\{1\};\\bm\{x\}\_\{0\}^\{\(i\)\};\\omega^\{\(i\)\}\),\\ldots,\\bm\{x\}^\{\(i\)\}\(t\_\{N\_\{t\}\};\\bm\{x\}\_\{0\}^\{\(i\)\};\\omega^\{\(i\)\}\)\\big\)\\Big\\\}\_\{i=1\}^\{N\_\{\\mathrm\{train\}\}\}, where𝒙\(i\)​\(tk\)∈ℝ10\\bm\{x\}^\{\(i\)\}\(t\_\{k\}\)\\in\\mathbb\{R\}^\{10\}is the snapshot of the trajectory at timetk=k​TNtt\_\{k\}=\\frac\{kT\}\{N\_\{t\}\}\.

For the testing set, we randomly select2020initial conditions\. For each initial condition,2020independent realizations of the random parameters are generated, producing empirical conditional distributions of the solution trajectories\. Therefore, the testing data naturally approximates the conditional distributionP⁡\(𝒙test​\(t,𝒙0,ω\)∣𝒙test​\(0,𝒙0,ω\)=𝒙0\)P\\left\(\\bm\{x\}^\{\\text\{test\}\}\(t;\\bm\{x\}\_\{0\};\\omega\)\\mid\\bm\{x\}^\{\\text\{test\}\}\(0;\\bm\{x\}\_\{0\};\\omega\)=\\bm\{x\}\_\{0\}\\right\), which serves as the reference distribution for evaluating the reconstructed stochastic dynamics\. For both the training set and the testing set, we set the time horizonT=1T=1and the time stepsNt=10N\_\{t\}=10in Eq\. \([3\.20](https://arxiv.org/html/2608.11613#S3.E20)\)\.

To reconstruct the unknown stochastic dynamics, we construct the approximate model:

d​𝒙^​\(t,𝒙0,ω^\)=a^​\(𝒙^​\(t,𝒙0,ω^\),ω^a\)​d​t\+b^​\(𝒙^​\(t,𝒙0,ω^\),ω^b\)​d​Bt,𝒙^​\(0,𝒙0,ω^\)=𝒙0,\\mbox\{d\}\\hat\{\\bm\{x\}\}\(t;\\bm\{x\}\_\{0\};\\hat\{\\omega\}\)=\\hat\{a\}\(\\hat\{\\bm\{x\}\}\(t;\\bm\{x\}\_\{0\};\\hat\{\\omega\}\),\\hat\{\\omega\}\_\{a\}\)\\,\\mbox\{d\}t\+\\hat\{b\}\(\\hat\{\\bm\{x\}\}\(t;\\bm\{x\}\_\{0\};\\hat\{\\omega\}\),\\hat\{\\omega\}\_\{b\}\)\\,\\mbox\{d\}B\_\{t\},\\,\\,\\hat\{\\bm\{x\}\}\(0;\\bm\{x\}\_\{0\};\\hat\{\\omega\}\)=\\bm\{x\}\_\{0\},\(3\.21\)wherea^​\(𝒙⁡\(t,𝒙0,ω^\),t,ω^a\)\\hat\{a\}\(\\bm\{x\}\(t;\\bm\{x\}\_\{0\};\\hat\{\\omega\}\),t;\\hat\{\\omega\}\_\{a\}\)andb^​\(𝒙⁡\(t,𝒙0,ω^\),t,ω^b\)\\hat\{b\}\(\\bm\{x\}\(t;\\bm\{x\}\_\{0\};\\hat\{\\omega\}\),t;\\hat\{\\omega\}\_\{b\}\)are two separate SNNs andω^=\(ω^a,ω^b\)\\hat\{\\omega\}=\(\\hat\{\\omega\}\_\{a\},\\hat\{\\omega\}\_\{b\}\)refers to the union of uncertain weights ina^\\hat\{a\}andb^\\hat\{b\}\.a^\\hat\{a\}andb^\\hat\{b\}in Eq\. \([3\.21](https://arxiv.org/html/2608.11613#S3.E21)\) take the 10\-dimensional state variable𝒙^​\(t,𝒙0,ω^\)\\hat\{\\bm\{x\}\}\(t;\\bm\{x\}\_\{0\};\\hat\{\\omega\}\)and timettas the input and then output 10\-dimensional approximate drift and diffusion functions to approximate the stochastic drift and diffusion coefficients in Eq\. \([3\.20](https://arxiv.org/html/2608.11613#S3.E20)\), respectively\. Since the dynamics ofwiw\_\{i\}in Eq\. \([3\.19](https://arxiv.org/html/2608.11613#S3.E19)\) are noise\-free, we enforce the last five components ofb^\\hat\{b\}in Eq\. \([3\.21](https://arxiv.org/html/2608.11613#S3.E21)\) to be 0\. During one trajectory simulation, a single realization of the network weightsω^\\hat\{\\omega\}is sampled and kept fixed throughout the entire time interval\. Consequently, the randomness remains trajectory\-wise rather than time\-wise, which is consistent with the assumption that the underlying uncertain parameters are realization\-dependent\. The two SNNs are trained simultaneously by minimizing a temporal\-average local Sinkhorn loss:

1T​∫0TSε,δe​\(𝒙⁡\(t,𝒙0,ω\),𝒙^​\(t,𝒙0,ω^\)\)​d​t\\displaystyle\\frac\{1\}\{T\}\\int\_\{0\}^\{T\}S\_\{\\varepsilon,\\delta\}^\{\\text\{e\}\}\(\\bm\{x\}\\big\(t;\\bm\{x\}\_\{0\};\\omega\),\\hat\{\\bm\{x\}\}\(t;\\bm\{x\}\_\{0\};\\hat\{\\omega\}\)\\big\)\\mbox\{d\}t≈1NT​∑i=1NTS¯ε,δe​\(𝒙⁡\(i​TNT,𝒙0,ω\),𝒙^​\(i​TNT,𝒙0,ω^\)\)\\displaystyle\\approx\\frac\{1\}\{N\_\{T\}\}\\sum\_\{i=1\}^\{N\_\{T\}\}\\overline\{S\}^\{\\text\{e\}\}\_\{\\varepsilon,\\delta\}\(\\bm\{x\}\\big\(\\tfrac\{iT\}\{N\_\{T\}\};\\bm\{x\}\_\{0\};\\omega\),\\hat\{\\bm\{x\}\}\(\\tfrac\{iT\}\{N\_\{T\}\};\\bm\{x\}\_\{0\};\\hat\{\\omega\}\)\\big\)\(3\.22\)=1NT​∑i=1NT∫𝒟Sε​\(μ𝒙0,δe​\(i​TNT\),μ^𝒙0,δe​\(i​TNT\)\)​νe​\(d​𝒙0\),\\displaystyle=\\frac\{1\}\{N\_\{T\}\}\\sum\_\{i=1\}^\{N\_\{T\}\}\\int\_\{\\mathcal\{D\}\}S\_\{\\varepsilon\}\\big\(\\mu^\{\\text\{e\}\}\_\{\\bm\{x\}\_\{0\},\\delta\}\(\\tfrac\{iT\}\{N\_\{T\}\}\),\\hat\{\\mu\}^\{\\text\{e\}\}\_\{\\bm\{x\}\_\{0\},\\delta\}\(\\tfrac\{iT\}\{N\_\{T\}\}\)\\big\)\\nu^\{\\text\{e\}\}\(\\mbox\{d\}\\bm\{x\}\_\{0\}\),where𝒙0\\bm\{x\}\_\{0\}represents the initial condition, andμ𝒙0,δe​\(t\)\\mu^\{e\}\_\{\\bm\{x\}\_\{0\},\\delta\}\(t\)andμ^𝒙0,δe​\(t\)\\widehat\{\\mu\}^\{e\}\_\{\\bm\{x\}\_\{0\},\\delta\}\(t\)denote the empirical distributions of𝒙⁡\(t,𝒙0,ω\)\\bm\{x\}\(t;\\bm\{x\}\_\{0\};\\omega\)and𝒙^​\(t,𝒙0,ω^\)\\widehat\{\\bm\{x\}\}\(t;\\bm\{x\}\_\{0\};\\widehat\{\\omega\}\), respectively, constructed from trajectories whose initial conditions fall within the selected neighborhood subsetℬ⁡\(𝒙0,δ\)⊆B⁡\(𝒙0,δ\)\\mathcal\{B\}\(\\bm\{x\}\_\{0\},\\delta\)\\subseteq B\(\\bm\{x\}\_\{0\},\\delta\)\.

The proposed temporal\-average local Sinkhorn loss Eq\. \([3\.22](https://arxiv.org/html/2608.11613#S3.E22)\) compares the empirical distributions of the predicted and reference trajectories at every observation time and then averages the discrepancies over the entire temporal domain\. Compared with the temporal\-average local squaredW2W\_\{2\}loss introduced in our previous work[xia2026local](https://arxiv.org/html/2608.11613#bib.bib31), which replaces the Sinkhorn divergence with the squaredW2W\_\{2\}distance in Eq\. \([3\.22](https://arxiv.org/html/2608.11613#S3.E22)\), the proposed Sinkhorn formulation provides a differentiable approximation with substantially improved computational efficiency while preserving the geometric structure of optimal transport\.

Figure 4:Comparison between the true and predicted stochastic dynamics for a representative testing sample in the stochastic FHN system\. Panels \(a\)\(b\) show the mean and standard deviation of trajectories ofv1​\(t\)v\_\{1\}\(t\)andw1​\(t\)w\_\{1\}\(t\)\. \(c\)\(d\) show the mean and standard deviation of the drift functions ofd​v1​\(t\)\\mbox\{d\}v\_\{1\}\(t\)andd​w1​\(t\)\\mbox\{d\}w\_\{1\}\(t\)in Eq\. \([3\.19](https://arxiv.org/html/2608.11613#S3.E19)\)\. \(e\) shows the mean and standard deviation of the diffusion functions ofd​v1​\(t\)\\mbox\{d\}v\_\{1\}\(t\)in Eq\. \([3\.19](https://arxiv.org/html/2608.11613#S3.E19)\)\. The shaded regions indicate one standard deviation estimated from multiple stochastic realizations\.To further evaluate the reconstruction quality of the proposed method, Fig\.[4](https://arxiv.org/html/2608.11613#S3.F4)compares the predicted stochastic dynamics with the corresponding reference solution for a representative testing sample\. In addition to the stochastic trajectories, the learned drift and diffusion functions are also compared with the ground truth\. Since the proposed model reconstructs conditional probability distributions rather than individual trajectories, the conditional mean and one\-standard\-deviation bands estimated from multiple stochastic realizations are reported\. As shown in Fig\.[4](https://arxiv.org/html/2608.11613#S3.F4)\(a\)\(b\), the predicted trajectories accurately reproduce both the temporal evolution and the associated uncertainty of the stochastic FHN system\. For both the membrane potential variablev1v\_\{1\}and the recovery variablew1w\_\{1\}, the predicted conditional means closely follow the reference solution, while the predicted uncertainty bands exhibit good agreement with the corresponding ground\-truth distributions\. Furthermore, the learned drift functions capture the nonlinear deterministic dynamics with high accuracy \(Fig\.[4](https://arxiv.org/html/2608.11613#S3.F4)\(c\)\(d\)\), and the diffusion functions provide reasonable approximations of the stochastic forcing, while larger errors in the predicted diffusion functions may be due to error accumulation \(shown in Fig\.[4](https://arxiv.org/html/2608.11613#S3.F4)\(e\)\)\. Overall, these results demonstrate that the proposed local Sinkhorn divergence is capable of simultaneously reconstructing the deterministic and stochastic components of the underlying dynamics from scattered conditional observations\.

We also compare minimizing our temporal\-average local Sinkhorn divergence Eq\. \([3\.22](https://arxiv.org/html/2608.11613#S3.E22)\) with minimizing the previous temporal\-average local squaredW2W\_\{2\}loss in[xia2026local](https://arxiv.org/html/2608.11613#bib.bib31)\. The following error metrics are used to quantify the errors of𝒙^\\hat\{\\bm\{x\}\},a^\\hat\{\{a\}\}, andb^\\hat\{\{b\}\}of the reconstructed SDE \([3\.21](https://arxiv.org/html/2608.11613#S3.E21)\) on the testing set, respectively:

error in𝒙^≔∑j=1Ntest∑i=1NTW22​\(μ𝒙0,jteste​\(i​TNT\),μ^𝒙0,jteste​\(i​TNT\)\)∑j=1Ntest∑i=1NTE​\[‖𝒙⁡\(i​TNT,𝒙0,jtest,ω\)‖2\],error ina^≔∑j=1Ntest∑i=1NTW22​\(ηa,𝒙0,jteste​\(i​TNT\),η^a,𝒙0,jteste​\(i​TNT\)\)∑j=1Ntest∑i=1NTE​\[‖a⁡\(𝒙⁡\(i​TNT,𝒙0,jtest,ω\),ω\)‖2\],\\displaystyle\\text\{error in\}~\\hat\{\\bm\{x\}\}\\coloneqq\\frac\{\\sum\_\{j=1\}^\{N\_\{\\text\{test\}\}\}\\sum\_\{i=1\}^\{N\_\{T\}\}W\_\{2\}^\{2\}\\big\(\\mu^\{\\text\{e\}\}\_\{\\bm\{x\}\_\{0,j\}^\{\\text\{test\}\}\}\(\\tfrac\{iT\}\{N\_\{T\}\}\),\\hat\{\\mu\}^\{\\text\{e\}\}\_\{\\bm\{x\}\_\{0,j\}^\{\\text\{test\}\}\}\(\\tfrac\{iT\}\{N\_\{T\}\}\)\\big\)\}\{\\sum\_\{j=1\}^\{N\_\{\\text\{test\}\}\}\\sum\_\{i=1\}^\{N\_\{T\}\}\\mbox\{E\}\\Big\[\\big\\\|\\bm\{x\}\(\\tfrac\{iT\}\{N\_\{T\}\};\\bm\{x\}\_\{0,j\}^\{\\text\{test\}\};\\omega\)\\big\\\|^\{2\}\\Big\]\},\\,\\,\\,\\,\\,\\,\\text\{error in\}~\\hat\{a\}\\coloneqq\\frac\{\\sum\_\{j=1\}^\{N\_\{\\text\{test\}\}\}\\sum\_\{i=1\}^\{N\_\{T\}\}W\_\{2\}^\{2\}\\big\(\\eta^\{\\text\{e\}\}\_\{a,\\bm\{x\}\_\{0,j\}^\{\\text\{test\}\}\}\(\\tfrac\{iT\}\{N\_\{T\}\}\),\\hat\{\\eta\}^\{\\text\{e\}\}\_\{a,\\bm\{x\}\_\{0,j\}^\{\\text\{test\}\}\}\(\\tfrac\{iT\}\{N\_\{T\}\}\)\\big\)\}\{\\sum\_\{j=1\}^\{N\_\{\\text\{test\}\}\}\\sum\_\{i=1\}^\{N\_\{T\}\}\\mbox\{E\}\\Big\[\\big\\\|\{a\}\(\\bm\{x\}\(\\tfrac\{iT\}\{N\_\{T\}\};\\bm\{x\}\_\{0,j\}^\{\\text\{test\}\};\\omega\),\\omega\)\\big\\\|^\{2\}\\Big\]\},\(3\.23\)error in​b^≔∑j=1Ntest∑i=1NTW22​\(ηb,𝒙0,jteste​\(i​TNT\),η^b,𝒙0,jteste​\(i​TNT\)\)∑j=1Ntest∑i=1NTE​\[‖b⁡\(𝒙⁡\(i​TNT,𝒙0,jtest,ω\),ω\)‖2\],\\displaystyle\\text\{error in\}~\\hat\{b\}\\coloneqq\\frac\{\\sum\_\{j=1\}^\{N\_\{\\text\{test\}\}\}\\sum\_\{i=1\}^\{N\_\{T\}\}W\_\{2\}^\{2\}\\big\(\\eta^\{\\text\{e\}\}\_\{b,\\bm\{x\}\_\{0,j\}^\{\\text\{test\}\}\}\(\\tfrac\{iT\}\{N\_\{T\}\}\),\\hat\{\\eta\}^\{\\text\{e\}\}\_\{b,\\bm\{x\}\_\{0,j\}^\{\\text\{test\}\}\}\(\\tfrac\{iT\}\{N\_\{T\}\}\)\\big\)\}\{\\sum\_\{j=1\}^\{N\_\{\\text\{test\}\}\}\\sum\_\{i=1\}^\{N\_\{T\}\}\\mbox\{E\}\\Big\[\\big\\\|\{b\}\(\\bm\{x\}\(\\tfrac\{iT\}\{N\_\{T\}\};\\bm\{x\}\_\{0,j\}^\{\\text\{test\}\};\\omega\),\\omega\)\\big\\\|^\{2\}\\Big\]\},where𝒙0,jtest\\bm\{x\}\_\{0,j\}^\{\\text\{test\}\}refers to thejtestj^\{\\text\{test\}\}center in the testing set,μ𝒙0,jteste​\(i​TNT\),μ^𝒙0,jteste​\(i​TNT\)\\mu^\{\\text\{e\}\}\_\{\\bm\{x\}\_\{0,j\}^\{\\text\{test\}\}\}\(\\tfrac\{iT\}\{N\_\{T\}\}\),\\hat\{\\mu\}^\{\\text\{e\}\}\_\{\\bm\{x\}\_\{0,j\}^\{\\text\{test\}\}\}\(\\tfrac\{iT\}\{N\_\{T\}\}\)denotes the empirical distributions of the ground\-truth and reconstructed trajectories at timei​TNT\\tfrac\{iT\}\{N\_\{T\}\}whose initial conditions are𝒙0,jtest\\bm\{x\}\_\{0,j\}^\{\\text\{test\}\},ηa,𝒙0,je\\eta^\{\\text\{e\}\}\_\{a,\\bm\{x\}\_\{0,j\}\}andη^a,𝒙0,je\\hat\{\\eta\}^\{\\text\{e\}\}\_\{a,\\bm\{x\}\_\{0,j\}\}refers to the empirical distribution ofa⁡\(𝒙⁡\(t,𝒙0,jtest,ω\),ω\)a\(\\bm\{x\}\(t;\\bm\{x\}\_\{0,j\}^\{\\text\{test\}\};\\omega\),\\omega\)anda^​\(𝒙⁡\(t,𝒙0,jtest,ω^\),t,ω^\)\\hat\{a\}\(\\bm\{x\}\(t;\\bm\{x\}\_\{0,j\}^\{\\text\{test\}\};\\hat\{\\omega\}\),t;\\hat\{\\omega\}\), andηb,𝒙0,je\\eta^\{\\text\{e\}\}\_\{b,\\bm\{x\}\_\{0,j\}\}andη^b,𝒙0,je\\hat\{\\eta\}^\{\\text\{e\}\}\_\{b,\\bm\{x\}\_\{0,j\}\}refers to the empirical distribution ofb⁡\(𝒙⁡\(t,𝒙0,jtest,ω\),ω\)b\(\\bm\{x\}\(t;\\bm\{x\}\_\{0,j\}^\{\\text\{test\}\};\\omega\),\\omega\)andb^​\(𝒙⁡\(t,𝒙0,jtest,ω^\),t,ω^\)\\hat\{b\}\(\\bm\{x\}\(t;\\bm\{x\}\_\{0,j\}^\{\\text\{test\}\};\\hat\{\\omega\}\),t;\\hat\{\\omega\}\), respectively\.

As shown in Table[3](https://arxiv.org/html/2608.11613#S3.T3), using the Sinkhorn divergence is more computationally efficient, and the learned drift and diffusion functions are more accurate compared to those learned when using the temporally decoupled local squaredW2W\_\{2\}as the loss function\. The saving in runtime of using the temporal local Sinkhorn divergence Eq\. \([3\.22](https://arxiv.org/html/2608.11613#S3.E22)\) is not significant compared to using the temporal local squaredW2W\_\{2\}loss, and the possible reason is that the time needed to numerically solve the SDE is much longer than evaluating either the temporal local Sinkhorn divergence loss or the temporal localW2W\_\{2\}loss\.

Table 3:Comparison of the local Sinkhorn divergence versus the local squaredW2W\_\{2\}loss to train an SNN for reconstructing the FHN model Eq\. \([3\.19](https://arxiv.org/html/2608.11613#S3.E19)\)\.Loss FunctionTraining Time \(s\)DriftW22W\_\{2\}^\{2\}ErrorDiffusionW22W\_\{2\}^\{2\}ErrorTest trajectoryW22W\_\{2\}^\{2\}ErrorLocal Sinkhorn67720\.11100\.10900\.06415Local squaredW2W\_\{2\}72560\.11950\.17910\.06444

## 4Summary and conclusion

In this paper, we proposed a local Sinkhorn divergence framework for conditional distribution reconstruction of random field models using SNNs\. By utilizing the debiased Sinkhorn divergence, the proposed method preserved the geometric structure of optimal transport while introducing entropic regularization, leading to a differentiable and computationally efficient distribution matching objective\. Compared with our previous local squaredW2W\_\{2\}formulation, the proposed local Sinkhorn loss achieved improved computational efficiency and scalability while maintaining high reconstruction accuracy for learning multidimensional random field models\. We further established theoretical error bounds for the proposed local Sinkhorn loss\. The resulting bounds explicitly characterized the trade\-off between approximation bias and statistical efficiency introduced by the entropic regularization parameter and implied that the proposed local Sinkhorn loss could partially alleviate the curse of dimensionality when learning multidimensional random field models by properly choosing the regularization parameter\. Numerical experiments on one\-dimensional conditional distribution reconstruction, stochastic Darcy flow, and stochastic FHN systems demonstrated that the proposed method accurately reconstructed random field models and outperformed several prevailing machine\-learning uncertainty quantification benchmarks\. Overall, the proposed local Sinkhorn approach offered a practical balance between statistical accuracy, computational efficiency, and theoretical interpretability, making it a promising framework for uncertainty quantification and probabilistic scientific machine learning\.

Several promising research directions deserve further investigation\. First, the neighborhood radiusδ\\deltaand the regularization parameterε\\varepsilonare currently selected empirically and remain fixed\. Developing adaptive strategies for simultaneously choosing and adjusting these two parameters according to the local complexity of the conditional distribution may further improve both accuracy and computational efficiency\. Recent advances in stabilized and multiscale Sinkhorn algorithms provide promising directions for adaptive regularization and efficient optimal transport computations[schmitzer2019stabilized](https://arxiv.org/html/2608.11613#bib.bib23);[feydy2019interpolating](https://arxiv.org/html/2608.11613#bib.bib8)\. Second, the current framework constructs neighborhoods using Euclidean distances\. More general neighborhood constructions, including anisotropic neighborhoods, manifold\-based neighborhoods, and graph\-based local structures, may significantly improve the approximation of conditional distributions with low\-dimensional intrinsic geometry[belkin2003laplacian](https://arxiv.org/html/2608.11613#bib.bib3);[tenenbaum2000global](https://arxiv.org/html/2608.11613#bib.bib27)\. Such adaptive neighborhood constructions may also lead to improved theoretical convergence guarantees\. Another important direction is to extend the proposed local Sinkhorn framework to large\-scale scientific machine learning problems, including stochastic partial differential equations, neural operator learning, probabilistic surrogate modeling, and Bayesian scientific machine learning\. Since the proposed local Sinkhorn loss is fully differentiable and naturally compatible with GPU\-based optimization, it is particularly suitable for these large\-scale applications[kovachki2023neuraloperator](https://arxiv.org/html/2608.11613#bib.bib18);[raissi2019physics](https://arxiv.org/html/2608.11613#bib.bib21);[papamakarios2021normalizing](https://arxiv.org/html/2608.11613#bib.bib19)\. Finally, it would be of considerable interest to further investigate the theoretical properties of local Sinkhorn learning in multidimensional settings\.

## Declaration of generative AI and AI\-assisted technologies in the manuscript preparation process

During the preparation of this work, the authors used ChatGPT in order to polish the writing and assist in checking the manuscript\. After using this tool/service, the authors reviewed and edited the content as needed and take full responsibility for the content of the published article\.

## Declaration of competing interest

The authors declare that they have no known competing financial interests or personal relationships that could have appeared to influence the work reported in this paper\.

## Data statement

The data used in this study were generated numerically according to the models and experimental settings described in the manuscript\. The data and implementation code that support the findings of this study will be made publicly available upon publication of the article\.

## Appendix AProof to Theorem[2\.1](https://arxiv.org/html/2608.11613#S2.Thmtheorem1)

Here, we prove Theorem[2\.1](https://arxiv.org/html/2608.11613#S2.Thmtheorem1)\. We first decompose the generalization error as:

\|𝒮¯ε,δe​\(𝒚𝒙,𝒚^𝒙\)−𝒲¯2​\(𝒚𝒙,𝒚^𝒙\)\|≤\|𝒮¯ε,δe​\(𝒚𝒙,𝒚^𝒙\)−𝒲¯2,δe​\(𝒚𝒙,𝒚^𝒙\)\|\+\|𝒲¯2,δe​\(𝒚𝒙,𝒚^𝒙\)−𝒲¯22​\(𝒚𝒙,𝒚^𝒙\)\|,\\displaystyle\\left\|\\overline\{\\mathcal\{S\}\}\_\{\\varepsilon,\\delta\}^\{\\text\{e\}\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\-\\overline\{\\mathcal\{W\}\}\_\{2\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\\right\|\\leq\\,\\left\|\\overline\{\\mathcal\{S\}\}\_\{\\varepsilon,\\delta\}^\{\\text\{e\}\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\-\\overline\{\\mathcal\{W\}\}\_\{2,\\delta\}^\{\\text\{e\}\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\\right\|\+\\left\|\\overline\{\\mathcal\{W\}\}\_\{2,\\delta\}^\{\\text\{e\}\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\-\\overline\{\\mathcal\{W\}\}\_\{2\}^\{2\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\\right\|,\(A\.1\)where𝒲¯2​\(𝒚𝒙,𝒚^𝒙\)\\overline\{\\mathcal\{W\}\}\_\{2\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)is defined in Eq\. \([2\.9](https://arxiv.org/html/2608.11613#S2.E9)\), and𝒲¯2,δe​\(𝒚𝒙,𝒚^𝒙\)\\overline\{\\mathcal\{W\}\}\_\{2,\\delta\}^\{\\text\{e\}\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)is the local squaredW2W\_\{2\}distance between the two random fields𝒚𝒙\\bm\{y\}\_\{\\bm\{x\}\}and𝒚^𝒙\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}defined as:

𝒲¯2,δe​\(𝒚𝒙,𝒚^𝒙\)≔∫𝒟W22​\(μ𝒙,δe,μ^𝒙,δe\)​ν​\(d​𝒙\)\.\\overline\{\\mathcal\{W\}\}\_\{2,\\delta\}^\{\\text\{e\}\}\\big\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\\big\)\\coloneqq\\int\_\{\\mathcal\{D\}\}W\_\{2\}^\{2\}\\big\(\\mu\_\{\\bm\{x\},\\delta\}^\{\\text\{e\}\},\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}^\{\\text\{e\}\}\\big\)\\nu\(\\mbox\{d\}\\bm\{x\}\)\.\(A\.2\)
The first term in Eq\. \([A\.1](https://arxiv.org/html/2608.11613#A1.E1)\) is the Sinkhorn regularization bias\. By\([genevay2019sample](https://arxiv.org/html/2608.11613#bib.bib11), Theorem 1\)and the boundedness assumption \(Assumption[2\.1](https://arxiv.org/html/2608.11613#S2.Thmassumption1)\), taking the Lipschitz constant of the quadratic cost functionccin Eq\. \([2\.4](https://arxiv.org/html/2608.11613#S2.E4)\) as2​M02M\_\{0\}yields

0≤Wε​\(α,β\)−W⁡\(α,β\)≤2​ε​log⁡\(2​e2​M02d​ε\),0\\leq W\_\{\\varepsilon\}\(\\alpha,\\beta\)\-W\(\\alpha,\\beta\)\\leq 2\\varepsilon\\log\\\!\\left\(\\frac\{2e^\{2\}M\_\{0\}^\{2\}\}\{\\sqrt\{d\}\\,\\varepsilon\}\\right\),for every relevant pair of probability measures\(α,β\)\(\\alpha,\\beta\)\. Since the debiased Sinkhorn divergence is obtained by subtracting the self\-transport terms, the same bound controls the difference between the local Sinkhorn loss and itsW22W\_\{2\}^\{2\}counterpart, namely

\|𝒮¯ε,δe​\(𝒚𝒙,𝒚^𝒙\)−𝒲¯2,δe​\(𝒚𝒙,𝒚^𝒙\)\|≤2​ε​log⁡\(2​e2​M02d​ε\)\.\\left\|\\overline\{\\mathcal\{S\}\}\_\{\\varepsilon,\\delta\}^\{\\text\{e\}\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\-\\overline\{\\mathcal\{W\}\}\_\{2,\\delta\}^\{\\text\{e\}\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\\right\|\\leq 2\\varepsilon\\log\\\!\\left\(\\frac\{2e^\{2\}M\_\{0\}^\{2\}\}\{\\sqrt\{d\}\\,\\varepsilon\}\\right\)\.\(A\.3\)
The second term in Eq\. \([A\.1](https://arxiv.org/html/2608.11613#A1.E1)\) is exactly the local squaredW2W\_\{2\}generalization error\. Therefore, applying\([xia2025efficient](https://arxiv.org/html/2608.11613#bib.bib30), Theorem 2\.3\)in our previous work gives

𝔼⁡\[\|𝒲¯2,δe​\(𝒚𝒙,𝒚^𝒙\)−𝒲¯2​\(𝒚𝒙,𝒚^𝒙\)\|\]≤4​M0N\+8​C​M0​𝔼​\[h⁡\(N⁡\(x,δ\),d\)\]\+16​M02​δ\.\\mathbb\{E\}\\\!\\left\[\\left\|\\overline\{\\mathcal\{W\}\}\_\{2,\\delta\}^\{\\text\{e\}\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\-\\overline\{\\mathcal\{W\}\}\_\{2\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\\right\|\\right\]\\leq\\frac\{4M\_\{0\}\}\{\\sqrt\{N\}\}\+8CM\_\{0\}\\,\\mathbb\{E\}\\\!\\bigl\[h\(N\(x,\\delta\),d\)\\bigr\]\+16M\_\{0\}^\{2\}\\,\\delta\.
Combining the Sinkhorn bias bound in Eq\. \([A\.3](https://arxiv.org/html/2608.11613#A1.E3)\) with the localW2W\_\{2\}generalization bound above proves Eq\. \([2\.13](https://arxiv.org/html/2608.11613#S2.E13)\)\.

## Appendix BProof to Theorem[2\.2](https://arxiv.org/html/2608.11613#S2.Thmtheorem2)

Here, we prove Theorem[2\.2](https://arxiv.org/html/2608.11613#S2.Thmtheorem2)\. We decompose the generalization error into three parts:

E​\[\|𝒮¯ε,δe​\(𝒚𝒙,𝒚^𝒙\)−𝒲¯22​\(𝒚𝒙,𝒚^𝒙\)\|\]≤\\displaystyle\\mbox\{E\}\\bigg\[\\left\|\\overline\{\\mathcal\{S\}\}\_\{\\varepsilon,\\delta\}^\{\\text\{e\}\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\-\\overline\{\\mathcal\{W\}\}\_\{2\}^\{2\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\\right\|\\bigg\]\\leqE​\[\|𝒮¯ε,δe​\(𝒚𝒙,𝒚^𝒙\)−𝒮¯ε,δ​\(𝒚𝒙,𝒚^𝒙\)\|\]\\displaystyle\\mbox\{E\}\\bigg\[\\left\|\\overline\{\\mathcal\{S\}\}\_\{\\varepsilon,\\delta\}^\{\\text\{e\}\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\-\\overline\{\\mathcal\{S\}\}\_\{\\varepsilon,\\delta\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\\right\|\\bigg\]\(B\.1\)\+E​\[\|∫DSε​\(μ𝒙,δ,μ^𝒙,δ\)​νe​\(d​𝒙\)−∫DW22​\(μ𝒙,δ,μ^𝒙,δ\)​νe​\(d​𝒙\)\|\]\\displaystyle\+\\mbox\{E\}\\bigg\[\\left\|\\int\_\{D\}S\_\{\\varepsilon\}\(\\mu\_\{\\bm\{x\},\\delta\},\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}\)\\nu^\{\\text\{e\}\}\(\\mbox\{d\}\\bm\{x\}\)\-\\int\_\{D\}W\_\{2\}^\{2\}\(\\mu\_\{\\bm\{x\},\\delta\},\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}\)\\nu^\{\\text\{e\}\}\(\\mbox\{d\}\\bm\{x\}\)\\right\|\\bigg\]\+E​\[\|∫DW22​\(μ𝒙,δ,μ^𝒙,δ\)​νe​\(d​𝒙\)−∫DW22​\(μ𝒙,μ^𝒙\)​νe​\(d​𝒙\)\|\]\.\\displaystyle\+\\mbox\{E\}\\bigg\[\\left\|\\int\_\{D\}W\_\{2\}^\{2\}\(\\mu\_\{\\bm\{x\},\\delta\},\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}\)\\nu^\{\\text\{e\}\}\(\\mbox\{d\}\\bm\{x\}\)\-\\int\_\{D\}W\_\{2\}^\{2\}\(\\mu\_\{\\bm\{x\}\},\\hat\{\\mu\}\_\{\\bm\{x\}\}\)\\nu^\{\\text\{e\}\}\(\\mbox\{d\}\\bm\{x\}\)\\right\|\\bigg\]\.
The first term is the averaged empirical Sinkhorn error\. For each𝒙\\bm\{x\}, by applying\([genevay2019sample](https://arxiv.org/html/2608.11613#bib.bib11), Theorem 3\), we have:

𝔼⁡\[\|𝒮ε​\(μ𝒙,δe,μ^𝒙,δe\)−𝒮ε​\(μ𝒙,δ,μ^𝒙,δ\)\|\]≤𝒪⁡\(e5​M02εN⁡\(𝒙,δ\)​\(1\+ε−⌊d/2⌋\)\)\.\\mathbb\{E\}\\\!\\left\[\\left\|\\mathcal\{S\}\_\{\\varepsilon\}\(\\mu\_\{\\bm\{x\},\\delta\}^\{\\text\{e\}\},\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}^\{\\text\{e\}\}\)\-\\mathcal\{S\}\_\{\\varepsilon\}\(\\mu\_\{\\bm\{x\},\\delta\},\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}\)\\right\|\\right\]\\leq\\mathcal\{O\}\\Big\(\\frac\{e^\{\\frac\{5M\_\{0\}^\{2\}\}\{\\varepsilon\}\}\}\{\\sqrt\{N\(\\bm\{x\},\\delta\)\}\}\\Bigl\(1\+\\varepsilon^\{\-\\lfloor d/2\\rfloor\}\\Bigr\)\\Big\)\.\(B\.2\)
The second term is the average regularization bias between the Sinkhorn divergence and the Wasserstein discrepancy\. From\([genevay2019sample](https://arxiv.org/html/2608.11613#bib.bib11), Theorem 1\), for each𝒙\\bm\{x\}, we have

E​\[\|𝒮ε​\(μ𝒙,δ,μ^𝒙,δ\)−𝒲22​\(μ𝒙,δ,μ^𝒙,δ\)\|\]≤2​ε​log⁡\(2​e2​M02d​ε\)\.\\mbox\{E\}\\bigg\[\\left\|\\mathcal\{S\}\_\{\\varepsilon\}\(\\mu\_\{\\bm\{x\},\\delta\},\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}\)\-\\mathcal\{W\}^\{2\}\_\{2\}\(\\mu\_\{\\bm\{x\},\\delta\},\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}\)\\right\|\\bigg\]\\leq 2\\varepsilon\\log\\\!\\left\(\\frac\{2e^\{2\}M\_\{0\}^\{2\}\}\{\\sqrt\{d\}\\,\\varepsilon\}\\right\)\.\(B\.3\)Therefore,

E​\[\|∫DSε​\(μ𝒙,δ,μ^𝒙,δ\)​νe​\(d​𝒙\)−∫DW22​\(μ𝒙,δ,μ^𝒙,δ\)​νe​\(d​𝒙\)\|\]≤2​ε​log⁡\(2​e2​M02d​ε\)\.\\mbox\{E\}\\bigg\[\\left\|\\int\_\{D\}S\_\{\\varepsilon\}\(\\mu\_\{\\bm\{x\},\\delta\},\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}\)\\nu^\{\\text\{e\}\}\(\\mbox\{d\}\\bm\{x\}\)\-\\int\_\{D\}W\_\{2\}^\{2\}\(\\mu\_\{\\bm\{x\},\\delta\},\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}\)\\nu^\{\\text\{e\}\}\(\\mbox\{d\}\\bm\{x\}\)\\right\|\\bigg\]\\leq 2\\varepsilon\\log\\\!\\left\(\\frac\{2e^\{2\}M\_\{0\}^\{2\}\}\{\\sqrt\{d\}\\,\\varepsilon\}\\right\)\.\(B\.4\)Finally, for each𝒙\\bm\{x\}, using the triangular inequality of theW2W\_\{2\}distance[clement2008elementary](https://arxiv.org/html/2608.11613#bib.bib5), we have:

\|W2​\(μ𝒙,δ,μ^𝒙,δ\)−W2​\(μ𝒙,μ^𝒙\)\|≤W2​\(μ𝒙,δ,μ𝒙\)\+W2​\(μ^𝒙,μ^𝒙,δ\)\.\|W\_\{2\}\(\\mu\_\{\\bm\{x\},\\delta\},\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}\)\-W\_\{2\}\(\\mu\_\{\\bm\{x\}\},\\hat\{\\mu\}\_\{\\bm\{x\}\}\)\|\\leq W\_\{2\}\(\\mu\_\{\\bm\{x\},\\delta\},\{\\mu\}\_\{\\bm\{x\}\}\)\+W\_\{2\}\(\\hat\{\\mu\}\_\{\\bm\{x\}\},\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}\)\.\(B\.5\)
Following the proof of\([xia2026local](https://arxiv.org/html/2608.11613#bib.bib31), Theorem 2\.3\)and the Lipschitz continuity condition in Assumption[2\.1](https://arxiv.org/html/2608.11613#S2.Thmassumption1), we have:

W2​\(μ𝒙,δ,μ𝒙\)\+W2​\(μ^𝒙,μ^𝒙,δ\)≤4​M0​δ\.W\_\{2\}\(\\mu\_\{\\bm\{x\},\\delta\},\{\\mu\}\_\{\\bm\{x\}\}\)\+W\_\{2\}\(\\hat\{\\mu\}\_\{\\bm\{x\}\},\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}\)\\leq 4M\_\{0\}\\delta\.\(B\.6\)Therefore,

\|W22\(μ𝒙,δ,μ^𝒙,δ\)−W22\(μ𝒙,μ^𝒙\)\\displaystyle\|W\_\{2\}^\{2\}\(\\mu\_\{\\bm\{x\},\\delta\},\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}\)\-W\_\{2\}^\{2\}\(\\mu\_\{\\bm\{x\}\},\\hat\{\\mu\}\_\{\\bm\{x\}\}\)\|≤\|W2\(μ𝒙,δ,μ^𝒙,δ\)−W2\(μ𝒙,μ^𝒙\)\|⋅\|W2\(μ𝒙,δ,μ^𝒙,δ\)\+W2\(μ𝒙,μ^𝒙\)\|\\displaystyle\|\\leq\|W\_\{2\}\(\\mu\_\{\\bm\{x\},\\delta\},\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}\)\-W\_\{2\}\(\\mu\_\{\\bm\{x\}\},\\hat\{\\mu\}\_\{\\bm\{x\}\}\)\|\\cdot\|W\_\{2\}\(\\mu\_\{\\bm\{x\},\\delta\},\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}\)\+W\_\{2\}\(\\mu\_\{\\bm\{x\}\},\\hat\{\\mu\}\_\{\\bm\{x\}\}\)\|\(B\.7\)≤2​\(2​M0\)​δ​\(2​sup‖y‖2\+2​sup‖y^‖2\)=16​M02​δ,\\displaystyle\\leq 2\(2M\_\{0\}\)\\delta\(2\\sup\\\|y\\\|\_\{2\}\+2\\sup\\\|\\hat\{y\}\\\|\_\{2\}\)=16M\_\{0\}^\{2\}\\delta,which leads to:

E​\[\|W22​\(μ𝒙,δ,μ^𝒙,δ\)−W22​\(μ𝒙,μ^𝒙\)\|\]≤16​M02​δ\\mbox\{E\}\\Big\[\\big\|W^\{2\}\_\{2\}\(\\mu\_\{\\bm\{x\},\\delta\},\\hat\{\\mu\}\_\{\\bm\{x\},\\delta\}\)\-W\_\{2\}^\{2\}\(\\mu\_\{\\bm\{x\}\},\\hat\{\\mu\}\_\{\\bm\{x\}\}\)\\big\|\\Big\]\\leq 16M\_\{0\}^\{2\}\\delta\(B\.8\)
Combining Eqs\. \([B\.8](https://arxiv.org/html/2608.11613#A2.E8)\), \([B\.1](https://arxiv.org/html/2608.11613#A2.E1)\), \([B\.2](https://arxiv.org/html/2608.11613#A2.E2)\), and \([B\.4](https://arxiv.org/html/2608.11613#A2.E4)\) yields Eq\. \([2\.17](https://arxiv.org/html/2608.11613#S2.E17)\)\.

Remark\.The error bound induced from using the neighborhood technique Eq\. \([B\.6](https://arxiv.org/html/2608.11613#A2.E6)\) only applies to the Wasserstein distance through devising a special coupling ofμ𝒙,δ\\mu\_\{\\bm\{x\},\\delta\}andμ𝒙\\mu\_\{\\bm\{x\}\}\. However, there is no direct error bound forSε​\(μ𝒙,δ,μ𝒙\)S\_\{\\varepsilon\}\(\\mu\_\{\\bm\{x\},\\delta\},\\mu\_\{\\bm\{x\}\}\)as the Sinkhorn divergence introduces the additional entropic Kullback\-Leibler divergence term\. Therefore, it is nontrivial to obtain a direct error bound of

\|𝒮¯ε,δe​\(𝒚𝒙,𝒚^𝒙\)−∫𝒟Sε​\(μ𝒙,μ^𝒙\)​ν​\(d​𝒙\)\|,\\left\|\\overline\{\\mathcal\{S\}\}\_\{\\varepsilon,\\delta\}^\{\\text\{e\}\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)\-\\int\_\{\\mathcal\{D\}\}S\_\{\\varepsilon\}\(\\mu\_\{\\bm\{x\}\},\\hat\{\\mu\}\_\{\\bm\{x\}\}\)\\nu\(\\mbox\{d\}\\bm\{x\}\)\\right\|,\(B\.9\)where∫𝒟Sε​\(μ𝒙,μ^𝒙\)​ν​\(d​𝒙\)\\int\_\{\\mathcal\{D\}\}S\_\{\\varepsilon\}\(\\mu\_\{\\bm\{x\}\},\\hat\{\\mu\}\_\{\\bm\{x\}\}\)\\nu\(\\mbox\{d\}\\bm\{x\}\)denotes the average Sinkhorn divergence between𝒚𝒙\\bm\{y\}\_\{\\bm\{x\}\}and𝒚^𝒙\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\.

## Appendix COptimization, training settings and hyperparameters

All numerical experiments are implemented in Python 3\.11 using the PyTorch deep learning framework\. The proposed Local Sinkhorn divergence is implemented using the GeomLoss package, while theW2W\_\{2\}distance is computed using the POT package\. Unless otherwise specified, all competing methods within each numerical example employ the same neural network architecture, optimizer, initialization strategy, and training settings\. All experiments are performed on a desktop equipped with an Intel Core i9\-13900KF CPU \(24 cores\) and 64 GB RAM\. The runtime reported in the numerical examples corresponds to the actual wall\-clock training time under the same hardware environment\. Training settings and hyperparameters are listed in Table[4](https://arxiv.org/html/2608.11613#A3.T4)\. In this work, theW2W\_\{2\}distance is evaluated numerically using theot\.emd2function in thePOTpackage[flamary2021pot](https://arxiv.org/html/2608.11613#bib.bib9)and the Sinkhorn divergence is evaluated numerically using theSamplesLossfunction in theGeomlosspackage[feydy2019interpolating](https://arxiv.org/html/2608.11613#bib.bib8)\.

Table 4:Training hyperparameters, neural network settings, and optimization configurations for all numerical examples\.Example[1](https://arxiv.org/html/2608.11613#Thmexample1)Example[2](https://arxiv.org/html/2608.11613#Thmexample2)Example[3](https://arxiv.org/html/2608.11613#Thmexample3)OptimizerAdamAdamAdamLearning rate5×10−35\\times 10^\{\-3\}5×10−45\\times 10^\{\-4\}1×10−21\\times 10^\{\-2\}Weight decay000Number of epochsEE200020000200Training samples20002000300Number of testing centersNcN\_\{c\}1002020Testing Realizations2010020Mini\-batch sizenbn\_\{b\}848Hidden layers232Hidden neurons3212832Activation functionReLUReLUGeLUResidual connection \(ResNet\)YesYesYesNeighborhood radiusδ\\delta0\.050\.050\.1Minimal samples required for each neighborhoodnminn\_\{\\min\}444Maximal samples recorded in each neighborhoodnmaxn\_\{\\max\}12864300Samples chosen in each neighborhoodnn323232Regularization coefficientε\\varepsilon0\.050\.030\.1SNN parameter initialization𝒩⁡\(0,0\.012\)\\mathcal\{N\}\(0,0\.01^\{2\}\)𝒩⁡\(0,0\.052\)\\mathcal\{N\}\(0,0\.05^\{2\}\)𝒩⁡\(0,0\.012\)\\mathcal\{N\}\(0,0\.01^\{2\}\)Unless otherwise specified, all model parameters are initialized using the default PyTorch initialization strategy\. The same initialization scheme, optimizer, learning rate, training epochs, and network architecture are adopted for all competing methods within each numerical example to ensure a fair comparison\.

## Appendix DDefinitions of Different Loss Metrics

Here, we summarize the loss functions used in Examples[1](https://arxiv.org/html/2608.11613#Thmexample1)and[2](https://arxiv.org/html/2608.11613#Thmexample2)\. Let

B⁡\(𝒙i,δ\)=\{𝒙j:‖𝒙j−𝒙i‖≤δ\}B\(\\bm\{x\}\_\{i\},\\delta\)=\\\{\\bm\{x\}\_\{j\}:\\\|\\bm\{x\}\_\{j\}\-\\bm\{x\}\_\{i\}\\\|\\leq\\delta\\\}denote the neighborhood centered at𝒙i\\bm\{x\}\_\{i\}\(Eq\. \([2\.10](https://arxiv.org/html/2608.11613#S2.E10)\)\)\. If a center𝒙i\\bm\{x\}\_\{i\}is selected when evaluating the loss function within an epoch, a subsetℬ⁡\(𝒙i,δ\)⊆B⁡\(𝒙i,δ\)\\mathcal\{B\}\(\\bm\{x\}\_\{i\},\\delta\)\\subseteq B\(\\bm\{x\}\_\{i\},\\delta\)is randomly chosen with\|ℬ⁡\(𝒙i,δ\)\|=min⁡\(n,\|B⁡\(𝒙i,δ\)\|\)\|\\mathcal\{B\}\(\\bm\{x\}\_\{i\},\\delta\)\|=\\min\(n,\|B\(\\bm\{x\}\_\{i\},\\delta\)\|\)\. The overall loss is obtained by averaging the corresponding local loss over all randomly selected neighborhood centers\.

- •1\. Local mean squared error \(Local MSE\) MSEδ​\(𝒚𝒙,𝒚^𝒙\)=1nb​∑i=1nb1\|ℬ⁡\(𝒙i,δ\)\|​∑𝒙j∈ℬ⁡\(𝒙i,δ\)‖y⁡\(𝒙j,ωj\)−𝒚^​\(𝒙j,ω^j\)‖2\.\\mathrm\{MSE\}\_\{\\delta\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)=\\frac\{1\}\{n\_\{b\}\}\\sum\_\{i=1\}^\{n\_\{b\}\}\\frac\{1\}\{\|\\mathcal\{B\}\(\\bm\{x\}\_\{i\},\\delta\)\|\}\\sum\_\{\\bm\{x\}\_\{j\}\\in\\mathcal\{B\}\(\\bm\{x\}\_\{i\},\\delta\)\}\\\|y\(\\bm\{x\}\_\{j\};\\omega\_\{j\}\)\-\\hat\{\\bm\{y\}\}\(\\bm\{x\}\_\{j\};\\hat\{\\omega\}\_\{j\}\)\\\|^\{2\}\.\(D\.1\)
- •2\. Local mean absolute error \(Local MAE\) MAEδ​\(𝒚𝒙,𝒚^𝒙\)=1nb​∑i=1nb1\|ℬ⁡\(𝒙i,δ\)\|​∑𝒙j∈ℬ⁡\(𝒙i,δ\)‖𝒚⁡\(𝒙j,ωj\)−𝒚^​\(𝒙j,ω^j\)‖\.\\mathrm\{MAE\}\_\{\\delta\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)=\\frac\{1\}\{n\_\{b\}\}\\sum\_\{i=1\}^\{n\_\{b\}\}\\frac\{1\}\{\|\\mathcal\{B\}\(\\bm\{x\}\_\{i\},\\delta\)\|\}\\sum\_\{\\bm\{x\}\_\{j\}\\in\\mathcal\{B\}\(\\bm\{x\}\_\{i\},\\delta\)\}\\\|\\bm\{y\}\(\\bm\{x\}\_\{j\};\\omega\_\{j\}\)\-\\hat\{\\bm\{y\}\}\(\\bm\{x\}\_\{j\};\\hat\{\\omega\}\_\{j\}\)\\\|\.\(D\.2\)
- •3\. Local energy distance \(Local ED\) The energy distance follows from[szekely2013energy](https://arxiv.org/html/2608.11613#bib.bib26)and is defined locally by EDδ​\(𝒚𝒙,𝒚^𝒙\)=1nb​∑i=1nbED⁡\(Yiδ,Y^iδ\),\\mathrm\{ED\}\_\{\\delta\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)=\\frac\{1\}\{n\_\{b\}\}\\sum\_\{i=1\}^\{n\_\{b\}\}\\mathrm\{ED\}\\\!\\left\(Y\_\{i\}^\{\\delta\},\\hat\{Y\}\_\{i\}^\{\\delta\}\\right\),\(D\.3\) where ED⁡\(Yiδ,Y^iδ\)\\displaystyle\\mathrm\{ED\}\\\!\\left\(Y\_\{i\}^\{\\delta\},\\hat\{Y\}\_\{i\}^\{\\delta\}\\right\)=2\|ℬ⁡\(𝒙i,δ\)\|2​∑𝒚p∈Yiδ,𝒚^q∈Y^iδ‖𝒚p−𝒚^q‖\\displaystyle=\\frac\{2\}\{\|\\mathcal\{B\}\(\\bm\{x\}\_\{i\},\\delta\)\|^\{2\}\}\\sum\_\{\\bm\{y\}\_\{p\}\\in Y\_\{i\}^\{\\delta\},\\hat\{\\bm\{y\}\}\_\{q\}\\in\\hat\{Y\}\_\{i\}^\{\\delta\}\}\\\|\\bm\{y\}\_\{p\}\-\\hat\{\\bm\{y\}\}\_\{q\}\\\|\(D\.4\)−1\|ℬ⁡\(𝒙i,δ\)\|2∑𝒚p∈Yiδ,𝒚^q∈Y^iδ∥𝒚p−𝒚q∥−1\|ℬ⁡\(𝒙i,δ\)\|2∑𝒚p∈Yiδ,𝒚^q∈Y^iδ∥𝒚^p−𝒚^q∥,\\displaystyle\-\\frac\{1\}\{\|\\mathcal\{B\}\(\\bm\{x\}\_\{i\},\\delta\)\|^\{2\}\}\\sum\_\{\\bm\{y\}\_\{p\}\\in Y\_\{i\}^\{\\delta\},\\hat\{\\bm\{y\}\}\_\{q\}\\in\\hat\{Y\}\_\{i\}^\{\\delta\}\}\\\|\\bm\{y\}\_\{p\}\-\\bm\{y\}\_\{q\}\\\|\-\\frac\{1\}\{\|\\mathcal\{B\}\(\\bm\{x\}\_\{i\},\\delta\)\|^\{2\}\}\\sum\_\{\\bm\{y\}\_\{p\}\\in Y\_\{i\}^\{\\delta\},\\hat\{\\bm\{y\}\}\_\{q\}\\in\\hat\{Y\}\_\{i\}^\{\\delta\}\}\\\|\\hat\{\\bm\{y\}\}\_\{p\}\-\\hat\{\\bm\{y\}\}\_\{q\}\\\|,with Yiδ=\{𝒚⁡\(𝒙j,ωj\):𝒙j∈ℬ⁡\(𝒙i,δ\)\},Y^iδ=\{𝒚^​\(𝒙j,ω^j\):𝒙j∈ℬ⁡\(𝒙i,δ\)\}\.Y\_\{i\}^\{\\delta\}=\\left\\\{\\bm\{y\}\(\\bm\{x\}\_\{j\};\\omega\_\{j\}\):\\bm\{x\}\_\{j\}\\in\\mathcal\{B\}\(\\bm\{x\}\_\{i\},\\delta\)\\right\\\},\\qquad\\hat\{Y\}\_\{i\}^\{\\delta\}=\\left\\\{\\hat\{\\bm\{y\}\}\(\\bm\{x\}\_\{j\};\\hat\{\\omega\}\_\{j\}\):\\bm\{x\}\_\{j\}\\in\\mathcal\{B\}\(\\bm\{x\}\_\{i\},\\delta\)\\right\\\}\.\(D\.5\)
- •4\. Local maximum mean discrepancy \(Local MMD\) The MMD follows from the kernel two\-sample test[gretton2012kernel](https://arxiv.org/html/2608.11613#bib.bib15)and is defined locally by: MMDδ​\(𝒚𝒙,𝒚^𝒙\)=1\|Γ\|​∑γ∈Γ1nb​∑i=1nbMMDγ​\(Yiδ,Y^iδ\),\\mathrm\{MMD\}\_\{\\delta\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)=\\frac\{1\}\{\|\\Gamma\|\}\\sum\_\{\\gamma\\in\\Gamma\}\\frac\{1\}\{n\_\{b\}\}\\sum\_\{i=1\}^\{n\_\{b\}\}\\mathrm\{MMD\}\_\{\\gamma\}\\\!\\left\(Y\_\{i\}^\{\\delta\},\\hat\{Y\}\_\{i\}^\{\\delta\}\\right\),\(D\.6\)where MMDγ​\(Yiδ,Y^iδ\)\\displaystyle\\mathrm\{MMD\}\_\{\\gamma\}\\\!\\left\(Y\_\{i\}^\{\\delta\},\\hat\{Y\}\_\{i\}^\{\\delta\}\\right\)=1\|ℬ⁡\(𝒙i,δ\)\|2​∑𝒚p∈Yiδ,𝒚^q∈Y^iδKγ​\(𝒚p,𝒚q\)\+1\|ℬ⁡\(𝒙i,δ\)\|2​∑𝒚p∈Yiδ,𝒚^q∈Y^iδKγ​\(𝒚^p,𝒚^q\)\\displaystyle=\\frac\{1\}\{\|\\mathcal\{B\}\(\\bm\{x\}\_\{i\},\\delta\)\|^\{2\}\}\\sum\_\{\\bm\{y\}\_\{p\}\\in Y\_\{i\}^\{\\delta\},\\hat\{\\bm\{y\}\}\_\{q\}\\in\\hat\{Y\}\_\{i\}^\{\\delta\}\}K\_\{\\gamma\}\(\\bm\{y\}\_\{p\},\\bm\{y\}\_\{q\}\)\+\\frac\{1\}\{\|\\mathcal\{B\}\(\\bm\{x\}\_\{i\},\\delta\)\|^\{2\}\}\\sum\_\{\\bm\{y\}\_\{p\}\\in Y\_\{i\}^\{\\delta\},\\hat\{\\bm\{y\}\}\_\{q\}\\in\\hat\{Y\}\_\{i\}^\{\\delta\}\}K\_\{\\gamma\}\(\\hat\{\\bm\{y\}\}\_\{p\},\\hat\{\\bm\{y\}\}\_\{q\}\)\(D\.7\)−2\|ℬ⁡\(𝒙i,δ\)\|2∑𝒚p∈Yiδ,𝒚^q∈Y^iδKγ\(𝒚p,𝒚^q\),\\displaystyle\-\\frac\{2\}\{\|\\mathcal\{B\}\(\\bm\{x\}\_\{i\},\\delta\)\|^\{2\}\}\\sum\_\{\\bm\{y\}\_\{p\}\\in Y\_\{i\}^\{\\delta\},\\hat\{\\bm\{y\}\}\_\{q\}\\in\\hat\{Y\}\_\{i\}^\{\\delta\}\}K\_\{\\gamma\}\(\\bm\{y\}\_\{p\},\\hat\{\\bm\{y\}\}\_\{q\}\),whereYiδY\_\{i\}^\{\\delta\}andY^iδ\\hat\{Y\}\_\{i\}^\{\\delta\}are defined in Eq\. \([D\.5](https://arxiv.org/html/2608.11613#A4.E5)\), andK⁡\(⋅,⋅\)K\(\\cdot,\\cdot\)denotes the Gaussian kernel: Kγ​\(𝒚,𝒚^\)=exp⁡\(−‖𝒚−𝒚^‖22​γ2\)\.K\_\{\\gamma\}\(\\bm\{y\},\\hat\{\\bm\{y\}\}\)=\\exp\\\!\\left\(\-\\frac\{\\\|\\bm\{y\}\-\\hat\{\\bm\{y\}\}\\\|^\{2\}\}\{2\\gamma^\{2\}\}\\right\)\.In Examples[1](https://arxiv.org/html/2608.11613#Thmexample1)and[2](https://arxiv.org/html/2608.11613#Thmexample2),Γ=\{0\.5,1\.0,2\.0,4\.0\}\\Gamma=\\\{0\.5,1\.0,2\.0,4\.0\\\}
- •5\. Local squaredW2W\_\{2\}distance W2,δ2​\(𝒚𝒙,𝒚^𝒙\)=1nb​∑i=1nbW22​\(μ𝒙i,δe,μ^𝒙i,δe\),W\_\{2,\\delta\}^\{2\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)=\\frac\{1\}\{n\_\{b\}\}\\sum\_\{i=1\}^\{n\_\{b\}\}W\_\{2\}^\{2\}\(\\mu\_\{\\bm\{x\}\_\{i\},\\delta\}^\{\\text\{e\}\},\\hat\{\\mu\}\_\{\\bm\{x\}\_\{i\},\\delta\}^\{\\text\{e\}\}\),\(D\.8\)whereμ𝒙i,δe\\mu\_\{\\bm\{x\}\_\{i\},\\delta\}^\{\\text\{e\}\}andμ^𝒙i,δe\\hat\{\\mu\}\_\{\\bm\{x\}\_\{i\},\\delta\}^\{\\text\{e\}\}are the empirical distributions associated with the samples𝒚⁡\(𝒙j,ωj\)\\bm\{y\}\(\\bm\{x\}\_\{j\};\\omega\_\{j\}\)inYiδY\_\{i\}^\{\\delta\}and𝒚^​\(𝒙j,ω^j\)\\hat\{\\bm\{y\}\}\(\\bm\{x\}\_\{j\};\\hat\{\\omega\}\_\{j\}\)inY^iδ\\hat\{Y\}\_\{i\}^\{\\delta\}defined in Eq\. \([D\.5](https://arxiv.org/html/2608.11613#A4.E5)\), respectively\.
- •6\. Local Sinkhorn divergence This loss is adopted throughout the numerical experiments unless otherwise specified\. Sε,δ​\(𝒚𝒙,𝒚^𝒙\)=1nb​∑i=1nbSε​\(μ𝒙i,δe,μ^𝒙i,δe\),S\_\{\\varepsilon,\\delta\}\(\\bm\{y\}\_\{\\bm\{x\}\},\\hat\{\\bm\{y\}\}\_\{\\bm\{x\}\}\)=\\frac\{1\}\{n\_\{b\}\}\\sum\_\{i=1\}^\{n\_\{b\}\}S\_\{\\varepsilon\}\(\\mu\_\{\\bm\{x\}\_\{i\},\\delta\}^\{\\text\{e\}\},\\hat\{\\mu\}\_\{\\bm\{x\}\_\{i\},\\delta\}^\{\\text\{e\}\}\),\(D\.9\)whereSε​\(μ,μ^\)S\_\{\\varepsilon\}\(\\mu,\\hat\{\\mu\}\)is the Sinkhorn divergence defined in Eq\. \([2\.6](https://arxiv.org/html/2608.11613#S2.E6)\), andμ𝒙i,δe\\mu\_\{\\bm\{x\}\_\{i\},\\delta\}^\{\\text\{e\}\}andμ^𝒙i,δe\\hat\{\\mu\}\_\{\\bm\{x\}\_\{i\},\\delta\}^\{\\text\{e\}\}are the empirical distributions associated with the samples𝒚⁡\(𝒙j,ωj\)\\bm\{y\}\(\\bm\{x\}\_\{j\};\\omega\_\{j\}\)inYiδY\_\{i\}^\{\\delta\}and𝒚^​\(𝒙j,ω^j\)\\hat\{\\bm\{y\}\}\(\\bm\{x\}\_\{j\};\\hat\{\\omega\}\_\{j\}\)inY^iδ\\hat\{Y\}\_\{i\}^\{\\delta\}defined in Eq\. \([D\.5](https://arxiv.org/html/2608.11613#A4.E5)\), respectively\.

## References

- \[1\]Juan A\. Acebrón, Adi R\. Bulsara, and W\.\-J\. Rappel\.Noisy FitzHugh–Nagumo model: From single elements to globally coupled networks\.Phys\. Rev\. E, 69:026202, 2004\.
- \[2\]Martin Arjovsky, Soumith Chintala, and Léon Bottou\.Wasserstein GAN\.2017\.
- \[3\]Mikhail Belkin and Partha Niyogi\.Laplacian eigenmaps for dimensionality reduction and data representation\.Neural Comput\., 15\(6\):1373–1396, 2003\.
- \[4\]Christopher M\. Bishop\.Mixture density networks\.1994\.
- \[5\]Philippe Clement and Wolfgang Desch\.An elementary proof of the triangle inequality for the Wasserstein metric\.Proc\. Am\. Math\. Soc\., 136\(1\):333–339, 2008\.
- \[6\]Marco Cuturi\.Sinkhorn distances: Lightspeed computation of optimal transport\.volume 26, 2013\.
- \[7\]Laurent Dinh et al\.Density estimation using real nvp\.InICLR, 2017\.
- \[8\]Jean Feydy, Thibault Séjourné, François\-Xavier Vialard, Shun\-ichi Amari, Alain Trouvé, and Gabriel Peyré\.Interpolating between optimal transport and maximum mean discrepancy using Sinkhorn divergences\.InProceedings of the 22nd International Conference on Artificial Intelligence and Statistics, pages 2681–2690, 2019\.
- \[9\]Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z\. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T\.H\. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J\. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer\.POT: Python optimal transport\.J\. Mach\. Learn\. Res\., 22\(78\):1–8, 2021\.
- \[10\]Charlie Frogner, Chiyuan Zhang, Hossein Mobahi, Mauricio Araya, and Tomaso Poggio\.Learning with a Wasserstein loss\.volume 28, 2015\.
- \[11\]Aude Genevay, Lénaïc Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré\.Sample complexity of Sinkhorn divergences\.InProceedings of the 22nd International Conference on Artificial Intelligence and Statistics, pages 1574–1583, 2019\.
- \[12\]Aude Genevay, Gabriel Peyré, and Marco Cuturi\.Learning generative models with sinkhorn divergences\.2018\.
- \[13\]Roger G\. Ghanem and Pol D\. Spanos\.Stochastic Finite Elements: A Spectral Approach\.Springer, 1991\.
- \[14\]Ian Goodfellow et al\.Generative adversarial nets\.volume 27, 2014\.
- \[15\]Arthur Gretton, Karsten M\. Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alexander J\. Smola\.A kernel two\-sample test\.J\. Mach\. Learn\. Res\., 13:723–773, 2012\.
- \[16\]Jonathan Ho, Ajay Jain, and Pieter Abbeel\.Denoising diffusion probabilistic models\.Advances in neural information processing systems, 33:6840–6851, 2020\.
- \[17\]Diederik P\. Kingma and Max Welling\.Auto\-encoding variational Bayes\.2014\.
- \[18\]Nikola Kovachki et al\.Neural operator: Learning maps between function spaces\.J\. Mach\. Learn\. Res\., 24:1–97, 2023\.
- \[19\]George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan\.Normalizing flows for probabilistic modeling and inference\.J\. Mach\. Learn\. Res\., 22:1–64, 2021\.
- \[20\]Gabriel Peyré and Marco Cuturi\.Computational Optimal Transport\.Now Publishers, 2019\.
- \[21\]Maziar Raissi, Paris Perdikaris, and George E\. Karniadakis\.Physics\-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations\.J\. Comput\. Phys\., 378:686–707, 2019\.
- \[22\]Danilo Rezende and Shakir Mohamed\.Variational inference with normalizing flows\.InICML, 2015\.
- \[23\]Bernhard Schmitzer\.Stabilized sparse scaling algorithms for entropy regularized transport problems\.SIAM J\. Sci\. Comput\., 41\(3\):A1443–A1481, 2019\.
- \[24\]Kihyuk Sohn et al\.Learning structured output representation using deep conditional generative models\.InNeurIPS, 2015\.
- \[25\]Timothy J\. Sullivan\.Introduction to Uncertainty Quantification\.Springer, 2015\.
- \[26\]Gábor J\. Székely and Maria L\. Rizzo\.Energy statistics: A class of statistics based on distances\.J\. Stat\. Plan\. Inference, 143\(8\):1249–1272, 2013\.
- \[27\]Joshua B\. Tenenbaum, Vin de Silva, and John C\. Langford\.A global geometric framework for nonlinear dimensionality reduction\.Science, 290\(5500\):2319–2323, 2000\.
- \[28\]Henry C\. Tuckwell and Roger Rodriguez\.Analytical and simulation results for stochastic FitzHugh–Nagumo neurons and neural networks\.J\. Comput\. Neurosci\., 5\(1\):91–113, 1998\.
- \[29\]Cédric Villani\.Optimal Transport: Old and New\.Springer, 2009\.
- \[30\]Mingtao Xia and Qijing Shen\.Efficient reconstruction of multidimensional random field models with heterogeneous data using stochastic neural networks\.arXiv preprint arXiv:2511\.13977, 2025\.
- \[31\]Mingtao Xia and Qijing Shen\.A local squared Wasserstein\-2 method for efficient reconstruction of models with uncertainty\.Mach\. Learn\.: Sci\. Technol\., 7\(3\):035001, 2026\.
- \[32\]Mingtao Xia, Qijing Shen, Philip Maini, Eamonn Gaffney, and Alex Mogilner\.A new local time\-decoupled squared Wasserstein\-2 method for training stochastic neural networks to reconstruct uncertain parameters in dynamical systems\.Neural Netw\., page 107893, 2025\.
- \[33\]Dongbin Xiu\.Numerical Methods for Stochastic Computations\.Princeton University Press, 2010\.