# A local squared Wasserstein-2 method for efficient reconstruction of models with uncertainty

arXiv: https://arxiv.org/abs/2406.06825 · HTML: https://mlsanity.com/paper/2406.06825/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 62 references that are in the map.
- First posted on arXiv 10 Jun 2024.

## Builds on (most cited first, 22 of 22)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Improved Training of Wasserstein GANs](https://mlsanity.com/paper/1704.00028/) (2017)
- [Diffusion Models in Vision: A Survey](https://mlsanity.com/paper/2209.04747/) (2022)
- [Data-driven Distributionally Robust Optimization Using the Wasserstein Metric: Performance Guarantees and Tractable Reformulations](https://mlsanity.com/paper/1505.05116/) (2015)
- [Train longer, generalize better: closing the generalization gap in large batch training of neural networks](https://mlsanity.com/paper/1705.08741/) (2017)
- [Generative Moment Matching Networks](https://mlsanity.com/paper/1502.02761/) (2015)
- [Generalized Sliced Wasserstein Distances](https://mlsanity.com/paper/1902.00434/) (2019)
- [Learning with a Wasserstein Loss](https://mlsanity.com/paper/1506.05439/) (2015)
- [Neural SDEs as Infinite-Dimensional GANs](https://mlsanity.com/paper/2102.03657/) (2021)
- [Bayesian Neural Networks: An Introduction and Survey](https://mlsanity.com/paper/2006.12024/) (2020)
- [Machine learning with data assimilation and uncertainty quantification for dynamical systems: a review](https://mlsanity.com/paper/2303.10462/) (2023)
- [Conditional Wasserstein Distances with Applications in Bayesian OT Flow Matching](https://mlsanity.com/paper/2403.18705/) (2024)
- [Wasserstein Generative Adversarial Uncertainty Quantification in Physics-Informed Neural Networks](https://mlsanity.com/paper/2108.13054/) (2021)
- [Wasserstein distributional robustness of neural networks](https://mlsanity.com/paper/2306.09844/) (2023)
- [Learning to solve inverse problems using Wasserstein loss](https://mlsanity.com/paper/1710.10898/) (2017)
- [Bridging POMDPs and Bayesian decision making for robust maintenance planning under model uncertainty: An application to railway systems](https://mlsanity.com/paper/2212.07933/) (2022)
- [The quadratic Wasserstein metric for inverse data matching](https://mlsanity.com/paper/1911.06911/) (2019)
- [An efficient Wasserstein-distance approach for reconstructing jump-diffusion processes using parameterized neural networks](https://mlsanity.com/paper/2406.01653/) (2024)
- [Implementation of batched Sinkhorn iterations for entropy-regularized Wasserstein loss](https://mlsanity.com/paper/1907.01729/) (2019)
- [Squared Wasserstein-2 Distance for Efficient Reconstruction of Stochastic Differential Equations](https://mlsanity.com/paper/2401.11354/) (2024)

## Built on by (most cited first, 2 of 2)

- [A generalized Wasserstein-2 distance approach for efficient reconstruction of random field models using stochastic neural networks](https://mlsanity.com/paper/2507.05143/) (2025)
- [A Local Sinkhorn Framework for Conditional Distribution Reconstruction of Multidimensional Random Fields](https://mlsanity.com/paper/2608.11613/) (2026)

## Abstract (arXiv)

In this paper, we propose a local squared Wasserstein-2 (W_2) method to solve the inverse problem of reconstructing models with uncertain latent variables or parameters. A key advantage of our approach is that it does not require prior information on the distribution of the latent variables or parameters in the underlying models. Instead, our method can efficiently reconstruct the distributions of the output associated with different inputs based on empirical distributions of observation data. We demonstrate the effectiveness of our proposed method across several uncertainty quantification (UQ) tasks, including linear regression with coefficient uncertainty, training neural networks with weight uncertainty, and reconstructing ordinary differential equations (ODEs) with a latent random variable.