# Optimal Neural Network Approximation of Wasserstein Gradient Direction via Convex Optimization

arXiv: https://arxiv.org/abs/2205.13098 · HTML: https://mlsanity.com/paper/2205.13098/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 44 references that are in the map.
- First posted on arXiv 26 May 2022.

## Builds on (most cited first, 30 of 30)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [CVXPY: A Python-Embedded Modeling Language for Convex Optimization](https://mlsanity.com/paper/1603.00943/) (2016)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [The Expressive Power of Neural Networks: A View from the Width](https://mlsanity.com/paper/1709.02540/) (2017)
- [OT-Flow: Fast and Accurate Continuous Normalizing Flows via Optimal Transport](https://mlsanity.com/paper/2006.00104/) (2020)
- [Neural Networks are Convex Regularizers: Exact Polynomial-time Convex Optimization Formulations for Two-layer Networks](https://mlsanity.com/paper/2002.10553/) (2020)
- [Proximal Optimal Transport Modeling of Population Dynamics](https://mlsanity.com/paper/2106.06345/) (2021)
- [Sliced-Wasserstein Flows: Nonparametric Generative Modeling via Optimal Transport and Diffusions](https://mlsanity.com/paper/1806.08141/) (2018)
- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)
- [Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks](https://mlsanity.com/paper/2106.00774/) (2021)
- [Variational Wasserstein gradient flow](https://mlsanity.com/paper/2112.02424/) (2021)
- [A Stein variational Newton method](https://mlsanity.com/paper/1806.03085/) (2018)
- [Dimension-independent likelihood-informed MCMC](https://mlsanity.com/paper/1411.3688/) (2014)
- [Alternating the Population and Control Neural Networks to Solve High-Dimensional Stochastic Mean-Field Games](https://mlsanity.com/paper/2002.10113/) (2020)
- [Vector-output ReLU Neural Network Problems are Copositive Programs: Convex Analysis of Two Layer Networks and Polynomial-time Algorithms](https://mlsanity.com/paper/2012.13329/) (2020)
- [Projected Stein Variational Gradient Descent](https://mlsanity.com/paper/2002.03469/) (2020)
- [HINT: Hierarchical Invertible Neural Transport for Density Estimation and Bayesian Inference](https://mlsanity.com/paper/1905.10687/) (2019)
- [Certified dimension reduction in nonlinear Bayesian inverse problems](https://mlsanity.com/paper/1807.03712/) (2018)
- [Projected Wasserstein gradient descent for high-dimensional Bayesian inference](https://mlsanity.com/paper/2102.06350/) (2021)
- [SDPNAL+: A Majorized Semismooth Newton-CG Augmented Lagrangian Method for Semidefinite Programming with Nonnegative Constraints](https://mlsanity.com/paper/1406.0942/) (2014)
- [Demystifying Batch Normalization in ReLU Networks: Equivalent Convex Optimization Models and Implicit Regularization](https://mlsanity.com/paper/2103.01499/) (2021)
- [Neural Variational Gradient Descent](https://mlsanity.com/paper/2107.10731/) (2021)
- [Wasserstein Proximal of GANs](https://mlsanity.com/paper/2102.06862/) (2021)
- [The Deep Minimizing Movement Scheme](https://mlsanity.com/paper/2109.14851/) (2021)
- [Neural Spectrahedra and Semidefinite Lifts: Global Convex Optimization of Polynomial Activation Neural Networks in Fully Polynomial-Time](https://mlsanity.com/paper/2101.02429/) (2021)
- [Approximate inference with Wasserstein gradient flows](https://mlsanity.com/paper/1806.04542/) (2018)
- [Scaling Up Bayesian Uncertainty Quantification for Inverse Problems using Deep Neural Networks](https://mlsanity.com/paper/2101.03906/) (2021)
- [Relative Entropy Gradient Sampler for Unnormalized Distributions](https://mlsanity.com/paper/2110.02787/) (2021)
- [Trust--Region Problems with Linear Inequality Constraints: Exact SDP Relaxation, Global Optimality and Robust Optimization](https://mlsanity.com/paper/1309.3000/) (2013)

## Built on by (most cited first, 9 of 9)

- [Particle-based Variational Inference with Preconditioned Functional Gradient Flow](https://mlsanity.com/paper/2211.13954/) (2022)
- [Sampling via Gradient Flows in the Space of Probability Measures](https://mlsanity.com/paper/2310.03597/) (2023)
- [Gradient Flows for Sampling: Mean-Field Models, Gaussian Approximations and Affine Invariance](https://mlsanity.com/paper/2302.11024/) (2023)
- [Mirror and Preconditioned Gradient Descent in Wasserstein Space](https://mlsanity.com/paper/2406.08938/) (2024)
- [Particle-based Variational Inference with Generalized Wasserstein Gradient Flow](https://mlsanity.com/paper/2310.16516/) (2023)
- [Score-based Neural Ordinary Differential Equations for Computing Mean Field Control Problems](https://mlsanity.com/paper/2409.16471/) (2024)
- [A score-based particle method for homogeneous Landau equation](https://mlsanity.com/paper/2405.05187/) (2024)
- [Leveraging Optimal Transport via Projections on Subspaces for Machine Learning Applications](https://mlsanity.com/paper/2311.13883/) (2023)
- [A Library of Mirrors: Deep Neural Nets in Low Dimensions are Convex Lasso Models with Reflection Features](https://mlsanity.com/paper/2403.01046/) (2024)

## Abstract (arXiv)

The computation of Wasserstein gradient direction is essential for posterior sampling problems and scientific computing. The approximation of the Wasserstein gradient with finite samples requires solving a variational problem. We study the variational problem in the family of two-layer networks with squared-ReLU activations, towards which we derive a semi-definite programming (SDP) relaxation. This SDP can be viewed as an approximation of the Wasserstein gradient in a broader function family including two-layer networks. By solving the convex SDP, we obtain the optimal approximation of the Wasserstein gradient direction in this class of functions. Numerical experiments including PDE-constrained Bayesian inference and parameter estimation in COVID-19 modeling demonstrate the effectiveness of the proposed method.