Training Free Guided Flow Matching with Optimal Control

2024-10-23

arXiv.org

AI Summary

This work addresses the unsolved challenge of achieving training-free, geometry-aware controllable generation on non-Euclidean manifolds—particularly SO(3)—using pretrained flow matching models. We propose a general guidance framework grounded in optimal control theory, which eliminates costly gradient backpropagation inherent in conventional differential sampling. Our approach is the first to integrate optimal control into flow matching guidance, providing unified convergence guarantees for both Euclidean spaces and SO(3), while subsuming existing ODE-based backpropagation methods as a special case. Key technical contributions include manifold-adapted flow matching modeling, gradient-free control optimization without model retraining, and explicit embedding of SO(3) geometric constraints. Experiments demonstrate substantial improvements over baselines in text-guided image editing, conditional molecular generation, and all-atom peptide design—achieving superior efficiency, numerical stability, and geometric consistency.

Technology Category

• Machine Learning: Learning with Manifolds
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Guiding ODE-based generative models using optimal control.
• Extending guided flow methods to complex geometries like SO(3).
• Providing theoretical convergence guarantees for guided flow matching.

Innovation

• Training-free guided flow matching framework
• Optimal control for ODE-based generation
• Supports Euclidean and SO(3) geometries