Authors: Luran Wang, Chaoran Cheng, Yizhen Liao, Yanru Qu, Ge Liu
Abstract: 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.
Source: http://arxiv.org/abs/2410.18070v1
