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