skip to main content

Title: Primal Dual Methods for Wasserstein Gradient Flows
Abstract Combining the classical theory of optimal transport with modern operator splitting techniques, we develop a new numerical method for nonlinear, nonlocal partial differential equations, arising in models of porous media, materials science, and biological swarming. Our method proceeds as follows: first, we discretize in time, either via the classical JKO scheme or via a novel Crank–Nicolson-type method we introduce. Next, we use the Benamou–Brenier dynamical characterization of the Wasserstein distance to reduce computing the solution of the discrete time equations to solving fully discrete minimization problems, with strictly convex objective functions and linear constraints. Third, we compute the minimizers by applying a recently introduced, provably convergent primal dual splitting scheme for three operators (Yan in J Sci Comput 1–20, 2018). By leveraging the PDEs’ underlying variational structure, our method overcomes stability issues present in previous numerical work built on explicit time discretizations, which suffer due to the equations’ strong nonlinearities and degeneracies. Our method is also naturally positivity and mass preserving and, in the case of the JKO scheme, energy decreasing. We prove that minimizers of the fully discrete problem converge to minimizers of the spatially continuous, discrete time problem as the spatial discretization is refined. We conclude with more » simulations of nonlinear PDEs and Wasserstein geodesics in one and two dimensions that illustrate the key properties of our approach, including higher-order convergence our novel Crank–Nicolson-type method, when compared to the classical JKO method. « less
; ; ;
Award ID(s):
1846854 1811012
Publication Date:
Journal Name:
Foundations of Computational Mathematics
Sponsoring Org:
National Science Foundation
More Like this
  1. In this paper, we consider electromagnetic (EM) wave propagation in nonlinear optical media in one spatial dimension. We model the EM wave propagation by the time- dependent Maxwell’s equations coupled with a system of nonlinear ordinary differential equations (ODEs) for the response of the medium to the EM waves. The nonlinearity in the ODEs describes the instantaneous electronic Kerr response and the residual Raman molecular vibrational response. The ODEs also include the single resonance linear Lorentz dispersion. For such model, we will design and analyze fully discrete finite difference time domain (FDTD) methods that have arbitrary (even) order in spacemore »and second order in time. It is challenging to achieve provable stability for fully discrete methods, and this depends on the choices of temporal discretizations of the nonlinear terms. In Bokil et al. (J Comput Phys 350:420–452, 2017), we proposed novel modifications of second-order leap-frog and trapezoidal temporal schemes in the context of discontinuous Galerkin methods to discretize the nonlinear terms in this Maxwell model. Here, we continue this work by developing similar time discretizations within the framework of FDTD methods. More specifically, we design fully discrete modified leap-frog FDTD methods which are proved to be stable under appropriate CFL conditions. These method can be viewed as an extension of the Yee-FDTD scheme to this nonlinear Maxwell model. We also design fully discrete trapezoidal FDTD methods which are proved to be unconditionally stable. The performance of the fully discrete FDTD methods are demonstrated through numerical experiments involving kink, antikink waves and third harmonic generation in soliton propagation.« less
  2. We develop a new computational framework to solve the partial differential equations (PDEs) governing the flow of the joint probability density functions (PDFs) in continuous-time stochastic nonlinear systems. The need for computing the transient joint PDFs subject to prior dynamics arises in uncertainty propagation, nonlinear filtering and stochastic control. Our methodology breaks away from the traditional approach of spatial discretization or function approximation – both of which, in general, suffer from the “curse-of-dimensionality”. In the proposed framework, we discretize time but not the state space. We solve infinite dimensional proximal recursions in the manifold of joint PDFs, which in themore »small time-step limit, is theoretically equivalent to solving the underlying transport PDEs. The resulting computation has the geometric interpretation of gradient flow of certain free energy functional with respect to the Wasserstein metric arising from the theory of optimal mass transport. We show that dualization along with an entropic regularization, leads to a cone-preserving fixed point recursion that is proved to be contractive in Thompson metric. A block co-ordinate iteration scheme is proposed to solve the resulting nonlinear recursions with guaranteed convergence. This approach enables remarkably fast computation for non-parametric transient joint PDF propagation. Numerical examples and various extensions are provided to illustrate the scope and efficacy of the proposed approach.« less
  3. Wasserstein gradient flows provide a powerful means of understanding and solving many diffusion equations. Specifically, Fokker-Planck equations, which model the diffusion of probability measures, can be understood as gradient descent over entropy functionals in Wasserstein space. This equivalence, introduced by Jordan, Kinderlehrer and Otto, inspired the so-called JKO scheme to approximate these diffusion processes via an implicit discretization of the gradient flow in Wasserstein space. Solving the optimization problem associated with each JKO step, however, presents serious computational challenges. We introduce a scalable method to approximate Wasserstein gradient flows, targeted to machine learning applications. Our approach relies on input-convex neuralmore »networks (ICNNs) to discretize the JKO steps, which can be optimized by stochastic gradient descent. Contrarily to previous work, our method does not require domain discretization or particle simulation. As a result, we can sample from the measure at each time step of the diffusion and compute its probability density. We demonstrate the performance of our algorithm by computing diffusions following the Fokker-Planck equation and apply it to unnormalized density sampling as well as nonlinear filtering.« less
  4. Abstract We revisit the problem of approximating minimizers of certain convex functionals subject to a convexity constraint by solutions of fourth order equations of Abreu type. This approximation problem was studied in previous articles of Carlier–Radice (Approximation of variational problems with a convexity constraint by PDEs of Abreu type. Calc. Var. Partial Differential Equations 58 (2019), no. 5, Art. 170) and the author (Singular Abreu equations and minimizers of convex functionals with a convexity constraint, arXiv:1811.02355v3, Comm. Pure Appl. Math. , to appear), under the uniform convexity of both the Lagrangian and constraint barrier. By introducing a new approximating scheme,more »we completely remove the uniform convexity of both the Lagrangian and constraint barrier. Our analysis is applicable to variational problems motivated by the original 2D Rochet–Choné model in the monopolist's problem in Economics, and variational problems arising in the analysis of wrinkling patterns in floating elastic shells in Elasticity.« less
  5. We propose a unified data-driven framework based on inverse optimal transport that can learn adaptive, nonlinear interaction cost function from noisy and incomplete empirical matching matrix and predict new matching in various matching contexts. We emphasize that the discrete optimal transport plays the role of a variational principle which gives rise to an optimization based framework for modeling the observed empirical matching data. Our formulation leads to a non-convex optimization problem which can be solved efficiently by an alternating optimization method. A key novel aspect of our formulation is the incorporation of marginal relaxation via regularized Wasserstein distance, significantly improvingmore »the robustness of the method in the face of noisy or missing empirical matching data. Our model falls into the category of prescriptive models, which not only predict potential future matching, but is also able to explain what leads to empirical matching and quantifies the impact of changes in matching factors. The proposed approach has wide applicability including predicting matching in online dating, labor market, college application and crowdsourcing. We back up our claims with numerical experiments on both synthetic data and real world data sets.« less