Stephen Zhang's homepage
About
I will be a Faculty Fellow at NYU Center for Data Science from 2028-29.
I am a PhD student at the University of Melbourne interested in probabilistic machine learning, generative modelling and applied mathematics to tackle inverse problems arising in complex systems, especially those arising in biology and physics.
In May-September 2025, I was an intern with Marco Cuturi at Apple MLR in Paris, working on optimal transport for flow-based generative models.
In April-September 2024, I visited Xiaojie Qiu's group at Stanford, working on generative modelling for single cell genomics.
Previously, I earned a MSc in Mathematics at the University of British Columbia.
Interests…
...theory and algorithms: generative modelling, optimal transport, inverse problems, deep learning, scientific computing
...applications: systems and computational biology, biophysics, scientific inverse problems
Selected papers
* denotes equal contribution.
On Fitting Flow Models with Large Sinkhorn Couplings
Stephen Zhang*, Alireza Mousavi-Hosseini*, Michal Klein, Marco Cuturi
TMLR, 2026
Flow Matching with Semidiscrete Couplings
Alireza Mousavi-Hosseini*, Stephen Zhang*, Michal Klein, Marco Cuturi
ICLR 2026
Learning non-equilibrium diffusions with Schrödinger bridges: from exactly solvable to simulation-free
Stephen Zhang and Michael Stumpf
NeurIPS 2025
Inferring stochastic dynamics with growth from cross-sectional data
Stephen Zhang, Suryanarayana Maddu, Xiaojie Qiu and Victor Chardès
NeurIPS 2025
Learning cell-specific networks from dynamics and geometry of single cells
Stephen Zhang and Michael Stumpf
Cell Systems, 2025
Towards a mathematical theory of trajectory inference
Hugo Lavenant*, Stephen Zhang*, Young-Heon Kim, Geoffrey Schiebinger
Annals of Applied Probability, 2024
Manifold learning with sparse regularised optimal transport
Stephen Zhang*, Gilles Mordant*, Tim Matsumoto, Geoffrey Schiebinger
Preprint, 2023
Trajectory Inference via Mean-field Langevin in Path Space
Lénaïc Chizat*, Stephen Zhang*, Matthieu Heitz, Geoffrey Schiebinger
NeurIPS 2022
Persistent exclusion processes: Inertia, drift, mixing, and correlation
Stephen Zhang, Aaron Chong and Barry Hughes
Physical Review E, 2020
Fun
The animations on the left illustrate measure transport from a standard Gaussian to the image of a Penrose triangle, using a flow matching generative model trained using (a) independently coupled data (I-FM), (b) minibatch optimal transport couplings (OT-FM) with batch size 256, and (c) using a semidiscrete OT coupling (SD-FM).
Semidiscrete Flow Matching trains a flow model that approximates the dynamical Benamou-Brenier optimal interpolation between the two measures.
While this toy example is two dimensional, in our recent paper Flow Matching with Semidiscrete Couplings we scale this to train better flow models on large (millions) datasets in high (~10k dimensions) datasets.