JKO schemes with general transport costs (2402.17681v1)
Abstract: We modify the JKO scheme, which is a time discretization of Wasserstein gradient flows, by replacing the Wasserstein distance with more general transport costs on manifolds. We show when the cost function has a mixed Hessian which defines a Riemannian metric, our modified JKO scheme converges under suitable conditions to the corresponding Riemannian Fokker--Planck equation. Thus on a Riemannian manifold one may replace the (squared) Riemannian distance with any cost function which induces the metric. Of interest is when the Riemannian distance is computationally intractable, but a suitable cost has a simple analytic expression. We consider the Fokker--Planck equation on compact submanifolds with the Neumann boundary condition and on complete Riemannian manifolds with a finite drift condition. As an application we consider Hessian manifolds, taking as a cost the Bregman divergence.
- From a large-deviations principle to the Wasserstein gradient flow: a new micro-macro passage. Communications in Mathematical Physics, 307:791–815, 2011.
- Martial Agueh. Existence of solutions to degenerate parabolic equations via the Monge-Kantorovich theory. Advances in Differential Equations, 10(3):309–360, 2005.
- Hessian Riemannian gradient flows in convex programming. SIAM Journal on Control and Optimization, 43(2):477–501, 2004.
- Shun-Ichi Amari. Information Geometry and Its Applications. Springer, 2016.
- Gradient Flows in Metric Spaces and in the Space of Probability Measures. Lectures in Mathematics ETH Zürich. Birkhäuser Verlag, Basel, second edition, 2008.
- Input convex neural networks. In International Conference on Machine Learning, pages 146–155. PMLR, 2017.
- Clustering with Bregman divergences. Journal of Machine Learning Research, 6(Oct):1705–1749, 2005.
- Ole Barndorff-Nielsen. Information and Exponential Families in Statistical Theory. John Wiley & Sons, 2014.
- Michéle Basseville. Divergence measures for statistical data processing—an annotated bibliography. Signal Processing, 93(4):621–633, 2013.
- Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters, 31(3):167–175, 2003.
- A geometric approach to apriori estimates for optimal transport maps. arXiv preprint arXiv:2311.10208, 2023.
- Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications on Pure and Applied Mathematics, 44(4):375–417, 1991.
- Proximal optimal transport modeling of population dynamics. In International Conference on Artificial Intelligence and Statistics, pages 6511–6528. PMLR, 2022.
- Geometric Modeling in Probability and Statistics. Springer, 2014.
- An explicit solution of information geodesic equations for the multivariate normal model. Statistics & Risk Modeling, 9(1-2):119–138, 1991.
- On Monge’s problem for Bregman-like cost functions. Journal of Convex Analysis, 14(3):647–655, 2007.
- Wasserstein mirror gradient flow as the limit of the Sinkhorn algorithm, 2023. arXiv 2307.16421.
- Matthias Erbar. The heat equation on manifolds as a gradient flow in the Wasserstein space. Ann. Inst. Henri Poincaré Probab. Stat., 46(1):1–23, 2010.
- A variational method for a class of parabolic PDEs. Ann. Sc. Norm. Super. Pisa Cl. Sci. (5), 10(1):207–252, 2011.
- A new transportation distance between non-negative measures, with applications to gradients flows with Dirichlet boundary conditions. J. Math. Pures Appl. (9), 94(2):107–130, 2010.
- An Invitation to Optimal Transport, Wasserstein distances, and Gradient Flows. EMS Press, 2021.
- Avner Friedman. Partial Differential Equations of Parabolic Type. Prentice-Hall, Inc., Englewood Cliffs, NJ, 1964.
- Hopfield neural network flow: A geometric viewpoint. IEEE Transactions on Neural Networks and Learning Systems, 31(11):4869–4880, 2020.
- The variational formulation of the fokker–planck equation. SIAM Journal on Mathematical Analysis, 29(1):1–17, 1998.
- Olav Kallenberg. Foundations of Modern Probability, volume 1. Springer, third edition.
- Olav Kallenberg. Foundations of Modern Probability, volume 2. Springer, third edition.
- Continuity, curvature, and the general covariance of optimal transportation. Journal of the European Mathematical Society, 12(4):1009–1040, 2010.
- John M Lee. Introduction to Riemannian Manifolds. Springer, second edition, 2018.
- Gradient descent with a general cost, 2023. arXiv 2305.04917.
- Regularity of potential functions of the optimal transportation problem. Archive for Rational Mechanics and Analysis, 177(2):151–183, 2005.
- Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pages 6672–6681. PMLR, 2020.
- Takao Matumoto. Any statistical manifold has a contrast function – on the C3superscript𝐶3{C}^{3}italic_C start_POSTSUPERSCRIPT 3 end_POSTSUPERSCRIPT-function taking the minimum at the diagonal of the product manifold. Hiroshima Mathematical Journal, 23(2):327–332, 1993.
- Robert J McCann. A convexity principle for interacting gases. Advances in Mathematics, 128(1):153–179, 1997.
- Robert John McCann. A convexity theory for interacting gases and equilibrium crystals. ProQuest LLC, Ann Arbor, MI, 1994. Thesis (Ph.D.)–Princeton University.
- Large-scale Wasserstein gradient flows. Advances in Neural Information Processing Systems, 34:15243–15256, 2021.
- Contraction of general transportation costs along solutions to Fokker-Planck equations with monotone drifts. J. Math. Pures Appl. (9), 95(1):18–35, 2011.
- Barrett O’Neill. Semi-Riemannian Geometry with Applications to Relativity. Academic Press, 1983.
- Felix Otto. Dynamics of labyrinthine pattern formation in magnetic fluids: a mean-field theory. Arch. Rational Mech. Anal., 141(1):63–103, 1998.
- Felix Otto. The geometry of dissipative evolution equations: the porous medium equation. Comm. Partial Differential Equations, 26(1-2):101–174, 2001.
- Soumik Pal. On the difference between entropic cost and the optimal transport cost. The Annals of Applied Probability, 34(1B):1003–1028, 2024.
- Multiplicative Schröodinger problem and the Dirichlet transport. Probability Theory and Related Fields, 178(1):613–654, 2020.
- Gabriel Peyré. Entropic approximation of Wasserstein gradient flows. SIAM Journal on Imaging Sciences, 8(4):2323–2351, 2015.
- Cale Rankin. A remark on the geometric interpretation of the A3w condition from optimal transport. Bull. Aust. Math. Soc., 108(1):162–165, 2023.
- Bregman-Wasserstein divergence: geometry and applications. arXiv preprint arXiv:2302.05833, 2023.
- R Tyrrell Rockafellar. Convex Analysis. Princeton University Press, 1970.
- Filippo Santambrogio. Optimal Transport for Applied Mathematicians. Birkäuser, 2015.
- Filippo Santambrogio. {{\{{Euclidean, metric, and Wasserstein}}\}} gradient flows: an overview. Bulletin of Mathematical Sciences, 7(1):87–154, 2017.
- Giuseppe Savaré. Gradient flows and diffusion semigroups in metric spaces under lower curvature bounds. C. R. Math. Acad. Sci. Paris, 345(3):151–154, 2007.
- Hirohiko Shima. The Geometry of Hessian Structures. World Scientific, 2007.
- Geometry of Hessian manifolds. Differential geometry and its applications, 7(3):277–290, 1997.
- On strict convexity and continuous differentiability of potential functions in optimal transportation. Arch. Ration. Mech. Anal., 192(3):403–418, 2009.
- Cédric Villani. Topics in Optimal Transportation. American Mathematical Society, 2003.
- Cédric Villani. Optimal Transport: Old and New. Springer, 2008.
- A variational perspective on accelerated methods in optimization. Proceedings of the National Academy of Sciences, 113(47):E7351–E7358, 2016.
- Pseudo-Riemannian geometry encodes information geometry in optimal transport. Information Geometry, 5(1):131–159, 2022.
- Tsallis and Rényi deformations linked via a new λ𝜆\lambdaitalic_λ-duality. IEEE Transactions on Information Theory, 68(8):5353–5373, 2022.
- Xicheng Zhang. Variational approximation for Fokker-Planck equation on Riemannian manifold. Probability Theory and Related Fields, 137(3-4):519–539, 2007.