Controlled generation with pre-trained Diffusion and Flow Matching models has vast applications. One strategy for guiding ODE-based generative models is through optimizing a target loss $R(x_1)$ while staying close to the prior distribution. Along this line, some recent work showed the effectiveness of guiding flow model by differentiating through its ODE sampling process. Despite the superior performance, the theoretical understanding of this line of methods is still preliminary, leaving space for algorithm improvement. Moreover, existing methods predominately focus on Euclidean data manifold, and there is a compelling need for guided flow methods on complex geometries such as SO(3), which prevails in high-stake scientific applications like protein design. We present OC-Flow, a general and theoretically grounded training-free framework for guided flow matching using optimal control. Building upon advances in optimal control theory, we develop effective and practical algorithms for solving optimal control in guided ODE-based generation and provide a systematic theoretical analysis of the convergence guarantee in both Euclidean and SO(3). We show that existing backprop-through-ODE methods can be interpreted as special cases of Euclidean OC-Flow. OC-Flow achieved superior performance in extensive experiments on text-guided image manipulation, conditional molecule generation, and all-atom peptide design.

翻译：利用预训练的扩散模型和流匹配模型进行受控生成具有广泛的应用前景。一种引导基于ODE的生成模型的策略是在保持接近先验分布的同时优化目标损失$R(x_1)$。沿着这一思路，近期的一些工作展示了通过对其ODE采样过程进行微分来引导流模型的有效性。尽管性能优越，但对此类方法的理论理解仍处于初步阶段，为算法改进留下了空间。此外，现有方法主要集中于欧几里得数据流形，而在复杂几何结构（如SO(3)）上迫切需要引导流方法，这在蛋白质设计等高价值科学应用中普遍存在。我们提出了OC-Flow，一个基于最优控制理论的、通用的、理论完备的无训练引导流匹配框架。基于最优控制理论的进展，我们开发了有效且实用的算法来解决基于ODE的引导生成中的最优控制问题，并对欧几里得空间和SO(3)上的收敛保证提供了系统的理论分析。我们证明了现有的通过ODE反向传播方法可以解释为欧几里得OC-Flow的特例。OC-Flow在文本引导图像编辑、条件分子生成和全原子肽设计的大量实验中取得了卓越的性能。