Neural Local Wasserstein Regression

• URL: http://arxiv.org/abs/2511.10824v1
• Date: Thu, 13 Nov 2025 21:54:18 GMT
• Title: Neural Local Wasserstein Regression
• Authors: Inga Girshfeld, Xiaohui Chen,
• Abstract summary: We study the estimation problem of distribution-on-distribution regression, where both predictors and responses are probability measures.<n>Existing approaches typically rely on a global optimal transport map or tangent-space linearization.<n>We propose a flexible nonparametric framework that models regression through locally defined transport maps in Wasserstein space.
• Score: 16.52489456261937
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Wasserstein Regression as a Variational Approximation of Probabilistic Trajectories through the Bernstein Basis [41.99844472131922]
  Existing approaches often ignore the geometry of the probability space or are computationally expensive.<n>A new method is proposed that combines the parameterization of probability trajectories using a Bernstein basis and the minimization of the Wasserstein distance between distributions.<n>The developed approach combines geometric accuracy, computational practicality, and interpretability.
  arXiv  Detail & Related papers  (2025-10-30T15:36:39Z)
• Likely Interpolants of Generative Models [12.543430183696648]
  Interpolation in generative models allows for controlled generation, model inspection, and more.<n>Most generative models lack a principal notion of interpolants without restrictive assumptions on either the model or data dimension.<n>We develop a general scheme that targets likely transition paths compatible with different metrics and probability distributions.
  arXiv  Detail & Related papers  (2025-10-30T08:46:53Z)
• Neural Optimal Transport Meets Multivariate Conformal Prediction [58.43397908730771]
  We propose a framework for conditional vectorile regression (CVQR)<n>CVQR combines neural optimal transport with quantized optimization, and apply it to predictions.
  arXiv  Detail & Related papers  (2025-09-29T19:50:19Z)
• Generalized Linear Mode Connectivity for Transformers [87.32299363530996]
  A striking phenomenon is linear mode connectivity (LMC), where independently trained models can be connected by low- or zero-loss paths.<n>Prior work has predominantly focused on neuron re-ordering through permutations, but such approaches are limited in scope.<n>We introduce a unified framework that captures four symmetry classes: permutations, semi-permutations, transformations, and general invertible maps.<n>This generalization enables, for the first time, the discovery of low- and zero-barrier linear paths between independently trained Vision Transformers and GPT-2 models.
  arXiv  Detail & Related papers  (2025-06-28T01:46:36Z)
• Deep Partially Linear Transformation Model for Right-Censored Survival Data [6.315323176162257]
  This paper introduces a deep partially linear transformation model (DPLTM) as a general and flexible regression framework.<n>The proposed method is capable of avoiding the curse of dimensionality while still retaining the interpretability of some cocensors of interest.<n> Comprehensive simulation studies demonstrate the impressive performance of the proposed procedure in terms of both accuracy and the predictive power.
  arXiv  Detail & Related papers  (2024-12-10T15:50:43Z)
• Deep Fréchet Regression [4.915744683251151]
  We propose a flexible regression model capable of handling high-dimensional predictors without imposing parametric assumptions.<n>The proposed model outperforms existing methods for non-Euclidean responses.
  arXiv  Detail & Related papers  (2024-07-31T07:54:14Z)
• Random Forest Weighted Local Fréchet Regression with Random Objects [18.128663071848923]
  We propose a novel random forest weighted local Fr'echet regression paradigm.<n>Our first method uses these weights as the local average to solve the conditional Fr'echet mean.<n>Second method performs local linear Fr'echet regression, both significantly improving existing Fr'echet regression methods.
  arXiv  Detail & Related papers  (2022-02-10T09:10:59Z)
• Probabilistic partition of unity networks: clustering based deep approximation [0.0]
  Partition of unity networks (POU-Nets) have been shown capable of realizing algebraic convergence rates for regression and solution of PDEs. We enrich POU-Nets with a Gaussian noise model to obtain a probabilistic generalization amenable to gradient-based generalizations of a maximum likelihood loss. We provide benchmarks quantifying performance in high/low-dimensions, demonstrating that convergence rates depend only on the latent dimension of data within high-dimensional space.
  arXiv  Detail & Related papers  (2021-07-07T08:02:00Z)
• MINIMALIST: Mutual INformatIon Maximization for Amortized Likelihood Inference from Sampled Trajectories [61.3299263929289]
  Simulation-based inference enables learning the parameters of a model even when its likelihood cannot be computed in practice. One class of methods uses data simulated with different parameters to infer an amortized estimator for the likelihood-to-evidence ratio. We show that this approach can be formulated in terms of mutual information between model parameters and simulated data.
  arXiv  Detail & Related papers  (2021-06-03T12:59:16Z)
• GELATO: Geometrically Enriched Latent Model for Offline Reinforcement Learning [54.291331971813364]
  offline reinforcement learning approaches can be divided into proximal and uncertainty-aware methods. In this work, we demonstrate the benefit of combining the two in a latent variational model. Our proposed metrics measure both the quality of out of distribution samples as well as the discrepancy of examples in the data.
  arXiv  Detail & Related papers  (2021-02-22T19:42:40Z)
• Identification of Probability weighted ARX models with arbitrary domains [75.91002178647165]
  PieceWise Affine models guarantees universal approximation, local linearity and equivalence to other classes of hybrid system. In this work, we focus on the identification of PieceWise Auto Regressive with eXogenous input models with arbitrary regions (NPWARX) The architecture is conceived following the Mixture of Expert concept, developed within the machine learning field.
  arXiv  Detail & Related papers  (2020-09-29T12:50:33Z)
• Provably Efficient Neural Estimation of Structural Equation Model: An Adversarial Approach [144.21892195917758]
  We study estimation in a class of generalized Structural equation models (SEMs) We formulate the linear operator equation as a min-max game, where both players are parameterized by neural networks (NNs), and learn the parameters of these neural networks using a gradient descent. For the first time we provide a tractable estimation procedure for SEMs based on NNs with provable convergence and without the need for sample splitting.
  arXiv  Detail & Related papers  (2020-07-02T17:55:47Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.