Neural Local Wasserstein Regression
This work addresses distributional regression for applications involving probability measures, offering a novel localized approach to overcome limitations of existing global or linearized methods.
The paper tackles the problem of distribution-on-distribution regression, where both predictors and responses are probability measures, by proposing Neural Local Wasserstein Regression, a flexible nonparametric framework that models regression through locally defined transport maps in Wasserstein space, and demonstrates its effectiveness in capturing nonlinear and high-dimensional distributional relationships in synthetic experiments and MNIST tasks.
We study the estimation problem of distribution-on-distribution regression, where both predictors and responses are probability measures. Existing approaches typically rely on a global optimal transport map or tangent-space linearization, which can be restrictive in approximation capacity and distort geometry in multivariate underlying domains. In this paper, we propose the \emph{Neural Local Wasserstein Regression}, a flexible nonparametric framework that models regression through locally defined transport maps in Wasserstein space. Our method builds on the analogy with classical kernel regression: kernel weights based on the 2-Wasserstein distance localize estimators around reference measures, while neural networks parameterize transport operators that adapt flexibly to complex data geometries. This localized perspective broadens the class of admissible transformations and avoids the limitations of global map assumptions and linearization structures. We develop a practical training procedure using DeepSets-style architectures and Sinkhorn-approximated losses, combined with a greedy reference selection strategy for scalability. Through synthetic experiments on Gaussian and mixture models, as well as distributional prediction tasks on MNIST, we demonstrate that our approach effectively captures nonlinear and high-dimensional distributional relationships that elude existing methods.