Conditional Variable Flow Matching: Transforming Conditional Densities with Amortized Conditional Optimal Transport

TL;DR

CVFM addresses learning conditional stochastic dynamics with unpaired conditioning by learning flows between conditional densities $p_t(x|y)$. It combines two conditional flows, a conditional Wasserstein distance, and a conditioning-loss kernel to realize amortized conditional optimal transport over the conditioning variable. Theoretical insights establish that CVFM aligns with a marginal flow matching objective under mild assumptions and demonstrate improved stability and convergence over prior conditional approaches. Empirical results across 2D toys, MNIST–FashionMNIST domain transfer, and spinodal-decomposition microstructure dynamics illustrate accurate conditional density evolution and robustness in high-dimensional settings.

Abstract

Forecasting conditional stochastic nonlinear dynamical systems is a fundamental challenge repeatedly encountered across the biological and physical sciences. While flow-based models can impressively predict the temporal evolution of probability distributions representing possible outcomes of a specific process, existing frameworks cannot satisfactorily account for the impact of conditioning variables on these dynamics. Amongst several limitations, existing methods require training data with paired conditions and are developed for discrete conditioning variables. We propose Conditional Variable Flow Matching (CVFM), a framework for learning flows transforming conditional distributions with amortization across continuous conditioning variables - permitting predictions across the conditional density manifold. This is accomplished through several novel advances. In particular, simultaneous sample conditioned flows over the main and conditioning variables, alongside a conditional Wasserstein distance combined with a loss reweighting kernel facilitating conditional optimal transport. Collectively, these advances allow for learning system dynamics provided measurement data whose states and conditioning variables are not in correspondence. We demonstrate CVFM on a suite of increasingly challenging problems, including discrete and continuous conditional mapping benchmarks, image-to-image domain transfer, and modeling the temporal evolution of materials internal structure during manufacturing processes. We observe that CVFM results in improved performance and convergence characteristics over alternative conditional variants.

Figures (12)

• Figure 1: (Left) Conditional time-dependent density evolution from 8-Gaussians to Moons through the SDE and ODE formulations of CVFM. (Right) Comparison of conditional flows learned using solely the proposed conditional kernel in isolation (red) and the proposed CVFM framework (blue). The kernel effectively further facilitates disentanglement of the flow conditioning during static conditional OT.
• Figure 2: CVFM results in lower error in Wasserstein-2 distance to target distribution across batch sizes and conditional cost weighting $\eta$, compared to COT-FM (Eq. (\ref{['eq:conditionalwasserstein']})) or the naïve conditional implementation CFM. Trajectories are colored by conditioning variable.
• Figure 3: CVSFM conditionally generated images from the FashionMNIST dataset. (Left) pane displays samples from the initial distributions $p_0(x|y)$, while the middle displays generated samples from $p_1(x|y)$. Relative positioning indicates paired samples. (Right) pane illustrates the improved convergence of CVSFM over COT-SFM in high-dimensional domain transfer for $\eta=10$ and $\eta=1000$. Displayed FID/LPIPS scores are computed per class and averaged.
• Figure 4: Collection of 5 randomly sampled trajectories from the test set in projections of PC1-PC2 and PC2-PC3 displaying (left) samples in blue from CVSFM-Exact, with the expected value as a function of time in red. Comparison of test marginal densities (middle) constructed as narrow subsets of first constituent's mobility ($\mathcal{M}_1$) parameter, in PC1-PC2 projection, and (right) marginal comparison of final state at $t=100$ for PC1. 128 samples were simulated for each $(x_0, y)$.
• Figure E.1: Comparison of obtained trajectories for various OT and SB modeling approaches in the synthetic datasets considered associated with Wasserstein-2 error and normalized path energy reported in Table \ref{['tab:toy_problem_results']}. Trajectories are colored by the conditioning variable. Source distributions are shown in black, target distributions are shown in maroon.

Theorems & Definitions (6)

• Theorem 3.1
• Theorem 3.2
• Proposition 3.1
• Theorem \ref{thm:generatevectorfield}
• Theorem \ref{thm:objectiveequivalence}
• Theorem 3.3