Constrained Flow Optimization via Sequential Fine-Tuning for Molecular Design

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.

Modern AI can now generate new molecules, proteins, and materials. But scientific use rarely needs only creativity — a drug candidate, for example, must bind its target strongly (the goal) while also being safe, stable, and chemically realistic (the rules that have to hold). Existing methods can push such generative models toward higher goal scores, but the rules then tend to be silently broken in the process, and there is no reliable way to dial the trade-off between the two. We introduce Constrained Flow Optimization, a method that adapts a pre-trained generative model to pursue a goal and respect a set of rules at the same time, without users having to manually tune the balance between them. The method comes with mathematical guarantees. On a molecular design task, it produces molecules with higher target properties while reliably staying within the physical constraints required for the task at hand.