Preconditioned Regularized Wasserstein Proximal Sampling

2025/09/01 by Tan, Hong Ye, Osher, Stanley, Li, Wuchen
#62G07 #65C05 #Computation (stat.CO) #FOS: Computer and information sciences #FOS: Mathematics #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Optimization and Control (math.OC)

paper · doi:10.48550/arxiv.2509.01685

Abstract

We consider sampling from a Gibbs distribution by evolving finitely many particles. We propose a preconditioned version of a recently proposed noise-free sampling method, governed by approximating the score function with the numerically tractable score of a regularized Wasserstein proximal operator. This is derived by a Cole--Hopf transformation on coupled anisotropic heat equations, yielding a kernel formulation for the preconditioned regularized Wasserstein proximal. The diffusion component of the proposed method is also interpreted as a modified self-attention block, as in transformer architectures. For quadratic potentials, we provide a discrete-time non-asymptotic convergence analysis and explicitly characterize the bias, which is dependent on regularization and independent of step-size. Experiments demonstrate acceleration and particle-level stability on various log-concave and non-log-concave toy examples to Bayesian total-variation regularized image deconvolution, and competitive/better performance on non-convex Bayesian neural network training when utilizing variable preconditioning matrices.

Citations

• DeepInverse: A Python package for solving imaging inverse problems with deep learning
• Splitting Regularized Wasserstein Proximal Algorithms for Nonsmooth Sampling Problems
• Diffusion at Absolute Zero: Langevin Sampling Using Successive Moreau Envelopes [conference paper]
• A Unified Perspective on the Dynamics of Deep Transformers
• Convergence of Noise-Free Sampling Algorithms with Regularized Wasserstein Proximals
• Laplace Meets Moreau: Smooth Approximation to Infimal Convolutions Using Laplace's Method
• FAdam: Adam is a natural gradient optimizer using diagonal empirical Fisher information
• Unsupervised Training of Convex Regularizers using Maximum Likelihood Estimation
• Tensor train based sampling algorithms for approximating regularized Wasserstein proximal operators
• A mathematical perspective on Transformers
• Noise-Free Sampling Algorithms via Regularized Wasserstein Proximals
• Subgradient Langevin Methods for Sampling from Non-smooth Potentials
• A kernel formula for regularized Wasserstein proximal operators
• Sinkformers: Transformers with Doubly Stochastic Attention
• Score-Based Generative Modeling through Stochastic Differential Equations
• Limitations of the Empirical Fisher Approximation for Natural Gradient Descent
• Neural Ordinary Differential Equations
• Mirrored Langevin Dynamics
• VEEGAN: Reducing Mode Collapse in GANs using Implicit Variational Learning
• Efficient Bayesian computation by proximal Markov chain Monte Carlo: when Langevin meets Moreau
• High-dimensional Bayesian inference via the Unadjusted Langevin Algorithm
• Adam: A Method for Stochastic Optimization
• Fast Image Recovery Using Variable Splitting and Constrained Optimization
• Nonlinear total variation based noise removal algorithms
• A computational fluid mechanics solution to the Monge-Kantorovich mass transfer problem

Related

• Noise-Free Sampling Algorithms via Regularized Wasserstein Proximals
• Convergence of Noise-Free Sampling Algorithms with Regularized Wasserstein Proximals
• Convergence of Noise-Free Sampling Algorithms with Regularized Wasserstein Proximals
• Splitting Regularized Wasserstein Proximal Algorithms for Nonsmooth Sampling Problems
• Tensor train based sampling algorithms for approximating regularized Wasserstein proximal operators