Constrained Flow Optimization via Sequential Fine Tuning for Molecular Design
This work provides a method for reliably balancing reward maximization and constraint satisfaction in generative models, which is crucial for real-world scientific discovery applications like molecular design, benefiting researchers and practitioners in drug discovery and materials science.
This paper addresses the challenge of optimizing generative foundation models for molecular design by simultaneously maximizing a reward function (e.g., binding affinity) and satisfying constraints (e.g., synthesizability). The proposed Constrained Flow Optimization (CFO) algorithm achieves consistent increases in reward while maintaining high constraint satisfaction across synthetic and molecular design tasks.
Adapting generative foundation models, in particular diffusion and flow models, to optimize given reward functions (e.g., binding affinity) while satisfying constraints (e.g., molecular synthesizability) is fundamental for their adoption in real-world scientific discovery applications such as molecular design or protein engineering. While recent works have introduced scalable methods for reward-guided fine-tuning of such models via reinforcement learning and control schemes, it remains an open problem how to algorithmically trade-off reward maximization and constraint satisfaction in a reliable and predictable manner. Motivated by this challenge, we first present a rigorous framework for Constrained Generative Optimization, which brings an optimization viewpoint to the introduced adaptation problem and retrieves the relevant task of constrained generation as a sub-case. Then, we introduce Constrained Flow Optimization (CFO), an algorithm that automatically and provably balances reward maximization and constraint satisfaction by reducing the original problem to sequential fine-tuning via established, scalable methods. We provide convergence guarantees for constrained generative optimization and constrained generation via CFO. Ultimately, we present an experimental evaluation of CFO on both synthetic, yet illustrative, settings, and a molecular design task. Across these evaluations, CFO achieves consistent increases in reward while ensuring high constraint satisfaction, showcasing its practical utility for constrained generative optimization.