Shortest-Path Flow Matching with Mixture-Conditioned Bases for OOD Generalization to Unseen Conditions

2026-01-16

AI Summary

This work addresses the limited out-of-distribution generalization of existing conditional generative models, which struggle to extrapolate effectively to unseen conditioning inputs. To overcome this challenge, the authors propose MixFlow, a novel framework that introduces a learnable descriptor-dependent mixture distribution as the base measure in conditional flow matching for the first time. By integrating shortest-path flow field modeling, MixFlow enables smooth interpolation and robust extrapolation under previously unobserved conditions. Empirical evaluations on tasks such as single-cell transcriptomic response prediction and high-content microscopic drug screening demonstrate that MixFlow significantly outperforms standard conditional flow matching baselines, exhibiting superior cross-domain generative generalization capabilities.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Mixed Discrete/Continuous Search

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications

Abstract

Robust generalization under distribution shift remains a key challenge for conditional generative modeling: conditional flow-based methods often fit the training conditions well but fail to extrapolate to unseen ones. We introduce SP-FM, a shortest-path flow-matching framework that improves out-of-distribution (OOD) generalization by conditioning both the base distribution and the flow field on the condition. Specifically, SP-FM learns a condition-dependent base distribution parameterized as a flexible, learnable mixture, together with a condition-dependent vector field trained via shortest-path flow matching. Conditioning the base allows the model to adapt its starting distribution across conditions, enabling smooth interpolation and more reliable extrapolation beyond the observed training range. We provide theoretical insights into the resulting conditional transport and show how mixture-conditioned bases enhance robustness under shift. Empirically, SP-FM is effective across heterogeneous domains, including predicting responses to unseen perturbations in single-cell transcriptomics and modeling treatment effects in high-content microscopy--based drug screening. Overall, SP-FM provides a simple yet effective plug-in strategy for improving conditional generative modeling and OOD generalization across diverse domains.

Problem

• out-of-distribution generalization
• conditional generative modeling
• distribution shift
• descriptor-controlled generation

Innovation

• Mixture-conditioned flow matching
• Out-of-distribution generalization
• Descriptor-controlled generation
• Learnable base distribution
• Shortest-path flow matching