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

Journal article

In this paper, we propose a local squared Wasserstein-2 (W2) 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, through minimizing our proposed local squared W2 loss function, linear regression models or neural networks can be directly trained to 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 tasks, including linear regression with coefficient uncertainty, training neural networks with weight uncertainty, and reconstructing ordinary differential equations with a latent random variable.

Authors: Xia, M; Shen, Q

• Publication status: Published
• Peer review status: Peer reviewed
• Publisher: IOP Publishing
• Journal: Machine Learning: Science and Technology
• Volume: 7
• Issue: 3
• Pages: 035001
• Article number: 035001
• Publication date: 2026-04-24
• Acceptance date: 2026-04-07
• DOI: 10.1088/2632-2153/ae5c59
• EISSN: 2632-2153
• ISSN: 2632-2153
• Language: English
• Keywords: neural network; linear regression; Wasserstein distance; uncertainty quantification; inverse problem