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

TL;DR

Conditional Variable Flow Matching (CVFM) introduces a novel framework for transforming conditional distributions using amortized flows across continuous conditioning variables, enabling flexible predictions in complex dynamical systems.

Contribution

CVFM is the first method to learn flows for conditional distributions with amortization over continuous variables, addressing limitations of existing models requiring paired data and discrete conditions.

Findings

Improved performance over alternative methods in various benchmarks.

Enhanced convergence characteristics in learning conditional dynamics.

Effective modeling of complex systems like material evolution during manufacturing.

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.