I am an associate professor in mathematics (tenured) at the KTH Royal Institute of Technology in Stockholm.
See here for a short bio and my journey so far.
I also write a non-research blog here. However, the update frequency depends on how busy I am at the moment.
Overall, I am interested in optimization and mathematical foundation of machine learning and statistical inference algorithms. I’m especially interested in principled analysis, e.g., optimal transport, PDE, calculus of variations, gradient flows dynamics, kernel methods. Recently, I have been working on (infinite-dimensional, gradient-based) optimization and dynamics of probability distributions. A central perspective here is that diffusion equations can be seen as steepest descent (gradient flows) in the geometry defined by optimal transportation. This turns questions about randomness into questions about optimization and the calculus of variations. My research interests sit at the interface of classical and modern approaches to those topics.
To get in touch, click the icon at the bottom of the page. There are sometimes delays in my response to emails, please be patient.
Recent talks (selected)
- July 20 - 31, 2026. I gave a lecture series on “Computational Gradient Flows and Optimal Transport” at Peking University, at the School of Mathematical Sciences. This is a Summer School on Applied Mathematics. The lecture notes (student scribed) are available as a PDF file entitled “Computational Gradient Flows and Optimal Transport: an Introduction”.
- SIAM UQ 2026 Minisymposium, Gradient Flows for Uncertainty Quantification: New Algorithms and Applications
- EPFL Bernoulli workshop “Particles, Flows & Maps for Sampling Complex Distributions”, 2025. Video recording available here
- Gradient Flows Face-to-Face Workshop in Granada, Spain, 2025. Organizers: Maria Bruna, José Alfredo Cañizo, José Antonio Carrillo, Antonio Esposito. Slides available here
- Workshop on Bayesian Analysis and Artificial Intelligence, Peking University, Beijing, China, 2025
- Banff International Research Station (BIRS) “Mathematical Analysis of Adversarial Machine Learning” workshop in Oaxaca, Mexico, from August 17 to August 22, 2025
- The Kantorovich Initiative Seminar, University of British Columbia: Kernel Approximation of Wasserstein and Fisher-Rao Gradient flows
- Invited talk slides for the EPFL talk, Nov 2024
- Workshop on Optimal Transport and PDEs, Program on The Mathematics of Data, Institute for Mathematical Sciences, National University of Singapore, 2024. Organizers: Philippe Rigollet, Afonso Bandeira, Subhro Ghosh
Open positions
- KTH Master thesis: if you are a master’s student already enrolled in KTH and interested in optimal transport/optimization (theory/applications), state-of-the-art machine learning/optimization, generative models, PDE/SDE, please feel free to reach out. Please attach your CV and transcripts (undergraduate and graduate courses) so I can judge the fit with the topics.
- There are currently no PhD or postdoc positions in my group. Future positions will be announced here when available.
News and updates
-
Lecture notes available: my lecture notes for my summer school lectures at the Peking University is available at here. Many thanks to the student who scribed it.
- New preprints:
- A new version of the inclusive-KL gradient flows, that includes new results such as inclusive/forward KL-Gaussian gradient flows, local-estimator Wasserstein gradient flows, is updated here: Inclusive KL Gradient Flows: Otto-Wasserstein, Fisher-Rao-Gaussian, and Local-Estimator Dynamics. Jia-Jie Zhu.
-
Workshop organization: 24 May 2026 - 29 May 2026: SwissMAP Workshop on “Computational Optimization Meets Gradient Flows and Optimal Transport”. Organized by Niao He (ETH Zurich), Yifan Hu (EPFL), Daniel Kuhn (EPFL), Jia-Jie Zhu (KTH)
- New publication accepted:
- Gradient Flow Sampler-based Distributionally Robust Optimization. ICML 2026, Joint work with Zusen Xu. Code
- Improved Stochastic Optimization of LogSumExp. ICML 2026, Joint work with Egor Gladin, Alexey Kroshnin, Pavel Dvurechensky.
- Conditional LDDMM flow matching: an application to uncertainty quantification. Computer Methods in Applied Mechanics and Engineering. Accepted for publication. Joint work with Sarah Katz, Francesco Romor, Alfonso Caiazzo.
- New publication in the special issue of the journal Communications on Pure and Applied Analysis on the topic of “Transport Equations in Optimization, Sampling, and Control”: Evolution of Gaussians in the Hellinger-Kantorovich-Boltzmann gradient flow. Joint work with Matthias Liero, Alexander Mielke, Oliver Tse.
-
New publication in Journal of Optimization Theory and Applications (JOTA): An Inexact Halpern Iteration with Application to Distributionally Robust Optimization. Joint work with Ling Liang and Kim-Chuan Toh.
- 28 July - 1 August 2025: ICSP 2025 invited session on interfacing optimal transport, gradient flows, diffusion models, and machine learning, optimization, at École des Ponts, IP Paris. Organized by Pavel Dvurechensky, Jia-Jie Zhu
- Speaker slides of the talks are available here
- New preprints:
- New work to appear in NeurIPS 2024: Globally convergent gradient flows for the MMD-minimization inference problem (a.k.a. MMD-flow). Preprint: Egor Gladin, Pavel Dvurechensky, Alexander Mielke, Jia-Jie Zhu. Interaction-Force Transport Gradient Flows Code: link, Slides (NeurIPS 2024)
- I’m serving as an area chair for AISTATS 2025. If you are interested in contributing to the community via reviewing a paper, please get in touch.
- Summer 2024. New third-party funding awarded: DFG Project on “Optimal Transport and Measure Optimization Foundation for Robust and Causal Machine Learning” within the Priority Program “Theoretical Foundations of Deep Learning” (SPP 2298).
- March 11th - 15th, 2024. I am organizing a Workshop on Optimal Transport from Theory to Applications – Interfacing Dynamical Systems, Optimization, and Machine Learning (OT-DOM) in Berlin, Germany. program and slides
- New preprints available:
- Approximation, Kernelization, and Entropy-Dissipation of Gradient Flows: from Wasserstein to Fisher-Rao. Joint work with Alexander Mielke.
- Analysis of Kernel Mirror Prox for Measure Optimization. Accepted for publication at AISTATS 2024. Joint work with Pavel Dvurechensky.
- I am teaching the nonparametric statistics course at Humboldt University of Berlin (at master level), co-lecturing with Vladimir Spokoinyi, in term 2023/24.
- Recent talks on gradient flow force-balance, especially in robust learning under (strong) structured distribution shifts, and conditional moment restriction for causal inference at
- EUCCO 2023, at Heidelberg University. Talk slides available
- I taught a mini-course on the optimization perspective of gradient flow dynamics, introducing (beginner-friendly) concepts of gradient flows in the Eulidean and Wasserstein space, at the Workshop of Intelligent Autonomous Learning Systems 2023. Lecture slides available:
- July 2023. I gave an invited talk at the ICML 2023 Workshop on Duality Principles for Modern Machine Learning. The slides are available here.
- July 2023. Accepted paper at CDC 2023: Propagating Kernel Ambiguity Sets in Nonlinear Data-driven Dynamics Models
- May 2023. Accepted paper at ICML 2023 (link to preprint): Heiner Kremer, Yassine Nemmour, Bernhard Sch ̈olkopf, and Jia-Jie Zhu. Estimation Beyond Data Reweighting: Kernel Method of Moments.
- May 2023. A couple of new preprints are available:
- Apr 2023. Gave a plenary talk at the Leibniz Institute for Agricultural Engineering and Bioeconomy Potsdam, during the workshop “Mathematical Modeling and Simulation” (MMS) Days.
- Served as area chair for AISTATS 2023.