Optimal Neural Network Approximation of Wasserstein Gradient Direction via Convex Optimization

TL;DR

This paper introduces a convex optimization approach to accurately approximate Wasserstein gradient directions using two-layer neural networks, enhancing posterior sampling and scientific computing tasks.

Contribution

It develops a semi-definite programming relaxation for the Wasserstein gradient approximation within neural network function classes, providing an optimal solution.

Findings

Effective in PDE-constrained Bayesian inference

Improves parameter estimation in COVID-19 modeling

Demonstrates superior approximation accuracy

Abstract

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.