Constrained Molecular Generation via Sequential Flow Model Fine-Tuning

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, 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.