Constrained Molecular Generation via Sequential Flow Model Fine-Tuning

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 or protein engineering. 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, in this work we first present a rigorous formulation for constrained generation. Then, we introduce Augmented Lagrangian Flows Fine-tuning (ALF$^2$), an augmented Lagrangian method that renders 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: 154