On Conditional Sampling with Joint Flow Matching

A transport map is versatile and useful for many downstream tasks, from training generative modeling to solving Bayesian inference problems. \cite{marzouk2016introduction} pioneered the \emph{measure transport} approach for sampling by introducing its connection with transport map $T_{\sharp} \rho = \mu$, where samples from $\mu$ can be easily drawn. When the transport map is a lower-triangular map or \emph{Knothe-Rosenblatt map}, we can also draw conditional samples $ \mu_{2|1}(\cdot|x_1)$ from the target distribution $\mu$ generalized by \cite{kovachki2020conditional}. This state-of-the-art sampling approach deviates from traditional methods such as MCMC or variational inference and has received many research interests in recent years. In our work, we introduce a new approach to approximate this transport map to perform conditional sampling tasks using a recent computational advance in generative modeling -- flow matching. Specifically, we use \citet{pooladian2023multisample}'s joint flow matching approach with a twisted Euclidean cost to ensure the triangular property of the map. We empirically verify our method through benchmark examples and quantifying the approximated map errors.