Publications

Articles in Refereed Journals

  • R.A. Vandermeulen, R. Saitenmacher, Generalized identifiability bounds for mixture models with grouped samples, IEEE Transactions on Information Theory, 70 (2024), pp. 2746--2758, DOI 10.1109/TIT.2024.3367433 .

Contributions to Collected Editions

  • P. Dvurechensky, J.-J. Zhu, Analysis of kernel mirror prox for measure optimization, in: Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, S. Dasgupta, S. Mandt, Y. Li, eds., 238 of Proceedings of Machine Learning Research, 2024, pp. 2350--2358.
    Abstract
    Kernel mirror prox and RKHS gradient flow for mixed functional Nash equilibrium Pavel Dvurechensky , Jia-Jie Zhu Abstract The theoretical analysis of machine learning algorithms, such as deep generative modeling, motivates multiple recent works on the Mixed Nash Equilibrium (MNE) problem. Different from MNE, this paper formulates the Mixed Functional Nash Equilibrium (MFNE), which replaces one of the measure optimization problems with optimization over a class of dual functions, e.g., the reproducing kernel Hilbert space (RKHS) in the case of Mixed Kernel Nash Equilibrium (MKNE). We show that our MFNE and MKNE framework form the backbones that govern several existing machine learning algorithms, such as implicit generative models, distributionally robust optimization (DRO), and Wasserstein barycenters. To model the infinite-dimensional continuous- limit optimization dynamics, we propose the Interacting Wasserstein-Kernel Gradient Flow, which includes the RKHS flow that is much less common than the Wasserstein gradient flow but enjoys a much simpler convexity structure. Time-discretizing this gradient flow, we propose a primal-dual kernel mirror prox algorithm, which alternates between a dual step in the RKHS, and a primal step in the space of probability measures. We then provide the first unified convergence analysis of our algorithm for this class of MKNE problems, which establishes a convergence rate of O(1/N ) in the deterministic case and O(1/√N) in the stochastic case. As a case study, we apply our analysis to DRO, providing the first primal-dual convergence analysis for DRO with probability-metric constraints.

Talks, Poster

  • J.-J. Zhu, Analysis of kernel mirror prox for measure optimization, 27th Conference on Artificial Intelligence and Statistics (AISTATS), May 2 - 4, 2024, Valencia, Spain, May 2, 2024.

  • J.-J. Zhu, Approximation and kernelization of gradient flow geometry: Fisher-Rao and Wasserstein, The Mathematics of Data: Workshop on Optimal Transport and PDEs, January 17 - 23, 2024, National University of Singapore, Institute for Mathematical Sciences, Singapore, January 22, 2024.

  • J.-J. Zhu, Flow and transport: Modern mathematical foundation for statistical machine learning, University of St. Gallen, Switzerland, October 17, 2024.

  • J.-J. Zhu, From distributional ambiguity to gradient flows: Wasserstein, Fisher-Rao, and kernel approximation, École Polytechnique Fédérale de Lausanne, Switzerland, November 28, 2024.

  • J.-J. Zhu, Gradient flows and kernelization in the Hellinger-Kantorovich (a.k.a. Wasserstein-Fisher-Rao) space, Europt 2024, 21st Conference on Advances in Continuous Optimization, June 26 - 28, 2024, Lund University, Department of Automatic Control, Sweden, June 28, 2024.

  • J.-J. Zhu, Kernel approximation of Wasserstein and Fisher-Rao Gradient flows, École Polytechnique Fédérale de Lausanne, Switzerland, November 27, 2024.

  • J.-J. Zhu, Kernel approximation of Wasserstein and Fisher-Rao gradient flows, The Thirty-Eighth Annual Conference on Neural Information Processing Systems (NeurIPS 2024), December 10 - 15, 2024, Neural Information Processing Systems Foundation, Vancouver, Canada, December 16, 2024.

  • J.-J. Zhu, Kernel approximation of Wasserstein-Fisher-Rao gradient flows, DFG SPP 2298 Annual Meeting, November 10 - 12, 2024, LMU München, Department of Mathematics, Tutzingen, November 11, 2024.

  • J.-J. Zhu, Kernelization, approximation, and entropy dissipation of gradient flows, January 24 - 26, 2024, RIKEN, Center for Advanced Intelligence Project, Japan.

  • J.-J. Zhu, Transport and Flow: The modern mathematics of distributional learning and optimization, Universität des Saarlandes, Saarland Informatics Campus, Saarbrücken, July 5, 2024.

External Preprints

  • E. Gladin, P. Dvurechensky, A. Mielke , J.-J. Zhu, Interaction-force transport gradient flows, Preprint no. arXiv:2405.17075, Cornell University, 2024, DOI 10.48550/arXiv.2405.17075 .
    Abstract
    This paper presents a new type of gradient flow geometries over non-negative and probability measures motivated via a principled construction that combines the optimal transport and interaction forces modeled by reproducing kernels. Concretely, we propose the interaction-force transport (IFT) gradient flows and its spherical variant via an infimal convolution of the Wasserstein and spherical MMD Riemannian metric tensors. We then develop a particle-based optimization algorithm based on the JKO-splitting scheme of the mass-preserving spherical IFT gradient flows. Finally, we provide both theoretical global exponential convergence guarantees and empirical simulation results for applying the IFT gradient flows to the sampling task of MMD-minimization studied by Arbel et al. [2019]. Furthermore, we prove that the spherical IFT gradient flow enjoys the best of both worlds by providing the global exponential convergence guarantee for both the MMD and KL energy.