Constrained Flow Optimization via Sequential Fine-Tuning for Molecular Design

Sven Gutjahr, Riccardo De Santi, Luca Schaufelberger, Kjell Jorner, Andreas Krause

Keywords: diffusion models, flow models, molecular design, fine-tuning, constrained generation

Abstract: Adapting generative foundation models to optimize rewards of interest (e.g., binding affinity) while satisfying constraints (e.g., molecular synthesizability) is of fundamental importance to render them applicable in real-world discovery campaigns such as molecular design. While recent works have introduced scalable methods for reward-guided fine-tuning of diffusion and flow models, it remains an open problem how to algorithmically trade off property maximization and constraint satisfaction in a reliable and predictable manner. Towards tackling this challenging problem, we first present a rigorous formulation for constrained generative optimization. Then, we introduce $\textbf{C}$onstrained $\textbf{F}$low $\textbf{O}$ptimization (CFO), an augmented Lagrangian method that renders it possible to arbitrarily control the aforementioned trade-off between reward maximization and constraint satisfaction. We provide convergence guarantees for the proposed scheme. Ultimately, we present an experimental evaluation on both synthetic, yet illustrative, settings, and a molecular design task optimizing molecular properties while constraining energy.

Submission Number: 85