Optimal Neural Network Approximation of Wasserstein Gradient Direction via Convex Optimization\ast

The calculation of the direction of the Wasserstein gradient is vital for addressing problems related to posterior sampling and scientific computing. To approximate the Wasserstein gradient using finite samples, it is necessary to solve a variation problem. Our study focuses on the variation problem within the framework of two-layer networks with squared ReLU activations. We present a semidefinite program (SDP) relaxation as a solution, which can be viewed as an approximation of the Wasserstein gradient for a broader range of functions, including two-layer networks. By solving the convex SDP, we achieve the best approximation of the Wasserstein gradient direction in this function class. We also provide conditions to ensure the relaxation is tight. Additionally, we propose methods for practical implementation, such as subsampling and dimension reduction. The effectiveness and efficiency of our proposed method are demonstrated through numerical experiments, including Bayesian inference with PDE constraints and parameter estimation in COVID-19 modeling.