Neural Local Wasserstein Regression

TL;DR

This paper introduces Neural Local Wasserstein Regression, a flexible nonparametric method for distribution-on-distribution regression that models local transport maps in Wasserstein space, overcoming limitations of global approaches.

Contribution

It proposes a novel localized framework using neural networks and kernel weighting to improve modeling of complex distributional relationships.

Findings

Effective on synthetic Gaussian and mixture models

Captures nonlinear high-dimensional relationships

Performs well on MNIST distributional prediction

Abstract

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.