Publications and Preprints
Deep JKO: time-implicit particle methods for general nonlinear gradient flows
- Wonjun Lee, Li Wang, Wuchen Li
- Preprint, 2023.
We develop novel neural network-based implicit particle methods to compute high-dimensional Wasserstein-type gradient flows with linear and nonlinear mobility functions. The main idea is to use the Lagrangian formulation in the Jordan--Kinderlehrer--Otto (JKO) framework, where the velocity field is approximated using a neural network. We leverage the formulations from the neural ordinary differential equation (neural ODE) in the context of continuous normalizing flow for efficient density computation. Additionally, we make use of an explicit recurrence relation for computing derivatives, which greatly streamlines the backpropagation process. Our methodology demonstrates versatility in handling a wide range of gradient flows, accommodating various potential functions and nonlinear mobility scenarios. Extensive experiments demonstrate the efficacy of our approach, including an illustrative example from Bayesian inverse problems. This underscores that our scheme provides a viable alternative solver for the Kalman-Wasserstein gradient flow.
Monotone generative modeling via a Gromov-Monge embedding
- Wonjun Lee, Yifei Yang, Dongmian Zou, Gilad Lerman
- Preprint, 2023.
Generative Adversarial Networks (GANs) are powerful tools for creating new content, but they face challenges such as sensitivity to starting conditions and mode collapse. To address these issues, we propose a deep generative model that utilizes the Gromov-Monge embedding (GME). It helps identify the low-dimensional structure of the underlying measure of the data and then map it, while preserving its geometry, into a measure in a low-dimensional latent space, which is then optimally transported to the reference measure. We guarantee the preservation of the underlying geometry by the GME and c-cyclical monotonicity of the generative map, where c is an intrinsic embedding cost employed by the GME. The latter property is a first step in guaranteeing better robustness to initialization of parameters and mode collapse. Numerical experiments demonstrate the effectiveness of our approach in generating high-quality images, avoiding mode collapse, and exhibiting robustness to different starting conditions.
Monotone discretizations of levelset convex geometric PDEs
- Jeff Calder, Wonjun Lee
- Preprint, 2023.
We introduce a novel algorithm that converges to level-set convex viscosity solutions of high-dimensional Hamilton-Jacobi equations. The algorithm is applicable to a broad class of curvature motion PDEs, as well as a recently developed Hamilton-Jacobi equation for the Tukey depth, which is a statistical depth measure of data points. A main contribution of our work is a new monotone scheme for approximating the direction of the gradient, which allows for monotone discretizations of pure partial derivatives in the direction of, and orthogonal to, the gradient. We provide a convergence analysis of the algorithm on both regular Cartesian grids and unstructured point clouds in any dimension and present numerical experiments that demonstrate the effectiveness of the algorithm in approximating solutions of the affine flow in two dimensions and the Tukey depth measure of high-dimensional datasets such as MNIST and FashionMNIST.
Mean Field Control Problems For Vaccine Distribution
- Wonjun Lee, Siting Liu, Wuchen Li, Stanley Osher
- Research in the Mathematical Sciences, 2022.
With the invention of the COVID-19 vaccine, shipping and distributing are crucial in controlling the pandemic. In this paper, we build a mean-field variational problem in a spatial domain, which controls the propagation of pandemic by the optimal transportation strategy of vaccine distribution. Here we integrate the vaccine distribution into the mean-field SIR model designed in [arXiv:2006.01249]. Numerical examples demonstrate that the proposed model provides practical strategies in vaccine distribution on a spatial domain.
Computational Mean-Field Information Dynamics Associated With Reaction Diffusion Equations
- Wuchen Li, Wonjun Lee, Stanley Osher
- Journal of Computational Physics, 2022.
We formulate and compute a class of mean-field information dynamics based on reaction diffusion equations. Given a class of nonlinear reaction diffusion and entropy type Lyapunov functionals, we study their gradient flow formulations. We write the “mean-field” metric space formalisms and derive Hamiltonian flows therein. These Hamiltonian flows follow saddle point systems of the proposed mean-field control problems. We apply primal-dual hybrid-gradient algorithms to compute the mean field information dynamics. Several numerical examples are provided.
Random Features for High-Dimensional Nonlocal Mean-Field Games
- Sudanshu Agrawal, Wonjun Lee, Samy Wu Fung, Levon Nurbekyan
- Journal of Computational Physics, Mar, 2022.
We propose an efficient solution approach for high-dimensional nonlocal mean-field game (MFG) systems based on the Monte Carlo approximation of interaction kernels via random features. We avoid costly space-discretizations of interaction terms in the state-space by passing to the feature-space. This approach allows for a seamless mean-field extension of virtually any single-agent trajectory optimization algorithm. Here, we extend the direct transcription approach in optimal control to the mean-field setting. We demonstrate the efficiency of our method by solving MFG problems in high-dimensional spaces which were previously out of reach for conventional non-deep-learning techniques.
Weakly-Supervised Convolutional Neural Networks for Vessel Segmentation in Cerebral Angiography
- Arvind Vepa, Andrew Choi, Noor Nakhaei, Wonjun Lee, Noah Stier, Andrew Vu, Greyson Jenkins, Xiaoyan Yang, Manjot Shergill, Moira Desphy, Kevin Delao, Mia Levy, Cristopher Garduno, Lacy Nelson, Wandi Liu, Fan Hung, Fabien Scalzo
- Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (WACV), 2022, pp. 585-594
Automated vessel segmentation in cerebral digital subtraction angiography (DSA) has significant clinical utility in the management of cerebrovascular diseases such as stroke diagnosis and detection of aneurysms. While deep learning is state-of-the-art in segmentation, a significant amount of labeled data is needed for training. Because of domain differences, pretrained networks cannot be applied to DSA data out-of-the-box. We propose a novel learning framework, which utilizes an active contour model for weak supervision and low-cost human-in-the-loop strategies to improve weak label quality. Our study produces several significant results, including state-of-the-art results for cerebral DSA vessel segmentation, which exceed human annotator quality, and an analysis of annotation cost and model performance trade-offs utilizing weak supervision strategies. Additionally, we will be publicly releasing code to reproduce our methodology and our dataset, the largest known high-quality annotated cerebral DSA vessel segmentation dataset.
Tropical Optimal Transport and Wasserstein Distances in Phylogenetic Tree Space
- Wonjun Lee, Wuchen Li, Bo Lin, Anthea Monod
- Information Geometry, June, 2021.
We study the problem of optimal transport on phylogenetic tree space from the perspective of tropical geometry, and thus define the Wasserstein-p distances for probability measures in this continuous metric measure space setting. With respect to the tropical metric -- a combinatorial metric on the space of phylogenetic trees -- the cases of p = 1, 2 are treated in detail, which give an efficient way to compute geodesics and also provide theoretical foundations for geometric insight a statistical framework on tree space. We construct explicit algorithms for the computation of the tropical Wasserstein-1 and 2 distances, and prove their convergence. Our results provide the first study of the Wasserstein distances and optimal transport on sets of phylogenetic trees. Several numerical examples are provided.
The Back-And-Forth Method For Wasserstein Gradient Flows
- Matt Jacobs, Wonjun Lee, Flavien Léger
- ESAIM: COCV, 27:28, 2021.
We present a method to efficiently compute Wasserstein gradient flows. Our approach is based on a generalization of the back-and-forth method (BFM) introduced in [arXiv:1905.12154] to solve optimal transport problems. We evolve the gradient flow by solving the dual problem to the JKO scheme. In general, the dual problem is much better behaved than the primal problem. This allows us to efficiently run large scale gradient flows simulations for a large class of internal energies including singular and non-convex energies.
Controlling Propagation of epidemics via mean-field games
- Wonjun Lee, Siting Liu, Hamidou Tembine, Wuchen Li, Stanley Osher
- SIAM Journal on Applied Math, Dec, 2020.
The coronavirus disease 2019 (COVID-19) pandemic is changing and impacting lives on a global scale. In this paper, we introduce a mean-field control modelin controlling the propagation of epidemics on a spatial domain. The control variable,the spatial velocity, is first introduced for the classical disease models, such as the SIRmodel. For this proposed model, we provide fast numerical algorithms based on proximal primal-dual methods. Numerical experiments demonstrate that the proposed modelillustrates how to separate infected patients in a spatial domain effectively.
Joint Task Assignment and Trajectory Optimization for a Mobile Robot Swarm by Mean-Field Game
- Yuhan Kang, Siting Liu, Wonjun Lee, Hongliang Zhang, Wuchen Li, Zhu Han
- IEEE GLOBECOM, Dec, 2020.
In recent years, there has been a growing interest in utilizing mobile robot swarm to execute several tasks at the same time. However, how to assign tasks to the swarm and optimize the trajectory of the robots scientifically to minimize energy consumption is still a big challenge. In this paper, we consider a mobile robot swarm system where a large number of robots are deployed by a centralized controller to execute a series of tasks, such as target detection tasks, cooperatively. The controller controls the velocity strategy of each robot, and makes corresponding task assignment decisions to minimize the overall cost of the robot swarm. Since the number of involving robots is large, it will be extremely difficult to consider the interaction between them. In this regard, we adopt the concept of mean-field term to approximate the behaviors and states of the robots, and formulate the joint task assignment and trajectory optimization problem as a mean-field game. To solve the problem efficiently, a primal-dual hybrid gradient algorithm is proposed to find the optimal trajectory and corresponding task assignment decisions for each robot. The numerical simulation results show the effectiveness of the proposed algorithm.
Energy-efficient Velocity Control for Massive Rotary-Wing UAVs: A Mean Field Game Approach
- H. Gao, Wonjun Lee, Wuchen Li, Zhu Han, Stanley Osher, Vincent Poor
- IEEE GLOBECOM, Dec, 2020.
When a disaster happens in the metropolitan area, the wireless communication systems are highly affected, degrading the efficiency of the search and rescue (SAR) mission. An emergent wireless network must be deployed quickly and efficiently to preserve human lives. Teams of low-altitude rotary-wing unmanned aerial vehicle (UAVs) is preferable to be utilized as an on-demand temporal wireless network because they are generally faster to deploy, flexible to reconfigure, and able to provide better communication services with short line-of-sight links. However, rotary-wing UAVs’ limited on-board batteries require that they need to recharge and reconfigure frequently during a mission. Therefore, we formulate the velocity control problem for massive rotary-wing UAVs as a Schrödinger bridge problem which can describe the frequent reconfiguration of UAVs. Then we transform it into a mean field game and solve it with the G-prox primal dual hybrid gradient (PDHG) method. Finally, we show the efficiency of our algorithm and analyze the influence of wind dynamics with numerical results.
Generalized unnormalized optimal transport and its fast algorithms
- Wonjun Lee, Rongjie Lai, Wuchen Li, Stanley Osher
- Journal of Computational Physics, Nov, 2020.
We introduce fast algorithms for generalized unnormalized optimal transport. To handle densities with different total mass, we consider a dynamic model, which mixes the $L^p$ optimal transport with $L^p$ distance. For $p=1$, we derive the corresponding $L^1$ generalized unnormalized Kantorovich formula. We further show that the problem becomes a simple $L^1$ minimization which is solved efficiently by a primal-dual algorithm. For $p=2$, we derive the $L^2$ generalized unnormalized Kantorovich formula, a new unnormalized Monge problem and the corresponding Monge-Amp\`ere equation. Furthermore, we introduce a new unconstrained optimization formulation of the problem. The associated gradient flow is essentially related to an elliptic equation which can be solved efficiently. Here the proposed gradient descent procedure together with the Nesterov acceleration involves the Hamilton-Jacobi equation which arises from the KKT conditions. Several numerical examples are presented to illustrate the effectiveness of the proposed algorithms.