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.
翻译:在分布偏移下的稳健泛化仍然是条件生成建模的一个关键挑战:基于条件流的方法通常能很好地拟合训练条件,但难以外推至未见条件。本文提出SP-FM——一种最短路径流匹配框架,通过将基分布与流场同时条件化于输入条件,以提升分布外(OOD)泛化性能。具体而言,SP-FM学习一个由灵活可学习的混合分布参数化的条件依赖型基分布,并结合一个通过最短路径流匹配训练的条件依赖型向量场。对基分布进行条件化使模型能够根据不同条件调整其起始分布,从而实现平滑插值并在超出观测训练范围时进行更可靠的外推。我们从理论上分析了由此产生的条件传输过程,并阐释了混合条件基如何增强分布偏移下的稳健性。实验表明,SP-FM在多个异质领域均表现有效,包括单细胞转录组学中对未见扰动的响应预测,以及基于高内涵显微镜的药物筛选中的治疗效果建模。总体而言,SP-FM为改进条件生成建模及跨领域OOD泛化提供了一种简单而有效的即插即用策略。