MixFlow: Mixture-Conditioned Flow Matching for Out-of-Distribution Generalization

TL;DR

MixFlow tackles the poor out-of-distribution generalization of conditional flow models by jointly learning a descriptor-conditioned Gaussian mixture base distribution and a descriptor-conditioned velocity field trained via shortest-path flow matching. Theoretical analysis shows that moving beyond a single Gaussian base reduces degrees of freedom in the transport problem, enabling better identifiability and stability for unseen conditions. Empirically, MixFlow yields substantial improvements over standard conditional flow matching across synthetic rotations, single-cell perturbations, and image-based phenotypic assays, demonstrating stronger distributional alignment under novel perturbations. This framework offers a principled and scalable approach to robust, controllable generative modeling in heterogeneous scientific domains, with potential applications in drug discovery and cellular perturbation prediction.

Abstract

Achieving robust generalization under distribution shift remains a central challenge in conditional generative modeling, as existing conditional flow-based methods often struggle to extrapolate beyond the training conditions. We introduce MixFlow, a conditional flow-matching framework for descriptor-controlled generation that directly targets this limitation by jointly learning a descriptor-conditioned base distribution and a descriptor-conditioned flow field via shortest-path flow matching. By modeling the base distribution as a learnable, descriptor-dependent mixture, MixFlow enables smooth interpolation and extrapolation to unseen conditions, leading to substantially improved out-of-distribution generalization. We provide analytical insights into the behavior of the proposed framework and empirically demonstrate its effectiveness across multiple domains, including prediction of responses to unseen perturbations in single-cell transcriptomic data and high-content microscopy-based drug screening tasks. Across these diverse settings, MixFlow consistently outperforms standard conditional flow-matching baselines. Overall, MixFlow offers a simple yet powerful approach for achieving robust, generalizable, and controllable generative modeling across heterogeneous domains.

Figures (3)

• Figure 1: Standard conditional flow matching versus MixFlow. (a) Conventional approaches fix the base distribution to a standard Gaussian and condition only the velocity field on the descriptor $y$, effectively learning independent flows for each condition. (b) MixFlow conditions both the base distribution (as a Gaussian mixture) and the velocity field on the descriptor. Similar conditions produce similar base distributions, requiring only small flow corrections and enabling smooth extrapolation to unseen conditions.
• Figure 2: Qualitative comparison of Conditional Flow Matching (CFM) and MixFlow on the synthetic letter–rotation benchmark. Each block shows samples at different integration times ($t=0$, $t=0.5$, $t=1$) and the corresponding target distribution. MixFlow produces smoother trajectories and more faithful reconstructions on both training and unseen rotations, while CFM exhibits drift and shape distortion, particularly on validation conditions.
• Figure 3: Ablation studies on transcriptomic (Combo-SciPlex) and morphological (BBBC021) datasets. (a) Effect of the number of GMM modes on generation quality. (b) Comparison of MixFlow and CFM robustness across increasing data dimensionality.

Theorems & Definitions (5)

• Definition 2: Mixture Wasserstein Distance
• Definition 3: Mixture Wasserstein Flow
• Lemma 4: Dual of Mixture Wasserstein Distance
• Definition 5: Subset sum condition
• Definition 9: Projection to subspace of GMMs with $I$ modes