Shortest-Path Flow Matching with Mixture-Conditioned Bases for OOD Generalization to Unseen Conditions
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2026
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866908827868200960 |
|---|---|
| author | Rubbi, Andrea Akbarnejad, Amir Sanian, Mohammad Vali Parast, Aryan Yazdan Asadollahzadeh, Hesam Amani, Arian Akhtar, Naveed Cooper, Sarah Bassett, Andrew Liò, Pietro Paavolainen, Lassi Vakili, Sattar Lotfollahi, Mo |
| author_facet | Rubbi, Andrea Akbarnejad, Amir Sanian, Mohammad Vali Parast, Aryan Yazdan Asadollahzadeh, Hesam Amani, Arian Akhtar, Naveed Cooper, Sarah Bassett, Andrew Liò, Pietro Paavolainen, Lassi Vakili, Sattar Lotfollahi, Mo |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2601_11827 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Shortest-Path Flow Matching with Mixture-Conditioned Bases for OOD Generalization to Unseen Conditions Rubbi, Andrea Akbarnejad, Amir Sanian, Mohammad Vali Parast, Aryan Yazdan Asadollahzadeh, Hesam Amani, Arian Akhtar, Naveed Cooper, Sarah Bassett, Andrew Liò, Pietro Paavolainen, Lassi Vakili, Sattar Lotfollahi, Mo Machine Learning Computer Vision and Pattern Recognition Image and Video Processing 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. |
| title | Shortest-Path Flow Matching with Mixture-Conditioned Bases for OOD Generalization to Unseen Conditions |
| topic | Machine Learning Computer Vision and Pattern Recognition Image and Video Processing |
| url | https://arxiv.org/abs/2601.11827 |